ImplicitSpecies.n can be space and time dependent

src/festim/hydrogen_transport_problem.py

@@ -271,6 +271,8 @@ def initialise(self):
                 self.settings.stepsize.initial_value, self.mesh.mesh
             )
 
+        self.create_implicit_species_value_fenics()
+
         self.define_temperature()
         self.define_boundary_conditions()
         self.create_source_values_fenics()
@@ -280,6 +282,16 @@ def initialise(self):
         self.create_solver()
         self.initialise_exports()
 
+    def create_implicit_species_value_fenics(self):
+        """For each implicit species, create the value_fenics"""
+        for reaction in self.reactions:
+            for reactant in reaction.reactant:
+                if isinstance(reactant, _species.ImplicitSpecies):
+                    reactant.create_value_fenics(
+                        mesh=self.mesh.mesh,
+                        t=self.t,
+                    )
+
     def create_species_from_traps(self):
         """Generate a species and reaction per trap defined in self.traps"""
 
@@ -729,11 +741,16 @@ def create_formulation(self):
     def update_time_dependent_values(self):
         super().update_time_dependent_values()
 
+        t = float(self.t)
+
+        for reaction in self.reactions:
+            for reactant in reaction.reactant:
+                if isinstance(reactant, _species.ImplicitSpecies):
+                    reactant.update_density(t=t)
+
         if not self.temperature_time_dependent:
             return
 
-        t = float(self.t)
-
         if isinstance(self.temperature_fenics, fem.Constant):
             self.temperature_fenics.value = self.temperature(t=t)
         elif isinstance(self.temperature_fenics, fem.Function):
@@ -898,6 +915,8 @@ def initialise(self):
                 self.settings.stepsize.initial_value, self.mesh.mesh
             )
 
+        self.create_implicit_species_value_fenics()
+
         self.define_temperature()
         self.create_source_values_fenics()
         self.create_flux_values_fenics()

src/festim/species.py

@@ -1,6 +1,11 @@
 from festim.subdomain.volume_subdomain import (
     VolumeSubdomain as _VolumeSubdomain,
 )
+from festim.helpers import as_fenics_constant
+
+from typing import List, Union
+import ufl
+from dolfinx import fem
 
 
 class Species:
@@ -105,7 +110,7 @@ class ImplicitSpecies:
     c = n - others
 
     Args:
-        n (float): the total concentration of the species
+        n (Union[float, callable]): the total concentration of the species
         others (list[Species]): the list of species from which the implicit
             species concentration is computed (c = n - others)
         name (str, optional): a name given to the species. Defaults to None.
