Add VTX Temperature Export (with rebased branch)

src/festim/hydrogen_transport_problem.py

@@ -5,6 +5,7 @@
 import basix
 import dolfinx
 import numpy.typing as npt
+import numpy as np
 import tqdm.autonotebook
 import ufl
 from dolfinx import fem
@@ -119,6 +120,8 @@ class HydrogenTransportProblem(problem.ProblemBase):
 
     """
 
+    _temperature_as_function: fem.Function
+
     def __init__(
         self,
         mesh: Mesh | None = None,
@@ -171,6 +174,8 @@ def __init__(
         self._element_for_traps = "DG"
         self.petcs_options = petsc_options
 
+        self._temperature_as_function = None
+
     @property
     def temperature(self):
         return self._temperature
@@ -370,6 +375,20 @@ def initialise_exports(self):
         a string, find species object in self.species"""
 
         for export in self.exports:
+            if isinstance(export, festim.VTXTemperatureExport):
+                self._temperature_as_function = (
+                    self._get_temperature_field_as_function()
+                )
+                self._vtxfiles.append(
+                    dolfinx.io.VTXWriter(
+                        self._temperature_as_function.function_space.mesh.comm,
+                        export.filename,
+                        self._temperature_as_function,
+                        engine="BP5",
+                    )
+                )
+                continue
+
             # if name of species is given then replace with species object
             if isinstance(export.field, list):
                 for idx, field in enumerate(export.field):
@@ -395,12 +414,14 @@ def initialise_exports(self):
                         engine="BP5",
                     )
                 )
+
         # compute diffusivity function for surface fluxes
 
         spe_to_D_global = {}  # links species to global D function
         spe_to_D_global_expr = {}  # links species to D expression
 
         for export in self.exports:
+
             if isinstance(export, exports.SurfaceQuantity):
                 if export.field in spe_to_D_global:
                     # if already computed then use the same D
@@ -421,6 +442,30 @@ def initialise_exports(self):
                 export.t = []
                 export.data = []
 
+    def _get_temperature_field_as_function(self) -> dolfinx.fem.Function:
+        """
+        Based on the type of the temperature_fenics attribute, converts
+        it as a Function to be used in VTX export
+
+        Returns:
+            the temperature field of the simulation
+        """
+        if isinstance(self.temperature_fenics, fem.Function):
+            return self.temperature_fenics
+        elif isinstance(self.temperature_fenics, fem.Constant):
+            # use existing function space if function already exists
+            if self._temperature_as_function is None:
+                V = dolfinx.fem.functionspace(self.mesh.mesh, ("P", 1))
+            else:
+                V = self._temperature_as_function.function_space
+            temperature_field = dolfinx.fem.Function(V)
+            temperature_expr = fem.Expression(
+                self.temperature_fenics,
+                get_interpolation_points(V.element),
+            )
+            temperature_field.interpolate(temperature_expr)
+            return temperature_field
+
     def define_D_global(self, species):
         """Defines the global diffusion coefficient for a given species
 
@@ -803,6 +848,13 @@ def post_processing(self):
                         species_not_updated.remove(export.field)
 
         for export in self.exports:
+            if (
+                isinstance(export, festim.VTXTemperatureExport)
+                and self.temperature_time_dependent
+            ):
+                self._temperature_as_function.interpolate(
+                    self._get_temperature_field_as_function()
+                )
             # TODO if export type derived quantity
             if isinstance(export, exports.SurfaceQuantity):
                 if isinstance(

test/test_vtx.py

@@ -97,7 +97,7 @@ def test_vtx_integration_with_h_transport_problem(tmpdir):
         F.SurfaceSubdomain1D(2, x=4.0),
     ]
     my_model.species = [F.Species("H")]
-    my_model.temperature = 500
+    my_model.temperature = lambda t: 500 + t
 
     filename = str(tmpdir.join("my_export.bp"))
     my_export = F.VTXSpeciesExport(filename, field=my_model.species[0])
@@ -110,6 +110,41 @@ def test_vtx_integration_with_h_transport_problem(tmpdir):
     assert len(my_model._vtxfiles) == 1
 
 
+@pytest.mark.parametrize(
+    "T",
+    [
+        500,
+        lambda t: 500 + t,
+        lambda x: 500 + 200 * x[0],
+        lambda x, t: 500 + 200 * x[0] + t,
+    ],
+)
+def test_vtx_temperature(T, tmpdir):
+    """Tests that VTX temperature exports work with HydrogenTransportProblem"""
+    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")
+    my_model.subdomains = [
+        F.VolumeSubdomain1D(1, borders=[0.0, 4.0], material=my_mat),
+        F.SurfaceSubdomain1D(1, x=0.0),
+        F.SurfaceSubdomain1D(2, x=4.0),
+    ]
+    my_model.species = [F.Species("H")]
+    my_model.temperature = T
+
+    filename = str(tmpdir.join("my_export.bp"))
+
+    my_export = F.VTXTemperatureExport(filename)
+    my_model.exports = [my_export]
+    my_model.settings = F.Settings(atol=1, rtol=0.1, final_time=2)
+    my_model.settings.stepsize = F.Stepsize(initial_value=1)
+
+    my_model.initialise()
+    assert len(my_model._vtxfiles) == 1
+
+    my_model.run()
+
+
 def test_field_attribute_is_always_list():
     """Test that the field attribute is always a list"""
     my_export = F.VTXSpeciesExport("my_export.bp", field=F.Species("H"))
@@ -130,3 +165,9 @@ def test_filename_raises_error_when_wrong_type():
     """Test that the filename attribute raises an error if the extension is not .bp"""
     with pytest.raises(TypeError):
         F.VTXSpeciesExport(1, field=[F.Species("H")])
+
+
+def test_filename_temp_raises_error_when_wrong_type():
+    """Test that the filename attribute for VTXTemperature export raises an error if the extension is not .bp"""
+    with pytest.raises(TypeError):
+        F.VTXTemperatureExport(1)