Merge pull request #1017 from festim-dev/is-it-time

Add `times` argument to ProfileExport

Author: RemDelaporteMathurin

src/festim/exports/profile_1d.py

@@ -10,27 +10,41 @@ class Profile1DExport:
         field: the species for which the profile is computed.
         subdomain: the volume subdomain to compute the profile on. If None, the profile
             is computed over the entire domain.
+        times: if provided, the profile will be exported at these timesteps.
+            Otherwise, exports at all timesteps. Defaults to None.
 
     Attributes:
         field: the species for which the profile is computed.
         subdomain: the volume subdomain to compute the profile on. If None, the profile
             is computed over the entire domain.
+        times: if provided, the profile will be exported at these timesteps. Otherwise,
+            exports at all timesteps.
         x: the coordinates along which the profile is computed.
         data: the computed profile data.
+        t: the list of time values at which the profile is computed.
     """
 
     x: np.ndarray
     data: list
     field: Species
     subdomain: VolumeSubdomain | None
+    times: list[float] | None
+    t: list[float]
     _dofs: np.ndarray
     _sort_coords: np.ndarray
 
-    def __init__(self, field: Species, subdomain: VolumeSubdomain = None):
+    def __init__(
+        self,
+        field: Species,
+        subdomain: VolumeSubdomain = None,
+        times: list[float] | None = None,
+    ):
         self.field = field
         self.data = []
+        self.t = []
         self.x = None
         self.subdomain = subdomain
+        self.times = times
 
         self._dofs = None
         self._sort_coords = None

src/festim/exports/vtx.py

@@ -49,27 +49,6 @@ def __init__(
     def filename(self):
         return self._filename
 
-    def is_it_time_to_export(self, current_time: float) -> bool:
-        """
-        Checks if the exported field should be written to a file or not based on the
-        current time and the times in `export.times`
-
-        Args:
-            current_time: the current simulation time
-
-        Returns:
-            bool: True if the exported field should be written to a file, else False
-        """
-
-        if self.times is None:
-            return True
-
-        for time in self.times:
-            if np.isclose(time, current_time, atol=0):
-                return True
-
-        return False
-
 
 class VTXTemperatureExport(ExportBaseClass):
     """Export temperature field functions to VTX file

src/festim/heat_transfer_problem.py

@@ -236,6 +236,12 @@ def post_processing(self):
         """Post processes the model"""
 
         for export in self.exports:
+            # skip if it isn't time to export
+            if hasattr(export, "times"):
+                if not helpers.is_it_time_to_export(
+                    current_time=float(self.t), times=export.times
+                ):
+                    continue
             # TODO if export type derived quantity
             if isinstance(export, exports.SurfaceQuantity):
                 raise NotImplementedError(
@@ -256,8 +262,7 @@ def post_processing(self):
                 export.write(float(self.t))
 
             if isinstance(export, exports.VTXTemperatureExport):
-                if export.is_it_time_to_export(float(self.t)):
-                    export.writer.write(float(self.t))
+                export.writer.write(float(self.t))
 
     def __del__(self):
         for export in self.exports:

src/festim/helpers.py

@@ -317,3 +317,29 @@ def nmm_interpolate(
         f_out.function_space, f_in.function_space, cells, padding=padding
     )
     f_out.interpolate_nonmatching(f_in, cells, interpolation_data=interpolation_data)
+
+
+def is_it_time_to_export(
+    times: list | None, current_time: float, atol=0, rtol=1.0e-5
+) -> bool:
+    """
+    Checks if the exported field should be written to a file or not based on the
+    current time and the times in `export.times`
+
+    Args:
+        current_time: the current simulation time
+        atol: absolute tolerance for time comparison
+        rtol: relative tolerance for time comparison
+        times: the times at which the field should be exported, if None, returns True
+
+    Returns:
+        bool: True if the exported field should be written to a file, else False
+    """
+    if times is None:
+        return True
+
+    for time in times:
+        if np.isclose(time, current_time, atol=atol, rtol=rtol):
+            return True
+
+    return False

src/festim/hydrogen_transport_problem.py

@@ -33,7 +33,11 @@
     subdomain as _subdomain,
 )
 from festim.advection import AdvectionTerm
-from festim.helpers import as_fenics_constant, get_interpolation_points
+from festim.helpers import (
+    as_fenics_constant,
+    get_interpolation_points,
+    is_it_time_to_export,
+)
 from festim.mesh import Mesh
 
 __all__ = ["HydrogenTransportProblemDiscontinuous", "HydrogenTransportProblem"]
@@ -916,28 +920,34 @@ def post_processing(self):
                         species_not_updated.remove(export.field)
 
         for export in self.exports:
+            # skip if it isn't time to export
+            if hasattr(export, "times"):
+                if not is_it_time_to_export(
+                    current_time=float(self.t), times=export.times
+                ):
+                    continue
+
             # handle VTX exports
             if isinstance(export, exports.ExportBaseClass):