@@ -116,13 +121,13 @@ 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
-
+        value_fenics: the total concentration as a fenics object
     """
 
     def __init__(
         self,
-        n: float,
-        others: list[Species] = None,
+        n: Union[float, callable],
+        others: List[Species] = None,
         name: str = None,
     ) -> None:
         self.name = name
@@ -144,7 +149,61 @@ def concentration(self):
                         f"Cannot compute concentration of {self.name} "
                         + f"because {other.name} has no solution."
                     )
-        return self.n - sum([other.solution for other in self.others])
+        return self.value_fenics - sum([other.solution for other in self.others])
+
+    def create_value_fenics(self, mesh, t: fem.Constant):
+        """Creates the value of the density as a fenics object and sets it to
+        self.value_fenics.
+        If the value is a constant, it is converted to a fenics.Constant.
+        If the value is a function of t, it is converted to a fenics.Constant.
+        Otherwise, it is converted to a ufl Expression
+
+        Args:
+            mesh (dolfinx.mesh.Mesh) : the mesh
+            t (dolfinx.fem.Constant): the time
+        """
+        x = ufl.SpatialCoordinate(mesh)
+
+        if isinstance(self.n, (int, float)):
+            self.value_fenics = as_fenics_constant(mesh=mesh, value=self.n)
+
+        elif isinstance(self.n, (fem.Function, ufl.core.expr.Expr)):
+            self.value_fenics = 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.value should return a float or an int, not {type(self.n(t=float(t)))} "
+                    )
+                self.value_fenics = as_fenics_constant(
+                    mesh=mesh, value=self.n(t=float(t))
+                )
+            else:
+                kwargs = {}
+                if "t" in arguments:
+                    kwargs["t"] = t
+                if "x" in arguments:
+                    kwargs["x"] = x
+
+                self.value_fenics = self.n(**kwargs)
+
+    def update_density(self, t):
+        """Updates the density value (only if the density is a function of time only)
+
+        Args:
+            t (float): the time
+        """
+        if isinstance(self.n, (fem.Function, ufl.core.expr.Expr)):
+            return
+
+        if callable(self.n):
+            arguments = self.n.__code__.co_varnames
+            if isinstance(self.value_fenics, fem.Constant) and "t" in arguments:
+                self.value_fenics.value = self.n(t=t)
 
 
 def find_species_from_name(name: str, species: list):

src/festim/trap.py

@@ -29,6 +29,7 @@ class Trap(_Species):
         volume (F.VolumeSubdomain1D): The volume subdomain where the trap is.
         trapped_concentration (_Species): The immobile trapped concentration
         trap_reaction (_Reaction): The reaction for trapping the mobile conc.
+        empty_trap_sites (F.ImplicitSpecies): The implicit species for the empty trap sites
 
     Usage:
         >>> import festim as F
@@ -94,10 +95,11 @@ def __init__(
     def create_species_and_reaction(self):
         """create the immobile trapped species object and the reaction for trapping"""
         self.trapped_concentration = _Species(name=self.name, mobile=False)
-        trap_site = _ImplicitSpecies(n=self.n, others=[self.trapped_concentration])
-
+        self.empty_trap_sites = _ImplicitSpecies(
+            n=self.n, others=[self.trapped_concentration]
+        )
         self.reaction = _Reaction(
-            reactant=[self.mobile_species, trap_site],
+            reactant=[self.mobile_species, self.empty_trap_sites],
             product=self.trapped_concentration,
             k_0=self.k_0,
             E_k=self.E_k,

test/system_tests/test_non_homogeneous_density.py

@@ -0,0 +1,139 @@
+import festim as F
+import ufl
+import numpy as np
+import pytest
+import dolfinx.fem as fem
+import basix
+
+step_function_space = lambda x: ufl.conditional(ufl.gt(x[0], 0.5), 10, 0.0)
+step_function_time_non_homogeneous = lambda x, t: ufl.conditional(
+    ufl.gt(t, 50), 10 - 10 * x[0], 0.0
+)
+step_function_time_homogeneous = lambda t: 10.0 if t > 50 else 2.0
+simple_space = lambda x: 10 - 10 * x[0]
+
+
+@pytest.mark.parametrize(
+    "density_func, expected_value",
+    [
+        (step_function_space, 5.0),
+        (simple_space, 5.0),
+        (step_function_time_non_homogeneous, 5.0),
+        (step_function_time_homogeneous, 10.0),
+    ],
+)
+def test_non_homogeneous_density(density_func, expected_value):
+    my_model = F.HydrogenTransportProblem()
+    my_model.mesh = F.Mesh1D(np.linspace(0, 1, 100))
+    my_mat = F.Material(name="mat", D_0=1, E_D=0)
+    vol = F.VolumeSubdomain1D(id=1, borders=[0, 1], material=my_mat)
+    left = F.SurfaceSubdomain1D(id=1, x=0)
+    right = F.SurfaceSubdomain1D(id=2, x=1)
+
+    my_model.subdomains = [vol, left, right]
+
+    H = F.Species("H")
+    trapped_H = F.Species("trapped_H", mobile=False)
+
+    empty = F.ImplicitSpecies(n=density_func, name="empty", others=[trapped_H])
+    my_model.species = [H, trapped_H]
+
+    my_model.reactions = [
+        F.Reaction(
+            reactant=[H, empty],
+            product=trapped_H,
+            k_0=1,
+            E_k=0,
+            p_0=0,
+            E_p=0,
+            volume=vol,
+        )
+    ]
+
+    my_model.temperature = 600
+
+    my_model.boundary_conditions = [
+        F.DirichletBC(subdomain=left, value=1, species=H),
+        F.DirichletBC(subdomain=right, value=1, species=H),
+    ]
+
+    my_model.settings = F.Settings(atol=1e-10, rtol=1e-10, final_time=100)
+    my_model.settings.stepsize = 1
+
+    total_trapped = F.TotalVolume(field=trapped_H, volume=vol)
+    my_model.exports = [
+        F.VTXSpeciesExport(filename="trapped_c.bp", field=trapped_H),
+        F.VTXSpeciesExport(filename="c.bp", field=H),
+        total_trapped,
+    ]
+
+    my_model.initialise()
+    my_model.run()
+
+    print(total_trapped.value)
+    print(expected_value)
+    assert np.isclose(total_trapped.value, expected_value)
+
+
+def test_density_as_function():
+    my_model = F.HydrogenTransportProblem()
+    my_model.mesh = F.Mesh1D(np.linspace(0, 1, 100))
+    my_mat = F.Material(name="mat", D_0=1, E_D=0)
+    vol = F.VolumeSubdomain1D(id=1, borders=[0, 1], material=my_mat)
+    left = F.SurfaceSubdomain1D(id=1, x=0)
+    right = F.SurfaceSubdomain1D(id=2, x=1)
+
+    my_model.subdomains = [vol, left, right]
+
+    H = F.Species("H")
+    trapped_H = F.Species("trapped_H", mobile=False)
+
+    degree = 1
+    element_CG = basix.ufl.element(
+        basix.ElementFamily.P,
+        my_model.mesh.mesh.basix_cell(),
+        degree,
+        basix.LagrangeVariant.equispaced,
+    )
+
+    density = fem.Function(fem.functionspace(my_model.mesh.mesh, element_CG))
+    density.interpolate(lambda x: 10 - 10 * x[0])
+
+    empty = F.ImplicitSpecies(n=density, name="empty", others=[trapped_H])
+    my_model.species = [H, trapped_H]
+
+    my_model.reactions = [
+        F.Reaction(
+            reactant=[H, empty],
+            product=trapped_H,
+            k_0=1,
+            E_k=0,
+            p_0=0,
+            E_p=0,
+            volume=vol,
+        )
+    ]
+
+    my_model.temperature = 600
+
+    my_model.boundary_conditions = [
+        F.DirichletBC(subdomain=left, value=1, species=H),
+        F.DirichletBC(subdomain=right, value=1, species=H),
+    ]
+
+    my_model.settings = F.Settings(atol=1e-10, rtol=1e-10, final_time=100)
+    my_model.settings.stepsize = 1
+
+    total_trapped = F.TotalVolume(field=trapped_H, volume=vol)
+    my_model.exports = [
+        F.VTXSpeciesExport(filename="trapped_c.bp", field=trapped_H),
+        F.VTXSpeciesExport(filename="c.bp", field=H),
+        total_trapped,
+    ]
+
+    my_model.initialise()
+    my_model.run()
+
+    expected_value = 5.0
+
+    assert np.isclose(total_trapped.value, expected_value)

test/test_species.py

@@ -90,8 +90,10 @@ def test_implicit_species_concentration():
     species1.solution = Function(V)
     species2.solution = Function(V)
 
+    implicit_species.create_value_fenics(mesh, t=0.0)
+
     # test the concentration of the implicit species
-    expected_concentration = implicit_species.n - (
+    expected_concentration = implicit_species.value_fenics - (
         species1.solution + species2.solution
     )
     assert implicit_species.concentration == expected_concentration
@@ -143,3 +145,18 @@ def test_create_species_and_reaction():
     # TEST
     assert isinstance(my_trap.trapped_concentration, F.Species)
     assert isinstance(my_trap.reaction, F.Reaction)
+
+
+def test_implicit_species_wrong_type():
+    """Test that a ValueError is raised when the value of n in an ImplicitSpecies is
+    a function of time only but doesn't return a float or an int.
+    """
+    mesh = create_unit_cube(MPI.COMM_WORLD, 10, 10, 10)
+
+    density = lambda t: "coucou"
+    species = F.ImplicitSpecies(n=density)
+
+    with pytest.raises(
+        ValueError, match="self.value should return a float or an int, not "
+    ):
+        species.create_value_fenics(mesh=mesh, t=0.0)