Export implicit species

festim/exports/vtx.py

@@ -39,6 +39,9 @@ def write(self, t: float):
         Args:
             t (float): the time of export
         """
+        for field in self.field:
+            if isinstance(field, F.ImplicitSpecies):
+                field.post_processing_solution.interpolate(field.concentration_expr)
         self.writer.write(t)
 
 
@@ -52,7 +55,7 @@ class VTXExport(VTXExportBase):
     Attributes:
         filename (str): the name of the output file
         writer (dolfinx.io.VTXWriter): the VTX writer
-        field (festim.Species, list of festim.Species): the field index to export
+        field (list of festim.Species): the field index to export
 
     Usage:
         >>> u = dolfinx.fem.Function(V)
@@ -74,16 +77,17 @@ def field(self):
     @field.setter
     def field(self, value):
         # check that field is festim.Species or list of festim.Species
-        if not isinstance(value, (F.Species, str)) and not isinstance(value, list):
+        accepted_types = (F.Species, str, F.ImplicitSpecies)
+        if not isinstance(value, accepted_types) and not isinstance(value, list):
             raise TypeError(
-                "field must be of type festim.Species or str or a list of festim.Species or str"
+                f"field must be of type {accepted_types} or a list of {accepted_types}, not {type(value)}"
             )
         # check that all elements of list are festim.Species
         if isinstance(value, list):
             for element in value:
-                if not isinstance(element, (F.Species, str)):
+                if not isinstance(element, accepted_types):
                     raise TypeError(
-                        "field must be of type festim.Species or str or a list of festim.Species or str"
+                        f"field must be of type {accepted_types} or a list of {accepted_types}, not {type(element)}"
                     )
         # if field is festim.Species, convert to list
         if not isinstance(value, list):

festim/hydrogen_transport_problem.py

@@ -245,6 +245,12 @@ def initialise(self):
                 self.settings.stepsize.initial_value, self.mesh.mesh
             )
 
+        # convert implicit species n to dolfinx
+        for reaction in self.reactions:
+            for reactant in reaction.reactant:
+                if isinstance(reactant, F.ImplicitSpecies):
+                    reactant.convert_n_to_dolfinx(function_space=self.V_DG_1, t=self.t)
+
         self.define_temperature()
         self.define_boundary_conditions()
         self.create_source_values_fenics()
@@ -460,6 +466,12 @@ def assign_functions_to_species(self):
                 spe.collapsed_function_space, _ = self.function_space.sub(
                     idx
                 ).collapse()
+        for reaction in self.reactions:
+            for reactant in reaction.reactant:
+                if isinstance(reactant, F.ImplicitSpecies):
+                    reactant.post_processing_solution = fem.Function(self.V_DG_1)
+                    if reactant.name:
+                        reactant.post_processing_solution.name = reactant.name
 
         for idx, spe in enumerate(self.species):
             spe.solution = sub_solutions[idx]
@@ -517,7 +529,7 @@ def define_boundary_conditions(self):
             if isinstance(bc.species, str):
                 # if name of species is given then replace with species object
                 bc.species = F.find_species_from_name(bc.species, self.species)
-            if isinstance(bc, F.DirichletBC):
+            if isinstance(bc, F.DirichletBCBase):
                 form = self.create_dirichletbc_form(bc)
                 self.bc_forms.append(form)
 
@@ -788,6 +800,11 @@ def update_time_dependent_values(self):
             elif self.temperature_time_dependent and source.temperature_dependent:
                 source.update(t=t)
 
+        for reaction in self.reactions:
+            for reactant in reaction.reactant:
+                if isinstance(reactant, F.ImplicitSpecies):
+                    reactant.update(t=t)
+
     def post_processing(self):
         """Post processes the model"""
 

festim/species.py

@@ -1,5 +1,7 @@
 from typing import List
 import festim as F
+import dolfinx.fem as fem
+import ufl
 
 
 class Species:
@@ -178,6 +180,14 @@ class ImplicitSpecies:
         others (List[Species]): the list of species from which the implicit
             species concentration is computed (c = n - others)
         concentration (form): the concentration of the species
+        n_as_dolfinx (dolfinx.fem.Constant): the total concentration of the
+            species as a dolfinx object
+        n_expr (dolfinx.fem.Expression): the expression of the total
+            concentration of the species
+        concentration_expr (dolfinx.fem.Expression): the expression of the
+            concentration of the species (c = n - others)
+        post_processing_solution (dolfinx.fem.Function): the solution for post
+            processing
 
     """
 
