Custom derived quantity

src/festim/__init__.py

@@ -26,6 +26,8 @@
 )
 from .exports.average_surface import AverageSurface
 from .exports.average_volume import AverageVolume
+from .exports.custom_quantity import CustomQuantity
+from .exports.derived_quantity import DerivedQuantity
 from .exports.maximum_surface import MaximumSurface
 from .exports.maximum_volume import MaximumVolume
 from .exports.minimum_surface import MinimumSurface

src/festim/exports/__init__.py

@@ -1,5 +1,7 @@
 from .average_surface import AverageSurface
 from .average_volume import AverageVolume
+from .custom_quantity import CustomQuantity
+from .derived_quantity import DerivedQuantity
 from .maximum_surface import MaximumSurface
 from .maximum_volume import MaximumVolume
 from .minimum_surface import MinimumSurface
@@ -23,6 +25,8 @@
     "AverageSurface",
     "AverageVolume",
     "CustomFieldExport",
+    "CustomQuantity",
+    "DerivedQuantity",
     "ExportBaseClass",
     "MaximumSurface",
     "MaximumVolume",

src/festim/exports/custom_quantity.py

@@ -0,0 +1,131 @@
+from collections.abc import Callable
+
+import ufl
+from dolfinx import fem
+from scifem import assemble_scalar
+
+from festim.exports.derived_quantity import DerivedQuantity
+from festim.subdomain.surface_subdomain import SurfaceSubdomain
+from festim.subdomain.volume_subdomain import VolumeSubdomain
+
+
+class CustomQuantity(DerivedQuantity):
+    r"""
+    Export CustomQuantity.
+
+    Args:
+        expr: function that returns a UFL expression
+        subdomain: subdomain on which the quantity is evaluated
+        title: title of the exported quantity
+        filename: name of the file to which the quantity is exported
+
+    Attributes:
+        expr: function that returns a UFL expression
+        subdomain: subdomain on which the quantity is evaluated
+        title: title of the exported quantity
+        filename: name of the file to which the quantity is exported
+        t: list of time values
+        data: list of values of the quantity
+
+    Usage:
+
+    .. testcode::
+
+        import numpy as np
+        import festim as F
+
+        material = F.Material(D_0=1, E_D=0)
+        volume = F.VolumeSubdomain(id=1, material=material)
+        surface = F.SurfaceSubdomain(id=1, locator=lambda x: np.isclose(x[1], 1))
+
+        def total_concentration(**kwargs):
+            return kwargs["A"] + kwargs["B"]
+
+        quantity = F.CustomQuantity(
+            expr=total_concentration,
+            subdomain=volume,
+            title="Total quantity",
+        )
+
+        surface_quantity = F.CustomQuantity(
+            expr=lambda **kwargs: -kwargs["D_A"] * ufl.dot(
+                ufl.grad(kwargs["A"]), kwargs["n"]
+            ),
+            subdomain=surface,
+            title="Surface flux",
+        )
+
+    The callable passed to ``expr`` receives keyword arguments assembled by the
+    problem class. Common entries are:
+
+    ``A``, ``B``, ...
+        Concentrations of the species present in the problem (here A and B).
+    ``n``
+        The facet normal on the selected surface subdomain.
+    ``T``
+        The temperature field.
+    ``D_A``, ``D_B``, ...
+        Species-specific diffusion coefficients.
+    ``D``
+        The diffusion coefficient data, either a single field for one species or
+        a dictionary keyed by species name when several species are present.
+    ``x``
+        The spatial coordinate (x[0], x[1], x[2]).
+
+    For a surface quantity, the returned UFL expression can represent a flux such as
+
+    .. math::
+
+        q = -D\,\nabla c \cdot n
+
+    and FESTIM will assemble
+
+    .. math::
+
+        Q = \int_{\Gamma} q\,\mathrm{d}\Gamma
+
+    over the selected surface subdomain.
+
+    The expression returned by ``expr`` is treated as an integrand and assembled over
+    the chosen subdomain.
+
+    .. math::
+
+        Q = \int_{\Omega} q\,\mathrm{d}\Omega
+
+    where ``q`` is the UFL expression returned by ``expr`` and ``\Omega`` is either a
+    surface or a volume subdomain.
+    """
+
+    def __init__(
+        self,
+        expr: Callable,
+        subdomain: SurfaceSubdomain | VolumeSubdomain,
+        title: str = "Custom Quantity",
+        filename: str | None = None,
+    ) -> None:
+        super().__init__(filename=filename)
+        self.expr = expr
+        self.subdomain = subdomain
+        self._title = title
+        self.ufl_expr = None
+
+    @property
+    def title(self):
+        return self._title
+
+    def compute(self, measure: ufl.Measure, entity_maps: dict | None = None):
+        """Computes the value of the custom quantity and appends it to the data list.
+
+        Args:
+            measure: volume or surface measure of the model
+            entity_maps: entity maps relating parent mesh and submesh
+        """
+        if self.ufl_expr is None:
+            raise ValueError("The UFL expression has not been evaluated yet.")
+
+        form = fem.form(
+            self.ufl_expr * measure(self.subdomain.id), entity_maps=entity_maps
+        )
+        self.value = assemble_scalar(form)
+        self.data.append(self.value)

src/festim/exports/derived_quantity.py

