Allow D to be func for multi vol subdomains

src/festim/hydrogen_transport_problem.py

@@ -546,11 +546,11 @@ def define_D_global(self, species):
 
         # if diffusion coeffient has been given as a function, use that
         if self.volume_subdomains[0].material.D:
-            if len(self.volume_subdomains) > 1:
-                raise NotImplementedError(
-                    "Giving the diffusion coefficient as a function is currently "
-                    "only supported for a single volume subdomain case"
-                )
+            # if len(self.volume_subdomains) > 1:
+            #     raise NotImplementedError(
+            #         "Giving the diffusion coefficient as a function is currently "
+            #         "only supported for a single volume subdomain case"
+            #     )
             return self.volume_subdomains[0].material.D, None
 
         D_0 = fem.Function(self.V_DG_0)
@@ -569,6 +569,33 @@ def define_D_global(self, species):
         D.interpolate(D_expr)
         return D, D_expr
 
+    def define_D_mult_subdomains(
+        self,
+        species,
+        volume,
+        mesh=None,
+        temperature=None,
+    ):
+        """Defines the global diffusion coefficient for a given species
+
+        Args:
+            species (F.Species): the species
+            volume (F.VolumeSubdomain): the volume
+
+        Returns:
+            dolfinx.fem.Function, dolfinx.fem.Expression: the global diffusion
+                coefficient and the expression of the global diffusion coefficient
+                for a given species
+        """
+        assert isinstance(species, _species.Species)
+
+        # if diffusion coeffient has been given as a function, use that
+        if volume.material.D:
+            D = volume.material.D
+        else:
+            D = volume.material.get_diffusion_coefficient(mesh, temperature, species)
+        return D
+
     def define_function_spaces(self, element_degree=1):
         """Creates the function space of the model, creates a mixed element if
         model is multispecies. Creates the main solution and previous solution
@@ -1597,13 +1624,17 @@ def mixed_term(u, v, n):
                             interface.id
                         ) - 0.5 * mixed_term(
                             v_b, (u_b / K_b - u_t / K_t), n_0
-                        ) * dInterface(interface.id)
+                        ) * dInterface(
+                            interface.id
+                        )
 
                         F_1 = +0.5 * mixed_term((u_b + u_t), v_t, n_0) * dInterface(
                             interface.id
                         ) - 0.5 * mixed_term(
                             v_t, (u_b / K_b - u_t / K_t), n_0
-                        ) * dInterface(interface.id)
+                        ) * dInterface(
+                            interface.id
+                        )
                         F_0 += (
                             2
                             * gamma
@@ -1706,12 +1737,26 @@ def initialise_exports(self):
         # compute diffusivity function for surface fluxes
         # for the discontinuous case, we don't use D_global as in
         # HydrogenTransportProblem
+
+        vol_to_D_global = {}
+
         for export in self.exports:
             if isinstance(export, exports.SurfaceQuantity):
                 volume = self.surface_to_volume[export.surface]
-                D = volume.material.get_diffusion_coefficient(
-                    self.mesh.mesh, self.temperature_fenics, export.field
-                )
+                if volume.id in vol_to_D_global:
+                    D = vol_to_D_global[volume.id]
+                else:
+                    D = self.define_D_mult_subdomains(
+                        species=export.field,
+                        volume=volume,
+                        mesh=self.mesh.mesh,
+                        temperature=self.temperature_fenics,
+                    )
+                    vol_to_D_global[volume.id] = D
+
+                #    D = volume.material.get_diffusion_coefficient(
+                #         self.mesh.mesh, self.temperature_fenics, export.field
+                #     )
                 # NOTE: maybe we need to make sure there are no functionspace clashes?
 
                 export.D = D
@@ -1819,9 +1864,9 @@ def post_processing(self):
                     export.write(t=float(self.t))
 
             elif isinstance(export, exports.Profile1DExport):
-                assert export.subdomain, (
-                    "Profile1DExport requires a subdomain to be set"
-                )
+                assert (
+                    export.subdomain
+                ), "Profile1DExport requires a subdomain to be set"
                 u = export.subdomain.u
                 if export._dofs is None:
                     index = self.subdomain_to_species[export.subdomain].index(

test/test_h_transport_problem.py

@@ -10,6 +10,8 @@
 
 import festim as F
 
+from festim import k_B
+
 test_mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
 x = ufl.SpatialCoordinate(test_mesh.mesh)
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy_mat")
@@ -1252,19 +1254,59 @@ def test_create_flux_values_fenics_multispecies():
     assert np.isclose(float(my_model.boundary_conditions[1].value_fenics), 11)
 
 
-def test_not_implemented_error_raised_with_D_as_function():
-    """Test that NotImplementedError is raised when D is a function and more than one
-    volume subdomain has been defined"""
+def test_D_as_function_multiple_subdomains():
+    """Test that if a function is given as diffusion coeff of a material for multiple subdomain case,
+    that it is passed to D_global_mult"""
 
     # BUILD
     test_mesh = F.Mesh1D(vertices=np.linspace(0, 1, num=101))
     V = fem.functionspace(test_mesh.mesh, ("Lagrange", 1))
-    D_func = fem.Function(V)
-    D_func.x.array[:] = 7.0
+    D_func1 = fem.Function(V)
+    D_func1.x.array[:] = 7.0
+    D_func2 = fem.Function(V)
+    D_func2.x.array[:] = 1.0
+
+    my_mat1 = F.Material(D=D_func1)
+    my_mat2 = F.Material(D=D_func2)
+    vol_1 = F.VolumeSubdomain1D(id=1, borders=[0, 0.5], material=my_mat1)
+    vol_2 = F.VolumeSubdomain1D(id=1, borders=[0.5, 1], material=my_mat2)
+    H = F.Species("H")
+    my_model = F.HydrogenTransportProblem(
+        mesh=test_mesh,
+        temperature=10,
+        subdomains=[vol_1, vol_2],
+        species=[F.Species("H")],
+    )
 
-    my_mat = F.Material(D=D_func)
-    vol_1 = F.VolumeSubdomain1D(id=1, borders=[0, 0.5], material=my_mat)
-    vol_2 = F.VolumeSubdomain1D(id=1, borders=[0.5, 1], material=my_mat)
+    my_model.define_function_spaces()
+    my_model.assign_functions_to_species()
+    D1 = my_model.define_D_mult_subdomains(
+        species=H,
+        volume=vol_1,
+    )
+    D2 = my_model.define_D_mult_subdomains(
+        species=H,
+        volume=vol_2,
+    )
+
+    assert D1 == D_func1
+    assert D2 == D_func2
+
+
+def test_D_not_function_multiple_subdomains():
+    """Test that if a float or int is given as diffusion coeff of a material for multiple subdomain case,
+    that it is passed to D_global_mult and defined correctly"""
+
+    # BUILD
+    test_mesh = F.Mesh1D(vertices=np.linspace(0, 1, num=101))
+    V = fem.functionspace(test_mesh.mesh, ("Lagrange", 1))
+    D_01 = 5
+    D_02 = 7
+
+    my_mat1 = F.Material(D_0=D_01, E_D=1)
+    my_mat2 = F.Material(D_0=D_02, E_D=1)
+    vol_1 = F.VolumeSubdomain1D(id=1, borders=[0, 0.5], material=my_mat1)
+    vol_2 = F.VolumeSubdomain1D(id=1, borders=[0.5, 1], material=my_mat2)
     H = F.Species("H")
     my_model = F.HydrogenTransportProblem(
         mesh=test_mesh,
@@ -1275,16 +1317,18 @@ def test_not_implemented_error_raised_with_D_as_function():
 
     my_model.define_function_spaces()
     my_model.assign_functions_to_species()
+    D1 = my_model.define_D_mult_subdomains(
+        species=H, volume=vol_1, mesh=test_mesh.mesh, temperature=my_model.temperature
+    )
+    D2 = my_model.define_D_mult_subdomains(
+        species=H, volume=vol_2, mesh=test_mesh.mesh, temperature=my_model.temperature
+    )
 
-    # TEST
-    with pytest.raises(
-        NotImplementedError,
-        match=(
-            "Giving the diffusion coefficient as a function is currently only "
-            "supported for a single volume subdomain case"
-        ),
-    ):
-        my_model.define_D_global(H)
+    D_analytical_1 = D_01 * np.exp(-1 / F.k_B / my_model.temperature)
+    D_analytical_2 = D_02 * np.exp(-1 / F.k_B / my_model.temperature)
+
+    assert np.isclose(float(D1), D_analytical_1)
+    assert np.isclose(float(D2), D_analytical_2)
 
 
 def test_define_D_global_with_D_as_function():