@@ -191,6 +201,11 @@ def __init__(
         self.n = n
         self.others = others
 
+        self.n_as_dolfinx = None
+        self.n_expr = None
+
+        self._concentration_expr = None
+
     def __repr__(self) -> str:
         return f"ImplicitSpecies({self.name}, {self.n}, {self.others})"
 
@@ -205,7 +220,76 @@ def concentration(self):
                     raise ValueError(
                         f"Cannot compute concentration of {self.name} because {other.name} has no solution"
                     )
-        return self.n - sum([other.solution for other in self.others])
+        return self.n_as_dolfinx - sum([other.solution for other in self.others])
+
+    @property
+    def concentration_expr(self):
+        if self._concentration_expr is None:
+            self._concentration_expr = fem.Expression(
+                self.concentration,
+                self.post_processing_solution.function_space.element.interpolation_points(),
+            )
+        return self._concentration_expr
+
+    def convert_n_to_dolfinx(
+        self, function_space: fem.FunctionSpaceBase, t: fem.Constant
+    ):
+        """Converts the total concentration of the species to a dolfinx object
+
+        Args:
+            function_space (dolfinx.fem.FunctionSpaceBase): the function space
+            t (dolfinx.fem.Constant): the time
+        """
+        mesh = function_space.mesh
+        x = ufl.SpatialCoordinate(mesh)
+
+        if isinstance(self.n, (int, float)):
+            n_as_dolfinx = F.as_fenics_constant(mesh=mesh, value=self.n)
+
+        elif callable(self.n):
+            arguments = self.n.__code__.co_varnames
+
+            if "t" in arguments and "x" not in arguments and "T" not in arguments:
+                # only t is an argument
+                if not isinstance(self.n(t=float(t)), (float, int)):
+                    raise ValueError(
+                        f"self.n should return a float or an int, not {type(self.n(t=float(t)))} "
+                    )
+                n_as_dolfinx = F.as_fenics_constant(mesh=mesh, value=self.n(t=float(t)))
+            else:
+                n_as_dolfinx = fem.Function(function_space)
+                kwargs = {}
+                if "t" in arguments:
+                    kwargs["t"] = t
+                if "x" in arguments:
+                    kwargs["x"] = x
+                if "T" in arguments:
+                    raise NotImplementedError(
+                        "Temperature dependent implicit species not implemented"
+                    )
+
+                # store the expression of the boundary condition
+                # to update the n later
+                self.n_expr = fem.Expression(
+                    self.n(**kwargs),
+                    function_space.element.interpolation_points(),
+                )
+                n_as_dolfinx.interpolate(self.n_expr)
+
+        self.n_as_dolfinx = n_as_dolfinx
+
+    def update(self, t):
+        """Updates the total concentration of the species
+
+        Args:
+            t (float): the time
+        """
+        if callable(self.n):
+            arguments = self.n.__code__.co_varnames
+            if isinstance(self.n_as_dolfinx, fem.Constant) and "t" in arguments:
+                self.n_as_dolfinx.value = self.value(t=t)
+            else:
+                self.n_as_dolfinx.interpolate(self.n_expr)
 
 
 def find_species_from_name(name: str, species: list):

mwe_time_dependent_trap.py

@@ -0,0 +1,51 @@
+import festim as F
+import ufl
+import numpy as np
+
+my_model = F.HydrogenTransportProblem()
+
+L = 1
+my_model.mesh = F.Mesh1D(np.linspace(0, L, num=500))
+
+mat = F.Material(D_0=1, E_D=0)
+
+left_subdomain = F.SurfaceSubdomain1D(id=1, x=0)
+right_subdomain = F.SurfaceSubdomain1D(id=2, x=L)
+volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=mat)
+
+H = F.Species(name="H")
+trapped_H = F.Species(name="H_t", mobile=False)
+
+time_dependent_density = lambda x: ufl.conditional(ufl.lt(x[0], 0.5), 10, 0)
+empty_traps = F.ImplicitSpecies(n=time_dependent_density, others=[trapped_H])
+
+my_model.species = [H, trapped_H]
+
+my_model.reactions = [
+    F.Reaction(
+        reactant=[H, empty_traps],
+        product=[trapped_H],
+        k_0=0.1,
+        E_k=0,
+        p_0=0,
+        E_p=0,
+        volume=volume,
+    ),
+]
+
+
+my_model.subdomains = [left_subdomain, right_subdomain, volume]
+
+my_model.temperature = 500
+
+my_model.boundary_conditions = [
+    F.FixedConcentrationBC(subdomain=left_subdomain, value=10, species=H),
+    F.FixedConcentrationBC(subdomain=right_subdomain, value=0, species=H),
+]
+
+my_model.settings = F.Settings(atol=1e-10, rtol=1e-10, final_time=100)
+my_model.settings.stepsize = F.Stepsize(0.2)
+
+my_model.exports = [F.VTXExport("results.bp", field=empty_traps)]
+my_model.initialise()
+my_model.run()

