support for mixed domain

Author: RemDelaporteMathurin

src/festim/exports/vtx.py

@@ -216,14 +216,28 @@ def set_dolfinx_expression(
             x = ufl.SpatialCoordinate(time._ufl_domain)
             kwargs["x"] = x
         if "T" in arguments:
-            kwargs["T"] = temperature
+            if (
+                isinstance(temperature, fem.Function)
+                and self.subdomain.sub_T is not None
+            ):
+                kwargs["T"] = (
+                    self.subdomain.sub_T
+                )  # NOTE I'm not sure that sub_T is updated at every time step
+            else:
+                kwargs["T"] = temperature
         # check if there are other arguments and if they are in species_dependent_value
         for arg in arguments:
             if arg in self.species_dependent_value:
                 spe = self.species_dependent_value[arg]
-                kwargs[arg] = spe.post_processing_solution
-                # TODO handle discontinuous case and use self.subdomain
-
+                if spe.subdomain_to_post_processing_solution:
+                    kwargs[arg] = spe.subdomain_to_post_processing_solution[
+                        self.subdomain
+                    ]
+                else:
+                    kwargs[arg] = spe.post_processing_solution
+            assert kwargs[arg] is not None, (
+                f"Argument {arg} not found in species_dependent_value"
+            )
         self.dolfinx_expression = fem.Expression(
             self.expression(**kwargs),
             get_interpolation_points(self.function.function_space.element),

src/festim/hydrogen_transport_problem.py

@@ -1117,6 +1117,7 @@ def __init__(
         self.interfaces = interfaces or []
         self.surface_to_volume = surface_to_volume or {}
         self.subdomain_to_species = {}  # maps subdomain to species defined in it
+        self.subdomain_to_V_CG1 = {}
 
     @property
     def method_interface(self):
@@ -1357,6 +1358,10 @@ def define_function_spaces(
         u = dolfinx.fem.Function(V)
         u_n = dolfinx.fem.Function(V)
 
+        self.subdomain_to_V_CG1[subdomain] = dolfinx.fem.functionspace(
+            subdomain.submesh, ("CG", 1)
+        )
+
         # store attributes in the subdomain object
         subdomain.u = u
         subdomain.u_n = u_n
@@ -1684,7 +1689,22 @@ def initialise_exports(self):
                         export.filename,
                         mesh=functions[0].function_space.mesh,
                     )
+            elif isinstance(export, exports.CustomField):
+                # need to find an appropriate function space on the right submesh
+                V = self.subdomain_to_V_CG1[export.subdomain]
+                export.function = fem.Function(V)
+                export.set_dolfinx_expression(
+                    temperature=self.temperature_fenics,  # need to pass the right temperature
+                    species=self.species,
+                    time=self.t,
+                )
 
+                export.writer = dolfinx.io.VTXWriter(
+                    comm=export.function.function_space.mesh.comm,
+                    filename=export.filename,
+                    output=export.function,
+                    engine="BP5",
+                )
         # compute diffusivity function for surface fluxes
         # for the discontinuous case, we don't use D_global as in
         # HydrogenTransportProblem
@@ -1726,11 +1746,11 @@ def post_processing(self):
                     continue
             # handle VTX exports
             if isinstance(export, exports.ExportBaseClass):
-                if not isinstance(export, exports.VTXSpeciesExport):
-                    raise NotImplementedError(
-                        f"Export type {type(export)} not implemented"
-                    )
-                if isinstance(export, exports.VTXSpeciesExport):
+                if isinstance(export, exports.CustomField):
+                    # update internal function
+                    export.function.interpolate(export.dolfinx_expression)
+                    export.writer.write(float(self.t))
+                elif isinstance(export, exports.VTXSpeciesExport):
                     if export._checkpoint:
                         for species in export.field:
                             post_processing_solution = (
@@ -1746,7 +1766,10 @@ def post_processing(self):
                             )
                     else:
                         export.writer.write(float(self.t))
-
+                else:
+                    raise NotImplementedError(
+                        f"Export type {type(export)} not implemented"
+                    )
             # handle derived quantities
             if isinstance(export, exports.SurfaceQuantity):
                 if isinstance(

test/test_vtx.py

@@ -260,3 +260,83 @@ def test_custom_field(tmp_path, expression):
     my_model.initialise()
 
     my_model.run()
+
+
+@pytest.mark.parametrize(
+    "expression",
+    [
+        lambda x: x[0] + x[1] * 2,
+        lambda T: T + 1,
+        lambda c_A, c_B: c_A + c_B,
+        lambda c_A, T: c_A * T,
+        # FIXME: the following don't work (easily) because of https://github.com/FEniCS/dolfinx/issues/3207
+        # lambda c_A, c_B, x: c_A * c_B + x[0],
+        # lambda c_A, T, x: c_A * T + x[0],
+        # lambda T, x: T + x[0],
+    ],
+)
+def test_custom_field_discontinuous(tmp_path, expression):
+    """
+    Test custom field export functionality.
+    This test checks that a custom field can be created with various types of
+    expressions.
+    """
+
+    my_model = F.HydrogenTransportProblemDiscontinuous()
+
+    mat = F.Material(D_0=1, E_D=0, K_S_0=1, E_K_S=0)
+
+    vol = F.VolumeSubdomain(id=1, material=mat)
+
+    top = F.SurfaceSubdomain(id=1, locator=lambda x: np.isclose(x[1], 1))
+    bottom = F.SurfaceSubdomain(id=2, locator=lambda x: np.isclose(x[1], 0))
+    left = F.SurfaceSubdomain(id=3, locator=lambda x: np.isclose(x[0], 0))
+    right = F.SurfaceSubdomain(id=4, locator=lambda x: np.isclose(x[0], 1))
+
+    my_model.subdomains = [vol, top, bottom, left, right]
+
+    my_model.surface_to_volume = {top: vol, bottom: vol, left: vol, right: vol}
+
+    dolfinx_mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
+    my_model.mesh = F.Mesh(dolfinx_mesh)
+
+    A = F.Species("A", subdomains=my_model.volume_subdomains)
+    B = F.Species("B", subdomains=my_model.volume_subdomains)
+    C = F.Species("C", subdomains=my_model.volume_subdomains)
+    D = F.Species("D", subdomains=my_model.volume_subdomains)
+
+    my_model.species = [A, B, C, D]
+
+    my_model.boundary_conditions = (
+        [
+            F.FixedConcentrationBC(species=A, subdomain=top, value=1),
+            F.FixedConcentrationBC(species=B, subdomain=left, value=1),
+        ]
+        + [
+            F.FixedConcentrationBC(species=C, subdomain=surf, value=0)
+            for surf in [top, bottom, left, right]
+        ]
+        + [
+            F.FixedConcentrationBC(species=D, subdomain=surf, value=0)
+            for surf in [top, bottom, left, right]
+        ]
+    )
+
+    my_model.temperature = lambda x: 300 + 100 * x[0]
+
+    my_model.settings = F.Settings(transient=False, atol=1e-9, rtol=1e-9)
+
+    custom_field = F.CustomField(
+        filename=tmp_path / "custom_field.bp",
+        expression=expression,
+        species_dependent_value={"c_A": A, "c_B": B},
+        subdomain=vol,
+    )
+
+    my_model.exports = [
+        custom_field,
+    ]
+
+    my_model.initialise()
+
+    my_model.run()