Merge branch 'fenicsx' into reactions-in-formulation

festim/hydrogen_transport_problem.py

@@ -22,7 +22,8 @@ class HydrogenTransportProblem:
         mesh (festim.Mesh): the mesh of the model
         subdomains (list of festim.Subdomain): the subdomains of the model
         species (list of festim.Species): the species of the model
-        temperature (float or fem.Constant): the temperature of the model
+        temperature (float, int, fem.Constant, fem.Function or callable): the
+            temperature of the model (K)
         sources (list of festim.Source): the hydrogen sources of the model
         boundary_conditions (list of festim.BoundaryCondition): the boundary
             conditions of the model
@@ -33,7 +34,8 @@ class HydrogenTransportProblem:
         mesh (festim.Mesh): the mesh of the model
         subdomains (list of festim.Subdomain): the subdomains of the model
         species (list of festim.Species): the species of the model
-        temperature (fem.Constant): the temperature of the model
+        temperature (float, int, fem.Constant, fem.Function or callable): the
+            temperature of the model (K)
         boundary_conditions (list of festim.BoundaryCondition): the boundary
             conditions of the model
         solver_parameters (dict): the solver parameters of the model
@@ -47,7 +49,15 @@ class HydrogenTransportProblem:
             model
         formulation (ufl.form.Form): the formulation of the model
         solver (dolfinx.nls.newton.NewtonSolver): the solver of the model
-        multispecies (bool): True if the model has more than one species
+        multispecies (bool): True if the model has more than one species.
+        temperature_fenics (fem.Constant or fem.Function): the
+            temperature of the model as a fenics object (fem.Constant or
+            fem.Function).
+        temperature_expr (fem.Expression): the expression of the temperature
+            that is used to update the temperature_fenics
+        temperature_time_dependent (bool): True if the temperature is time
+            dependent
+
 
     Usage:
         >>> import festim as F
