Custom Field export

src/festim/__init__.py

@@ -35,7 +35,12 @@
 from .exports.total_surface import TotalSurface
 from .exports.total_volume import TotalVolume
 from .exports.volume_quantity import VolumeQuantity
-from .exports.vtx import ExportBaseClass, VTXSpeciesExport, VTXTemperatureExport
+from .exports.vtx import (
+    ExportBaseClass,
+    VTXSpeciesExport,
+    VTXTemperatureExport,
+    CustomField,
+)
 from .exports.xdmf import XDMFExport
 from .heat_transfer_problem import HeatTransferProblem
 from .helpers import (

src/festim/exports/__init__.py

@@ -10,12 +10,13 @@
 from .total_surface import TotalSurface
 from .total_volume import TotalVolume
 from .volume_quantity import VolumeQuantity
-from .vtx import ExportBaseClass, VTXSpeciesExport, VTXTemperatureExport
+from .vtx import ExportBaseClass, VTXSpeciesExport, VTXTemperatureExport, CustomField
 from .xdmf import XDMFExport
 
 __all__ = [
     "AverageSurface",
     "AverageVolume",
+    "CustomField",
     "ExportBaseClass",
     "MaximumSurface",
     "MaximumVolume",

src/festim/exports/vtx.py

@@ -1,9 +1,14 @@
+import inspect
 import warnings
+from collections.abc import Callable
 from pathlib import Path
+from typing import Union
 
+import ufl
 from dolfinx import fem, io
 
-from festim.species import Species
+from festim.helpers import get_interpolation_points
+from festim.species import Species, ImplicitSpecies
 from festim.subdomain.volume_subdomain import VolumeSubdomain
 
 
@@ -172,3 +177,150 @@ def get_functions(self) -> list[fem.Function]:
                             field.subdomain_to_post_processing_solution[self._subdomain]
                         )
                 return outfiles
+
+
+class CustomField(ExportBaseClass):
+    """Export a custom field to a VTX file
+
+    Args:
+        filename: The name of the output file
+        expression: A function evaluating the custom field. Positional
+            arguments of the function can be "t" (time), "x" (spatial coordinate),
+            "T" (temperature), or any key from the `species_dependent_value` dictionary.
+        species_dependent_value: A dictionary mapping argument names
+            in `expression` to Species objects. Defaults to None.
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps. Defaults to None.
+        subdomain: The volume subdomain on which the custom
+            field is evaluated. Defaults to None.
+        checkpoint: If True, the export will be a checkpoint file using
+            adios4dolfinx and won't be readable by ParaView. Default is False.
+
+    Attributes:
+        filename: The name of the output file
+        expression: A function evaluating the custom field.
+        species_dependent_value: A dictionary mapping argument names to Species objects.
+        subdomain: The volume subdomain on which the custom field is evaluated.
+        checkpoint: If True, the export will be a checkpoint file.
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps.
+        function: the function containing the custom field values
+        writer: The VTXWriter object used to write the file
+        dolfinx_expression: the dolfinx expression used to evaluate the function
+    """
+
+    function: fem.Function
+    writer: io.VTXWriter
+    dolfinx_expression: fem.Expression
+    expression: Callable
+    species_dependent_value: dict[str, Species]
+    subdomain: VolumeSubdomain
+    checkpoint: bool
+
+    def __init__(
+        self,
+        filename: Union[str, Path],
+        expression: Callable,
+        species_dependent_value: Union[dict[str, Species], None] = None,
+        times: Union[list[float], list[int], None] = None,
+        subdomain: VolumeSubdomain = None,
+        checkpoint: bool = False,
+    ):
+        super().__init__(
+            filename=filename,
+            times=times,
+            ext=".bp",
+        )
+        self.expression = expression
+        self.species_dependent_value = species_dependent_value or {}
+        self.checkpoint = checkpoint
+        self.subdomain = subdomain
+
+    def set_dolfinx_expression(
+        self,
+        temperature: fem.Constant | fem.Function,
+        time: fem.Constant,
+    ):
+        """
+        Set the dolfinx expression used to evaluate the custom field. This is done by
+        evaluating the user-provided expression with the appropriate arguments and using
+        the result to create a dolfinx expression.
+
+        Args:
+            temperature: The temperature field to use in the expression
+            time: The time to use in the expression
+        """
+        # check if we are in a mixed domain/discontinuous case
+        mixed_domain = any(
+            spe.subdomain_to_post_processing_solution
+            for spe in self.species_dependent_value.values()
+        ) or (self.subdomain.sub_T if self.subdomain else None)
+
+        # get the arguments of the user-provided expression
+        arguments = inspect.signature(self.expression).parameters
+
+        # create a dictionary mapping the arguments to the appropriate values
+        kwargs = {}
+        if "t" in arguments:
+            kwargs["t"] = time
+        if "x" in arguments:
+            x = ufl.SpatialCoordinate(self.function.function_space.mesh)
+            kwargs["x"] = x
+        if "T" in arguments:
+            if isinstance(temperature, fem.Function) and mixed_domain:
+                # fem.Function in mixed domain/discontinuous case, use sub_T
+                # NOTE I'm not sure that sub_T is updated at every time step
+                kwargs["T"] = self.subdomain.sub_T
+            else:
+                # else use the provided temperature
+                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]
+                if isinstance(spe, ImplicitSpecies):
+                    raise NotImplementedError(
+                        "Custom fields depending on implicit species are not"
+                        "implemented yet."
+                    )
+                if mixed_domain:
+                    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.check_valid_inputs(kwargs, mixed_domain)
+
+        # evaluate the user-provided expression with the appropriate arguments and create a
+        # dolfinx.fem.Expression
+        self.dolfinx_expression = fem.Expression(
+            self.expression(**kwargs),
+            get_interpolation_points(self.function.function_space.element),
+        )
+
+    def check_valid_inputs(self, kwargs: dict, mixed_domain: bool):
+        """
+        Check if we are in the mixed domain/discontinuous case and if the user-provided
+        expression is valid in this case.
+        dolfinx.fem.Expression does not support co-dim 0 submeshes and time is defined
+        on the parent mesh, so we cannot have time-dependent custom fields in the mixed
+        domain/discontinuous case.
+
+        When https://github.com/FEniCS/dolfinx/issues/3207 is resolved we should be
+        able to support this
+        """
+
+        # check the domain of all kwargs and check that they are the same
+
+        if mixed_domain and "t" in kwargs:
+            raise NotImplementedError(
+                "Time-dependent custom fields are not implemented in the case of a "
+                "mixed domain/discontinuous case."
+                "dolfinx.fem.Expression does not support co-dim 0 submeshes and time is"
+                "defined on the parent mesh."
+                "See https://github.com/FEniCS/dolfinx/issues/3207 for more details."
+            )