@@ -0,0 +1,60 @@
+import csv
+from abc import ABC, abstractmethod
+
+
+class DerivedQuantity(ABC):
+    """Base class for all derived quantities.
+
+    Attributes:
+        filename: name of the file to which the quantity is exported
+        t: list of time values
+        data: list of values of the quantity
+    """
+
+    filename: str | None
+    t: list[float]
+    data: list[float]
+
+    def __init__(self, filename: str | None = None) -> None:
+        self.filename = filename
+        self.t = []
+        self.data = []
+        self._first_time_export = True
+
+    @property
+    @abstractmethod
+    def title(self):
+        pass
+
+    @abstractmethod
+    def compute(self, *args, **kwargs):
+        pass
+
+    @property
+    def filename(self):
+        return self._filename
+
+    @filename.setter
+    def filename(self, value):
+        if value is None:
+            self._filename = None
+        elif not isinstance(value, str):
+            raise TypeError("filename must be of type str")
+        elif not value.endswith(".csv") and not value.endswith(".txt"):
+            raise ValueError("filename must end with .csv or .txt")
+        self._filename = value
+
+    def write(self, t):
+        """If the filename doesnt exist yet, create it and write the header, then append
+        the time and value to the file."""
+
+        if self.filename is not None:
+            if self._first_time_export:
+                header = ["t(s)", f"{self.title}"]
+                with open(self.filename, mode="w+", newline="") as file:
+                    writer = csv.writer(file)
+                    writer.writerow(header)
+                self._first_time_export = False
+            with open(self.filename, mode="a", newline="") as file:
+                writer = csv.writer(file)
+                writer.writerow([t, self.value])

src/festim/exports/maximum_surface.py

@@ -19,6 +19,8 @@ class MaximumSurface(SurfaceQuantity):
         see `festim.SurfaceQuantity`
     """
 
+    facet_meshtags: dolfinx.mesh.MeshTags
+
     @property
     def title(self):
         return f"Maximum {self.field.name} surface {self.surface.id}"

src/festim/exports/maximum_volume.py

@@ -19,6 +19,8 @@ class MaximumVolume(VolumeQuantity):
         see `festim.VolumeQuantity`
     """
 
+    volume_meshtags: dolfinx.mesh.MeshTags
+
     @property
     def title(self):
         return f"Maximum {self.field.name} volume {self.volume.id}"

src/festim/exports/minimum_surface.py

@@ -19,6 +19,8 @@ class MinimumSurface(SurfaceQuantity):
         see `festim.SurfaceQuantity`
     """
 
+    facet_meshtags: dolfinx.mesh.MeshTags
+
     @property
     def title(self):
         return f"Minimum {self.field.name} surface {self.surface.id}"

src/festim/exports/minimum_volume.py

@@ -19,6 +19,8 @@ class MinimumVolume(VolumeQuantity):
         see `festim.VolumeQuantity`
     """
 
+    volume_meshtags: dolfinx.mesh.MeshTags
+
     @property
     def title(self):
         return f"Minimum {self.field.name} volume {self.volume.id}"

src/festim/exports/surface_quantity.py

@@ -1,10 +1,11 @@
-import csv
+from abc import abstractmethod
 
+from festim.exports.derived_quantity import DerivedQuantity
 from festim.species import Species
 from festim.subdomain.surface_subdomain import SurfaceSubdomain
 
 
-class SurfaceQuantity:
+class SurfaceQuantity(DerivedQuantity):
     """Export SurfaceQuantity.
 
     Args:
@@ -28,29 +29,23 @@ class SurfaceQuantity:
     data: list[float]
 
     def __init__(
-        self, field: Species, surface: SurfaceSubdomain, filename: str | None = None
+        self,
+        field: Species | str,
+        surface: SurfaceSubdomain | int,
+        filename: str | None = None,
     ) -> None:
+        super().__init__(filename=filename)
         self.field = field
         self.surface = surface
-        self.filename = filename
-
-        self.t = []
-        self.data = []
-        self._first_time_export = True
 
     @property
-    def filename(self):
-        return self._filename
-
-    @filename.setter
-    def filename(self, value):
-        if value is None:
-            self._filename = None
-        elif not isinstance(value, str):
-            raise TypeError("filename must be of type str")
-        elif not value.endswith(".csv") and not value.endswith(".txt"):
-            raise ValueError("filename must end with .csv or .txt")
-        self._filename = value
+    @abstractmethod
+    def title(self):
+        pass
+
+    @abstractmethod
+    def compute(self, *args, **kwargs):
+        pass
 
     @property
     def surface(self):
@@ -60,6 +55,7 @@ def surface(self):
     def surface(self, value):
         if not isinstance(value, int | SurfaceSubdomain) or isinstance(value, bool):
             raise TypeError("surface should be an int or F.SurfaceSubdomain")
+
         self._surface = value
 
     @property
@@ -68,23 +64,7 @@ def field(self):
 
     @field.setter
     def field(self, value):
-        # check that field is festim.Species
         if not isinstance(value, Species | str):
-            raise TypeError("field must be of type festim.Species")
+            raise TypeError("field must be of type F.Species or str")
 
         self._field = value
-
-    def write(self, t):
-        """If the filename doesnt exist yet, create it and write the header, then append
-        the time and value to the file."""
-
-        if self.filename is not None:
-            if self._first_time_export:
-                header = ["t(s)", f"{self.title}"]
-                with open(self.filename, mode="w+", newline="") as file:
-                    writer = csv.writer(file)
-                    writer.writerow(header)
-                self._first_time_export = False
-            with open(self.filename, mode="a", newline="") as file:
-                writer = csv.writer(file)
-                writer.writerow([t, self.value])