@@ -101,6 +111,7 @@ def __init__(
         self.formulation = None
         self.volume_subdomains = []
         self.bc_forms = []
+        self.temperature_fenics = None
 
     @property
     def temperature(self):
@@ -110,8 +121,39 @@ def temperature(self):
     def temperature(self, value):
         if value is None:
             self._temperature = value
+        elif isinstance(value, (float, int, fem.Constant, fem.Function)):
+            self._temperature = value
+        elif callable(value):
+            self._temperature = value
+        else:
+            raise TypeError(
+                f"Value must be a float, int, fem.Constant, fem.Function, or callable"
+            )
+
+    @property
+    def temperature_fenics(self):
+        return self._temperature_fenics
+
+    @temperature_fenics.setter
+    def temperature_fenics(self, value):
+        if value is None:
+            self._temperature_fenics = value
+            return
+        elif not isinstance(value, (fem.Constant, fem.Function)):
+            raise TypeError(f"Value must be a fem.Constant or fem.Function")
+        self._temperature_fenics = value
+
+    @property
+    def temperature_time_dependent(self):
+        if self.temperature is None:
+            return False
+        if isinstance(self.temperature, fem.Constant):
+            return False
+        if callable(self.temperature):
+            arguments = self.temperature.__code__.co_varnames
+            return "t" in arguments
         else:
-            self._temperature = F.as_fenics_constant(value, self.mesh.mesh)
+            return False
 
     @property
     def multispecies(self):
@@ -139,11 +181,68 @@ def initialise(self):
         self.t = fem.Constant(self.mesh.mesh, 0.0)
         self.dt = self.settings.stepsize.get_dt(self.mesh.mesh)
 
+        self.define_temperature()
         self.define_boundary_conditions()
         self.create_formulation()
         self.create_solver()
         self.defing_export_writers()
 
+    def define_temperature(self):
+        """Sets the value of temperature_fenics_value. The type depends on
+        self.temperature. If self.temperature is a function on t only, create
+        a fem.Constant. Else, create an dolfinx.fem.Expression (stored in
+        self.temperature_expr) to be updated, a dolfinx.fem.Function object
+        is created from the Expression (stored in self.temperature_fenics_value).
+        Raise a ValueError if temperature is None.
+        """
+        # check if temperature is None
+        if self.temperature is None:
+            raise ValueError("the temperature attribute needs to be defined")
+
+        # if temperature is a float or int, create a fem.Constant
+        elif isinstance(self.temperature, (float, int)):
+            self.temperature_fenics = F.as_fenics_constant(
+                self.temperature, self.mesh.mesh
+            )
+        # if temperature is a fem.Constant or function, pass it to temperature_fenics
+        elif isinstance(self.temperature, (fem.Constant, fem.Function)):
+            self.temperature_fenics = self.temperature
+
+        # if temperature is callable, process accordingly
+        elif callable(self.temperature):
+            arguments = self.temperature.__code__.co_varnames
+            if "t" in arguments and "x" not in arguments and "T" not in arguments:
+                # only t is an argument
+                self.temperature_fenics = F.as_fenics_constant(
+                    mesh=self.mesh.mesh, value=self.temperature(t=float(self.t))
+                )
+            else:
+                x = ufl.SpatialCoordinate(self.mesh.mesh)
+                degree = 1
+                element_temperature = basix.ufl.element(
+                    basix.ElementFamily.P,
+                    self.mesh.mesh.basix_cell(),
+                    degree,
+                    basix.LagrangeVariant.equispaced,
+                )
+                function_space_temperature = fem.FunctionSpace(
+                    self.mesh.mesh, element_temperature
+                )
+                self.temperature_fenics = fem.Function(function_space_temperature)
+                kwargs = {}
+                if "t" in arguments:
+                    kwargs["t"] = self.t
+                if "x" in arguments:
+                    kwargs["x"] = x
+
+                # store the expression of the temperature
+                # to update the temperature_fenics later
+                self.temperature_expr = fem.Expression(
+                    self.temperature(**kwargs),
+                    function_space_temperature.element.interpolation_points(),
+                )
+                self.temperature_fenics.interpolate(self.temperature_expr)
+
     def defing_export_writers(self):
         """Defines the export writers of the model, if field is given as
         a string, find species object in self.species"""
@@ -162,10 +261,11 @@ def define_function_space(self):
         """Creates the function space of the model, creates a mixed element if
         model is multispecies. Creates the main solution and previous solution
         function u and u_n."""
+        degree = 1
         element_CG = basix.ufl.element(
             basix.ElementFamily.P,
             self.mesh.mesh.basix_cell(),
-            1,
+            degree,
             basix.LagrangeVariant.equispaced,
         )
         if not self.multispecies:
@@ -291,7 +391,7 @@ def create_dirichletbc_form(self, bc):
 
         bc.create_value(
             mesh=self.mesh.mesh,
-            temperature=self.temperature,
+            temperature=self.temperature_fenics,
             function_space=function_space_value,
             t=self.t,
         )
@@ -340,7 +440,7 @@ def create_formulation(self):
 
             for vol in self.volume_subdomains:
                 D = vol.material.get_diffusion_coefficient(
-                    self.mesh.mesh, self.temperature, spe
+                    self.mesh.mesh, self.temperature_fenics, spe
                 )
                 if spe.mobile:
                     self.formulation += dot(D * grad(u), grad(v)) * self.dx(vol.id)
@@ -424,9 +524,7 @@ def iterate(
         self.progress.update(self.dt.value)
         self.t.value += self.dt.value
 
-        # update boundary conditions
-        for bc in self.boundary_conditions:
-            bc.update(float(self.t))
+        self.update_time_dependent_values()
 
         self.solver.solve(self.u)
 
@@ -435,7 +533,7 @@ def iterate(
         if not skip_post_processing:
             if not self.multispecies:
                 D_D = self.subdomains[0].material.get_diffusion_coefficient(
-                    self.mesh.mesh, self.temperature, self.species[0]
+                    self.mesh.mesh, self.temperature_fenics, self.species[0]
                 )
                 cm = self.u
                 self.species[0].post_processing_solution = self.u
@@ -450,10 +548,10 @@ def iterate(
 
                 cm_1, cm_2 = self.u.split()
                 D_1 = self.subdomains[0].material.get_diffusion_coefficient(
-                    self.mesh.mesh, self.temperature, self.species[0]
+                    self.mesh.mesh, self.temperature_fenics, self.species[0]
                 )
                 D_2 = self.subdomains[0].material.get_diffusion_coefficient(
-                    self.mesh.mesh, self.temperature, self.species[1]
+                    self.mesh.mesh, self.temperature_fenics, self.species[1]
                 )
                 surface_flux_1 = form(D_1 * dot(grad(cm_1), self.mesh.n) * self.ds(2))
                 surface_flux_2 = form(D_2 * dot(grad(cm_2), self.mesh.n) * self.ds(2))
@@ -469,3 +567,14 @@ def iterate(
 
         # update previous solution
         self.u_n.x.array[:] = self.u.x.array[:]
+
+    def update_time_dependent_values(self):
+        t = float(self.t)
+        if self.temperature_time_dependent:
+            if isinstance(self.temperature_fenics, fem.Constant):
+                self.temperature_fenics.value = self.temperature(t=t)
+            else:
+                self.temperature_fenics.interpolate(self.temperature_expr)
+
+        for bc in self.boundary_conditions:
+            bc.update(t)

test/test_dirichlet_bc.py

@@ -297,6 +297,7 @@ def test_species_predefined():
     my_model.species = [F.Species("H")]
     my_bc = F.DirichletBC(subdomain, 1.0, "J")
     my_model.boundary_conditions = [my_bc]
+    my_model.temperature = 1
     my_model.settings = F.Settings(atol=1, rtol=0.1)
     my_model.settings.stepsize = 1