test/test_species.py

@@ -3,7 +3,7 @@
 import ufl
 import numpy as np
 import pytest
-from dolfinx.fem import functionspace, Function
+from dolfinx.fem import functionspace, Function, Constant
 from dolfinx.mesh import create_unit_cube
 from mpi4py import MPI
 
@@ -88,8 +88,10 @@ def test_implicit_species_concentration():
     species1.solution = Function(V)
     species2.solution = Function(V)
 
+    implicit_species.convert_n_to_dolfinx(function_space=V, t=Constant(mesh, 0.0))
+
     # test the concentration of the implicit species
-    expected_concentration = implicit_species.n - (
+    expected_concentration = implicit_species.n_as_dolfinx - (
         species1.solution + species2.solution
     )
     assert implicit_species.concentration == expected_concentration

test/test_vtx.py

@@ -45,7 +45,12 @@ def test_vtx_export_two_functions(tmpdir):
         my_export.write(t)
 
 
-def test_vtx_integration_with_h_transport_problem(tmpdir):
+H = F.Species("H")
+implicit = F.ImplicitSpecies(name="empty", others=[H], n=20)
+
+
+@pytest.mark.parametrize("species_to_export", [H, implicit])
+def test_vtx_integration_with_h_transport_problem(tmpdir, species_to_export):
     my_model = F.HydrogenTransportProblem()
     my_model.mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
     my_mat = F.Material(D_0=1, E_D=0, name="mat")
@@ -54,11 +59,20 @@ def test_vtx_integration_with_h_transport_problem(tmpdir):
         F.SurfaceSubdomain1D(1, x=0.0),
         F.SurfaceSubdomain1D(2, x=4.0),
     ]
-    my_model.species = [F.Species("H")]
+    my_model.species = [H]
+    my_model.reactions = [
+        F.Reaction(
+            reactant=[H, implicit],
+            product=[],
+            k_0=1,
+            E_k=0,
+            volume=my_model.subdomains[0],
+        )
+    ]
     my_model.temperature = 500
 
     filename = str(tmpdir.join("my_export.bp"))
-    my_export = F.VTXExport(filename, field=my_model.species[0])
+    my_export = F.VTXExport(filename, field=species_to_export)
     my_model.exports = [my_export]
     my_model.settings = F.Settings(atol=1, rtol=0.1)
     my_model.settings.stepsize = F.Stepsize(initial_value=1)