src/festim/helpers.py

@@ -34,6 +34,7 @@ def as_fenics_constant(
         )
 
 
+# TODO change this to accept species dependent values
 def as_mapped_function(
     value: Callable,
     function_space: fem.FunctionSpace | None = None,
@@ -68,6 +69,7 @@ def as_mapped_function(
     return value(**kwargs)
 
 
+# TODO change this to accept species dependent values
 def as_fenics_interp_expr_and_function(
     value: Callable,
     function_space: dolfinx.fem.function.FunctionSpace,

src/festim/hydrogen_transport_problem.py

@@ -449,6 +449,21 @@ def initialise_exports(self):
                     else:
                         adios4dolfinx.write_mesh(export.filename, mesh=self.mesh.mesh)
 
+                elif isinstance(export, exports.CustomField):
+                    export.function = fem.Function(self.V_CG_1)
+                    export.set_dolfinx_expression(
+                        temperature=self.temperature_fenics,
+                        time=self.t,
+                    )
+
+                    export.writer = dolfinx.io.VTXWriter(
+                        comm=export.function.function_space.mesh.comm,
+                        filename=export.filename,
+                        output=export.function,
+                        engine="BP5",
+                    )
+                    continue
+
             elif isinstance(export, exports.SurfaceQuantity | exports.VolumeQuantity):
                 # raise not implemented error if the derived quantity don't match the
                 # type of mesh eg. SurfaceFlux is used with cylindrical mesh
@@ -629,6 +644,7 @@ def define_function_spaces(self, element_degree: int = 1):
         )
         self.V_DG_0 = fem.functionspace(self.mesh.mesh, element_DG0)
         self.V_DG_1 = fem.functionspace(self.mesh.mesh, element_DG1)
+        self.V_CG_1 = fem.functionspace(self.mesh.mesh, ("CG", 1))
 
         self.u = fem.Function(self.function_space)
         self.u_n = fem.Function(self.function_space)
@@ -646,6 +662,7 @@ def assign_functions_to_species(self):
             spe.sub_function_space = self.function_space.sub(idx)
             spe.sub_function = self.u.sub(idx)  # TODO add this to discontinuous class
             spe.post_processing_solution = self.u.sub(idx).collapse()
+            spe.post_processing_solution.name = spe.name
             spe.collapsed_function_space, spe.map_sub_to_main_solution = (
                 self.function_space.sub(idx).collapse()
             )
@@ -981,6 +998,10 @@ def post_processing(self):
                         self._get_temperature_field_as_function()
                     )
                     export.writer.write(float(self.t))
+                elif isinstance(export, exports.CustomField):
+                    # update internal function
+                    export.function.interpolate(export.dolfinx_expression)
+                    export.writer.write(float(self.t))
 
             # TODO if export type derived quantity
             if isinstance(export, exports.SurfaceQuantity):
@@ -1095,6 +1116,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):
@@ -1335,6 +1357,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
@@ -1662,7 +1688,21 @@ 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
+                    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
@@ -1704,11 +1744,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 = (
@@ -1724,7 +1764,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(