Merge pull request #1011 from jhdark/alt_coordinate_support

Cylindrical and Spherical coordinate support

Author: jhdark

src/festim/exports/surface_flux.py

@@ -3,25 +3,42 @@
 from scifem import assemble_scalar
 
 from festim.exports.surface_quantity import SurfaceQuantity
+from festim.species import Species
+from festim.subdomain.surface_subdomain import SurfaceSubdomain
 
 
 class SurfaceFlux(SurfaceQuantity):
     """Computes the flux of a field on a given surface
 
     Args:
-        field (festim.Species): species for which the surface flux is computed
-        surface (festim.SurfaceSubdomain1D): surface subdomain
-        filename (str, optional): name of the file to which the surface flux is exported
+        field: species for which the surface flux is computed
+        surface: surface subdomain
+        filename: name of the file to which the surface flux is exported
 
     Attributes:
         see `festim.SurfaceQuantity`
     """
 
+    field: Species
+    surface: SurfaceSubdomain
+    filename: str
+
+    title: str
+    value: float
+    data: list[float]
+
     @property
     def title(self):
         return f"{self.field.name} flux surface {self.surface.id}"
 
-    def compute(self, u, ds: ufl.Measure, entity_maps=None):
+    def __init__(
+        self, field: Species, surface: SurfaceSubdomain, filename: str | None = None
+    ) -> None:
+        super().__init__(field=field, surface=surface, filename=filename)
+
+    def compute(
+        self, u: fem.Function | ufl.indexed.Indexed, ds: ufl.Measure, entity_maps=None
+    ):
         """Computes the value of the flux at the surface
 
         Args:

src/festim/exports/surface_quantity.py

@@ -8,19 +8,28 @@ class SurfaceQuantity:
     """Export SurfaceQuantity
 
     Args:
-        field (festim.Species): species for which the surface flux is computed
-        surface (festim.SurfaceSubdomain1D): surface subdomain
-        filename (str, optional): name of the file to which the surface flux is exported
+        field: species for which the surface flux is computed
+        surface: surface subdomain
+        filename: name of the file to which the surface flux is exported
 
     Attributes:
-        field (festim.Species): species for which the surface flux is computed
-        surface (festim.SurfaceSubdomain): surface subdomain
-        filename (str): name of the file to which the surface flux is exported
-        t (list): list of time values
-        data (list): list of values of the surface quantity
+        field: species for which the surface flux is computed
+        surface: surface subdomain
+        filename: name of the file to which the surface flux is exported
+        t: list of time values
+        data: list of values of the surface quantity
     """
 
-    def __init__(self, field, surface, filename: str | None = None) -> None:
+    field: Species
+    surface: SurfaceSubdomain
+    filename: str | None
+
+    t: list[float]
+    data: list[float]
+
+    def __init__(
+        self, field: Species, surface: SurfaceSubdomain, filename: str | None = None
+    ) -> None:
         self.field = field
         self.surface = surface
         self.filename = filename

src/festim/exports/volume_quantity.py

@@ -1,26 +1,33 @@
 import csv
-import os
 
-from festim import species
+from festim.species import Species
+from festim.subdomain.volume_subdomain import VolumeSubdomain
 
 
 class VolumeQuantity:
     """Export VolumeQuantity
 
     Args:
-        field (festim.Species): species for which the volume quantity is computed
-        volume (festim.VolumeSubdomain): volume subdomain
-        filename (str, optional): name of the file to which the volume quantity is exported
+        field: species for which the volume quantity is computed
+        volume: volume subdomain
+        filename: name of the file to which the volume quantity is exported
 
     Attributes:
-        field (festim.Species): species for which the volume quantity is computed
-        volume (festim.VolumeSubdomain): volume subdomain
-        filename (str): name of the file to which the volume quantity is exported
-        t (list): list of time values
-        data (list): list of values of the volume quantity
+        field: species for which the volume quantity is computed
+        volume: volume subdomain
+        filename: name of the file to which the volume quantity is exported
+        t: list of time values
+        data: list of values of the volume quantity
     """
 
-    def __init__(self, field, volume, filename: str = None) -> None:
+    field: Species
+    volume: VolumeSubdomain
+    filename: str | None
+
+    t: list[float]
+    data: list[float]
+
+    def __init__(self, field, volume, filename: str | None = None) -> None:
         self.field = field
         self.volume = volume
         self.filename = filename
@@ -50,7 +57,7 @@ def field(self):
     @field.setter
     def field(self, value):
         # check that field is festim.Species
-        if not isinstance(value, (species.Species, str)):
+        if not isinstance(value, Species | str):
             raise TypeError("field must be of type festim.Species")
 
         self._field = value

src/festim/hydrogen_transport_problem.py

@@ -421,6 +421,15 @@ def initialise_exports(self):
                     else:
                         adios4dolfinx.write_mesh(export.filename, mesh=self.mesh.mesh)
 
+            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
+                if self.mesh.coordinate_system != "cartesian":
+                    raise NotImplementedError(
+                        f"Derived quantity exports are not implemented for "
+                        f"{self.mesh.coordinate_system} meshes"
+                    )
+
             # if name of species is given then replace with species object
             if isinstance(export.field, list):
                 for idx, field in enumerate(export.field):
@@ -784,9 +793,24 @@ def create_formulation(self):
                     self.mesh.mesh, self.temperature_fenics, spe
                 )
                 if spe.mobile:
-                    self.formulation += ufl.dot(D * ufl.grad(u), ufl.grad(v)) * self.dx(
-                        vol.id
-                    )
+                    if self.mesh.coordinate_system == "cartesian":
+                        self.formulation += ufl.dot(
+                            D * ufl.grad(u), ufl.grad(v)
+                        ) * self.dx(vol.id)
+                    elif self.mesh.coordinate_system == "cylindrical":
+                        r = ufl.SpatialCoordinate(self.mesh.mesh)[0]
+                        self.formulation += (
+                            r
+                            * ufl.dot(D * ufl.grad(u), ufl.grad(v / r))
+                            * self.dx(vol.id)
+                        )
+                    elif self.mesh.coordinate_system == "spherical":
+                        r = ufl.SpatialCoordinate(self.mesh.mesh)[0]
+                        self.formulation += (
+                            r**2
+                            * ufl.dot(D * ufl.grad(u), ufl.grad(v / r**2))
+                            * self.dx(vol.id)
+                        )
 
                 if self.settings.transient:
                     self.formulation += ((u - u_n) / self.dt) * v * self.dx(vol.id)
@@ -1319,7 +1343,24 @@ def create_subdomain_formulation(self, subdomain: _subdomain.VolumeSubdomain):
                 form += ((u - u_n) / self.dt) * v * self.dx(subdomain.id)
 
             if spe.mobile:
-                form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * self.dx(subdomain.id)
+                if self.mesh.coordinate_system == "cartesian":
+                    form += ufl.dot(D * ufl.grad(u), ufl.grad(v)) * self.dx(
+                        subdomain.id
+                    )
+                elif self.mesh.coordinate_system == "cylindrical":
+                    r = ufl.SpatialCoordinate(self.mesh.mesh)[0]
+                    form += (
+                        r
+                        * ufl.dot(D * ufl.grad(u), ufl.grad(v / r))
+                        * self.dx(subdomain.id)
+                    )
+                elif self.mesh.coordinate_system == "spherical":
+                    r = ufl.SpatialCoordinate(self.mesh.mesh)[0]
+                    form += (
+                        r**2
+                        * ufl.dot(D * ufl.grad(u), ufl.grad(v / r**2))
+                        * self.dx(subdomain.id)
+                    )
 
         # add reaction terms
         for reaction in self.reactions:

src/festim/mesh/mesh.py

@@ -10,7 +10,9 @@ class Mesh:
     Mesh class
 
     Args:
-        mesh The mesh. Defaults to None.
+        mesh: The mesh. Defaults to None.
+        coordinate_system: the coordinate system of the mesh ("cartesian",
+            "cylindrical", "spherical"). Defaults to "cartesian".
 
     Attributes:
         mesh The mesh
@@ -20,9 +22,18 @@ class Mesh:
             of the mesh.
     """
 
-    _mesh: dolfinx.mesh.Mesh
+    mesh: dolfinx.mesh.Mesh
+    coordinate_system: str | None
 
-    def __init__(self, mesh: dolfinx_Mesh | None = None):
+    vdim: int
+    fdim: int
+    n: ufl.FacetNormal
+
+    def __init__(
+        self,
+        mesh: dolfinx_Mesh | None = None,
+        coordinate_system: str | None = "cartesian",
+    ):
         self.mesh = mesh
         if self._mesh is not None:
             # create cell to facet connectivity
@@ -34,6 +45,9 @@ def __init__(self, mesh: dolfinx_Mesh | None = None):
             self._mesh.topology.create_connectivity(
                 self._mesh.topology.dim - 1, self._mesh.topology.dim
             )
+        self.coordinate_system = coordinate_system
+
+        self.check_mesh_dim_coords()
 
     @property
     def mesh(self):
@@ -46,6 +60,20 @@ def mesh(self, value):
         else:
             raise TypeError("Mesh must be of type dolfinx.mesh.Mesh")
 
+    @property
+    def coordinate_system(self):
+        return self._coordinate_system
+
+    @coordinate_system.setter
+    def coordinate_system(self, value):
+        if value in ["cartesian", "cylindrical", "spherical"]:
+            self._coordinate_system = value
+        else:
+            raise ValueError(
+                "coordinate_system must be one of 'cartesian', "
+                "'cylindrical', or 'spherical'"
+            )
+
     @property
     def vdim(self):
         if self._mesh is None:
@@ -144,3 +172,16 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
         )
 
         return facet_meshtags, volume_meshtags
+
+    def check_mesh_dim_coords(self):
+        """Checks if the used coordinates can be applied for geometry with the specified
+        dimensions"""
+
+        if self.coordinate_system == "spherical" and self.vdim != 1:
+            raise AttributeError(
+                "spherical coordinates can be used for one-dimensional domains only"
+            )
+        if self.coordinate_system == "cylindrical" and self.vdim == 3:
+            raise AttributeError(
+                "cylindrical coordinates cannot be used for 3D domains"
+            )

test/benchmark.py

@@ -1,3 +1,4 @@
+import tempfile
 import time
 
 from mpi4py import MPI
@@ -32,7 +33,6 @@
     exp,
     grad,
 )
-import tempfile
 
 
 def fenics_test_permeation_problem(mesh_size=1001):