-                if export.is_it_time_to_export(float(self.t)):
-                    if isinstance(export, exports.VTXSpeciesExport):
-                        if export._checkpoint:
-                            for field in export.field:
-                                adios4dolfinx.write_function(
-                                    export.filename,
-                                    field.post_processing_solution,
-                                    time=float(self.t),
-                                    name=field.name,
-                                )
-                        else:
-                            export.writer.write(float(self.t))
-                    elif (
-                        isinstance(export, festim.VTXTemperatureExport)
-                        and self.temperature_time_dependent
-                    ):
-                        self._temperature_as_function.interpolate(
-                            self._get_temperature_field_as_function()
-                        )
+                if isinstance(export, exports.VTXSpeciesExport):
+                    if export._checkpoint:
+                        for field in export.field:
+                            adios4dolfinx.write_function(
+                                export.filename,
+                                field.post_processing_solution,
+                                time=float(self.t),
+                                name=field.name,
+                            )
+                    else:
                         export.writer.write(float(self.t))
+                elif (
+                    isinstance(export, festim.VTXTemperatureExport)
+                    and self.temperature_time_dependent
+                ):
+                    self._temperature_as_function.interpolate(
+                        self._get_temperature_field_as_function()
+                    )
+                    export.writer.write(float(self.t))
 
             # TODO if export type derived quantity
             if isinstance(export, exports.SurfaceQuantity):
@@ -1000,6 +1010,7 @@ def post_processing(self):
                 else:
                     c = self.u.x.array[export._dofs]
                 export.data.append(c)
+                export.t.append(float(self.t))
 
 
 class HydrogenTransportProblemDiscontinuous(HydrogenTransportProblem):
@@ -1593,6 +1604,12 @@ def post_processing(self):
                 collapsed_function.x.array[:] = u.x.array[v0_to_V]
 
         for export in self.exports:
+            # skip if it isn't time to export
+            if hasattr(export, "times"):
+                if not is_it_time_to_export(
+                    current_time=float(self.t), times=export.times
+                ):
+                    continue
             # handle VTX exports
             if isinstance(export, exports.ExportBaseClass):
                 if not isinstance(export, exports.VTXSpeciesExport):
@@ -1605,8 +1622,7 @@ def post_processing(self):
                             f"Export type {type(export)} not implemented "
                             f"for mixed-domain approach"
                         )
-                if export.is_it_time_to_export(float(self.t)):
-                    export.writer.write(float(self.t))
+                export.writer.write(float(self.t))
 
             # handle derived quantities
             if isinstance(export, exports.SurfaceQuantity):
@@ -1676,6 +1692,7 @@ def post_processing(self):
                 c = u.x.array[export._dofs][export._sort_coords]
 
                 export.data.append(c)
+                export.t.append(float(self.t))
 
     def iterate(self):
         """Iterates the model for a given time step"""

test/test_helpers.py

@@ -297,3 +297,28 @@ def example_func(t):
         match=f"A time dependent advection field should return an fem.Function, not a <class 'ufl.algebra.Product'>",
     ):
         test_value.convert_input_value(function_space=test_function_space, t=t)
+
+
+@pytest.mark.parametrize(
+    "input, expected_output",
+    [
+        (2, True),
+        (3, True),
+        (1, True),
+        (-1, False),
+        (0, False),
+        (5, False),
+        (1.5, False),
+    ],
+)
+def test_is_it_time_to_export(input, expected_output):
+    times = [1, 2, 3]
+    assert (
+        F.helpers.is_it_time_to_export(current_time=input, times=times)
+        == expected_output
+    )
+
+
+def test_is_it_time_to_export_when_times_not_given():
+    times = None
+    assert F.helpers.is_it_time_to_export(current_time=1.0, times=times)

test/test_profile.py

@@ -1,8 +1,10 @@
 import festim as F
 import numpy as np
+import pytest
 
 
-def test_profile():
+@pytest.mark.parametrize("times", [None, [2, 5], [5]])
+def test_profile(times):
     my_model = F.HydrogenTransportProblem()
 
     protium = F.Species("H")
@@ -41,17 +43,21 @@ def test_profile():
     my_model.settings.stepsize = F.Stepsize(1)
 
     my_model.exports = [
-        F.Profile1DExport(protium),
-        F.Profile1DExport(deuterium),
+        F.Profile1DExport(protium, times=times),
+        F.Profile1DExport(deuterium, times=times),
     ]
 
     my_model.initialise()
     my_model.run()
 
     assert my_model.exports[0].x is not None
     assert my_model.exports[1].x is not None
-    assert len(my_model.exports[0].data) > 0
-    assert len(my_model.exports[1].data) > 0
+    if times is None:
+        assert len(my_model.exports[0].data) > 0
+        assert len(my_model.exports[1].data) > 0
+    else:
+        assert len(my_model.exports[0].data) == len(times)
+        assert len(my_model.exports[1].data) == len(times)
 
 
 def test_profile_single_species():

test/test_vtx.py

@@ -179,29 +179,3 @@ 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)
-
-
-@pytest.mark.parametrize(
-    "input, expected_output",
-    [
-        (2, True),
-        (3, True),
-        (1, True),
-        (-1, False),
-        (0, False),
-        (5, False),
-        (1.5, False),
-    ],
-)
-def test_is_it_time_to_export(tmpdir, input, expected_output):
-    filename = str(tmpdir.join("my_T_export.bp"))
-    my_export = F.ExportBaseClass(times=[1, 2, 3], ext=".bp", filename=filename)
-
-    assert my_export.is_it_time_to_export(input) == expected_output
-
-
-def test_is_it_time_to_export_when_times_not_given(tmpdir):
-    filename = str(tmpdir.join("my_T_export.bp"))
-    my_export = F.ExportBaseClass(ext=".bp", filename=filename)
-
-    assert my_export.is_it_time_to_export(1.0)