Merge pull request #990 from jhdark/milestones_from_exports

Export at specific times

Author: jhdark

src/festim/__init__.py

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

src/festim/exports/vtx.py

@@ -1,18 +1,37 @@
 import warnings
 from pathlib import Path
+from typing import Optional
 
-from dolfinx import fem
+import numpy as np
+from dolfinx import fem, io
 
-from festim.species import Species as _Species
-from festim.subdomain.volume_subdomain import (
-    VolumeSubdomain as _VolumeSubdomain,
-)
+from festim.species import Species
+from festim.subdomain.volume_subdomain import VolumeSubdomain
 
 
 class ExportBaseClass:
+    """Export functions to VTX file
+
+    Args:
+        filename: The name of the output file
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps. Defaults to None.
+
+    Attributes:
+        filename: The name of the output file
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps. Defaults to None.
+    """
+
     _filename: Path | str
+    writer: io.VTXWriter
 
-    def __init__(self, filename: str | Path, ext: str) -> None:
+    def __init__(
+        self,
+        filename: str | Path,
+        ext: str,
+        times: Optional[list[float] | list[int] | None] = None,
+    ):
         name = Path(filename)
         if name.suffix != ext:
             warnings.warn(
@@ -21,86 +40,141 @@ def __init__(self, filename: str | Path, ext: str) -> None:
             name = name.with_suffix(ext)
 
         self._filename = name
+        if times:
+            self.times = sorted(times)
+        else:
+            self.times = times
 
     @property
     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):
-    def __init__(self, filename: str | Path):
-        super().__init__(filename, ".bp")
+    """Export temperature field functions to VTX file
+
+    Args:
+        filename: The name of the output file
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps. Defaults to None.
+
+    Attributes:
+        filename: The name of the output file
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps. Defaults to None.
+        writer: The VTXWriter object used to write the file
+    """
+
+    writer: io.VTXWriter
+
+    def __init__(
+        self,
+        filename: str | Path,
+        times: Optional[list[float] | list[int] | None] = None,
+    ):
+        super().__init__(filename, ".bp", times)
 
 
 class VTXSpeciesExport(ExportBaseClass):
-    """Export functions to VTX file
+    """Export species field functions to VTX file
 
     Args:
         filename: The name of the output file
         field: Set of species to export
-        subdomain: A field can be defined on multiple domains.
-            This arguments specifies what subdomains we export on.
-            If `None` we export on all domains.
-        checkpoint: If True, the export will be a checkpoint file
-            using adios4dolfinx and won't be readable by ParaView.
-            Default is False.
+        subdomain: A field can be defined on multiple domains. This arguments specifies
+            what subdomains we export on. If `None` we export on all domains.
+        checkpoint: If True, the export will be a checkpoint file using adios4dolfinx
+            and won't be readable by ParaView. Default is False.
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps. Defaults to None.
+
+    Attributes:
+        filename: The name of the output file
+        field: Set of species to export
+        times: if provided, the field will be exported at these timesteps. Otherwise
+            exports at all timesteps. Defaults to None.
+        writer: The VTXWriter object used to write the file
     """
 
-    field: list[_Species]
-    _subdomain: _VolumeSubdomain
+    _subdomain: VolumeSubdomain
     _checkpoint: bool
+    writer: io.VTXWriter
 
     def __init__(
         self,
         filename: str | Path,
-        field: _Species | list[_Species],
-        subdomain: _VolumeSubdomain = None,
+        field: Species | list[Species],
+        subdomain: VolumeSubdomain = None,
         checkpoint: bool = False,
-    ) -> None:
-        super().__init__(filename, ".bp")
+        times: Optional[list[float] | list[int] | None] = None,
+    ):
+        super().__init__(filename, ".bp", times)
         self.field = field
         self._subdomain = subdomain
         self._checkpoint = checkpoint
 
     @property
-    def field(self) -> list[_Species]:
+    def field(self) -> list[Species]:
         return self._field
 
     @field.setter
-    def field(self, value: _Species | list[_Species]):
+    def field(self, value: Species | list[Species]):
         """
         Update the field to export.
 
-        Note:
-            This also creates a new writer with the updated field.
-
         Args:
             value: The species to export
 
         Raises:
             TypeError: If input field is not a Species or a list of Species
+
+        Note:
+            This also creates a new writer with the updated field.
         """
         # check that all elements of list are festim.Species
         if isinstance(value, list):
             for element in value:
-                if not isinstance(element, (_Species, str)):
+                if not isinstance(element, Species | str):
                     raise TypeError(
-                        "field must be of type festim.Species or a list of festim.Species or str"
+                        "field must be of type festim.Species or a list of "
+                        "festim.Species or str"
                     )
             val = value
-        elif isinstance(value, _Species):
+        elif isinstance(value, Species):
             val = [value]
         else:
             raise TypeError(
-                "field must be of type festim.Species or a list of festim.Species or str",
+                "field must be of type festim.Species or a list of festim.Species or "
+                "str",
                 f"got {type(value)}.",
             )
         self._field = val
 
     def get_functions(self) -> list[fem.Function]:
         """
-        Returns list of species for a given subdomain.
-        If using legacy mode, return the whole species.
+        Returns list of species for a given subdomain. If using legacy mode, return the
+        whole species.
         """
 
         legacy_output: bool = False

src/festim/heat_transfer_problem.py

@@ -28,7 +28,6 @@ def __init__(
         )
 
         self.initial_condition = initial_condition
-        self._vtxfile: VTXWriter | None = None
 
     @property
     def sources(self):
@@ -49,15 +48,13 @@ def boundary_conditions(self, value):
         if not all(
             isinstance(
                 bc,
-                (
-                    boundary_conditions.FixedTemperatureBC,
-                    boundary_conditions.HeatFluxBC,
-                ),
+                boundary_conditions.FixedTemperatureBC | boundary_conditions.HeatFluxBC,
             )
             for bc in value
         ):
             raise TypeError(
-                "boundary_conditions must be a list of festim.FixedTemperatureBC or festim.HeatFluxBC objects"
+                "boundary_conditions must be a list of festim.FixedTemperatureBC "
+                "or festim.HeatFluxBC objects"
             )
         self._boundary_conditions = value
 
@@ -228,7 +225,7 @@ def initialise_exports(self):
                     "XDMF export is not implemented yet for heat transfer problems"
                 )
             if isinstance(export, exports.VTXTemperatureExport):
-                self._vtxfile = VTXWriter(
+                export.writer = VTXWriter(
                     self.u.function_space.mesh.comm,
                     export.filename,
                     [self.u],
@@ -242,7 +239,8 @@ def post_processing(self):
             # TODO if export type derived quantity
             if isinstance(export, exports.SurfaceQuantity):
                 raise NotImplementedError(
-                    "SurfaceQuantity export is not implemented yet for heat transfer problems"
+                    "SurfaceQuantity export is not implemented yet "
+                    "for heat transfer problems"
                 )
                 export.compute(
                     self.mesh.n,
@@ -257,9 +255,11 @@ def post_processing(self):
             if isinstance(export, exports.XDMFExport):
                 export.write(float(self.t))
 
-        if self._vtxfile is not None:
-            self._vtxfile.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))
 
     def __del__(self):
-        if self._vtxfile is not None:
-            self._vtxfile.close()
+        for export in self.exports:
+            if isinstance(export, exports.VTXTemperatureExport):
+                export.writer.close()