Merge main into release

.github/workflows/test_suite_main.yml

@@ -3,7 +3,8 @@ name: 'Run Animate test suite (main)'
 on:
   # Run test suite whenever the main branch is updated
   push:
-    branches: [main]
+    branches:
+      - main
     paths:
       - '**.py'
       - '**.cxx'
@@ -13,7 +14,6 @@ on:
 
   # Run test suite whenever commits are pushed to an open PR
   pull_request:
-    branches: [main]
     paths:
       - '**.py'
       - '**.cxx'

.github/workflows/test_suite_release.yml

@@ -3,13 +3,13 @@ name: 'Run Animate test suite (release)'
 on:
   # Run test suite whenever the release branch is updated
   push:
-    branches: [release]
+    branches:
+      - release
     paths:
       - '.github/workflows/test_suite_release.yml'
 
   # Run test suite whenever commits that change this workflow are pushed to an open PR
   pull_request:
-    branches: [release]
     paths:
       - '.github/workflows/test_suite_release.yml'
 

animate/adapt.py

@@ -1,4 +1,5 @@
 import abc
+import gc
 import os
 from functools import cached_property
 from shutil import rmtree
@@ -97,10 +98,9 @@ def adapted_mesh(self):
         :rtype: :class:`firedrake.mesh.MeshGeometry`
         """
         self.metric.enforce_spd(restrict_sizes=True, restrict_anisotropy=True)
-        size = self.metric.dat.dataset.layout_vec.getSizes()
-        data = self.metric.dat._data[: size[0]]
+        data = self.metric.dat.data_ro_with_halos
         v = PETSc.Vec().createWithArray(
-            data, size=size, bsize=self.metric.dat.cdim, comm=self.mesh.comm
+            data, size=data.size, bsize=self.metric.dat.cdim, comm=COMM_SELF
         )
         reordered = to_petsc_local_numbering(v, self.metric.function_space())
         v.destroy()
@@ -176,7 +176,7 @@ def adapt(mesh, *metrics, name=None, serialise=None, remove_checkpoints=True):
     """
     nprocs = COMM_WORLD.size
 
-    dim = mesh.topological_dimension()
+    dim = mesh.topological_dimension
     if serialise is None:
         serialise = nprocs > 1 and dim != 3
     elif not serialise and dim != 3:
@@ -193,13 +193,20 @@ def adapt(mesh, *metrics, name=None, serialise=None, remove_checkpoints=True):
         chk_fpath = os.path.join(chk_dir, "adapted_mesh_checkpoint.h5")
         metric_name = "tmp_metric"
         save_checkpoint(chk_fpath, metric, metric_name)
+        # Ensure all processes are finished writing
+        COMM_WORLD.barrier()
 
         if COMM_WORLD.rank == 0:
             metric0 = load_checkpoint(chk_fpath, mesh.name, metric_name, comm=COMM_SELF)
             adaptor0 = MetricBasedAdaptor(metric0._mesh, metric0, name=name)
             with fchk.CheckpointFile(chk_fpath, "w", comm=COMM_SELF) as chk:
                 chk.save_mesh(adaptor0.adapted_mesh)
+        # Ensure rank 0 is finished writing
         COMM_WORLD.barrier()
+        # Garbage collection might be called at different times on different ranks due
+        # to diverging paths, which appears to stall in final cleanup on Python
+        # system exit. Ensure everything is in-sync again at this point
+        gc.collect()
 
         # In parallel, load from the checkpoint
         if not os.path.exists(chk_fpath):

animate/cython/numbering.pyx

@@ -26,25 +26,27 @@ cdef extern from "petscis.h" nogil:
 def to_petsc_local_numbering(PETSc.Vec vec, V):
     """
     Reorder a PETSc Vec corresponding to a Firedrake Function w.r.t.
-    the PETSc natural numbering.
+    the PETSc natural numbering, i.e. the numbering consistent with
+    that of the DMPlex topological points.
 
-    :arg vec: the PETSc Vec to reorder; must be a global vector
+    :arg vec: the PETSc Vec to reorder; must be a local vector, i.e.
+              a sequential vector that includes all (owned and halo) DoFs
     :arg V: the FunctionSpace of the Function which the Vec comes from
     :ret out: a copy of the Vec, ordered with the PETSc natural numbering
     """
-    cdef int dim, idx, start, end, p, d, k
+    cdef int dim, idx, lsize, p, d, k
     cdef PetscInt dof, off
     cdef PETSc.Vec out
     cdef PETSc.Section section
     cdef np.ndarray varray, oarray
 
-    section = V.dm.getGlobalSection()
+    section = V.dm.getLocalSection()
     out = vec.duplicate()
     varray = vec.array_r
     oarray = out.array
     dim = V.value_size
     idx = 0
-    start, end = vec.getOwnershipRange()
+    lsize = vec.getSize()
     for p in range(*section.getChart()):
         CHKERR(PetscSectionGetDof(section.sec, p, &dof))
         if dof > 0:
@@ -53,7 +55,12 @@ def to_petsc_local_numbering(PETSc.Vec vec, V):
             off *= dim
             for d in range(dof):
                 for k in range(dim):
-                    oarray[idx] = varray[off + dim * d + k - start]
+                    oarray[idx] = varray[off + dim * d + k]
                     idx += 1
-    assert idx == (end - start)
+    if idx != lsize:
+        raise ValueError(
+           f"Number of local section entries not the same as vector size"
+           f"({idx} vs. {lsize}). Need to provide local vector including halo DoFs."
+        )
+
     return out

animate/interpolation.py

@@ -324,7 +324,7 @@ def clement_interpolant(source, target_space=None, boundary=False):
     if rank not in (0, 1, 2):
         raise ValueError(f"Rank-{rank + 1} tensors are not supported.")
     mesh = Vs.mesh()
-    dim = mesh.topological_dimension()
+    dim = mesh.topological_dimension
 
     # Process target space
     Vt = target_space

animate/math.py

@@ -73,7 +73,7 @@ def construct_basis(vector, normalise=True):
         if not isinstance(vector, ufl.core.expr.Expr):
             raise TypeError(f"Expected UFL Expr, not '{type(vector)}'.")
         as_vector = ufl.as_vector
-        dim = ufl.domain.extract_unique_domain(vector).topological_dimension()
+        dim = ufl.domain.extract_unique_domain(vector).topological_dimension
 
     if dim not in (2, 3):
         raise ValueError(f"Dimension {dim} not supported.")

animate/metric.py

@@ -89,7 +89,7 @@ def __init__(self, function_space, *args, **kwargs):
         # Check that we have an appropriate tensor P1 function
         fs = self.function_space()
         mesh = fs.mesh()
-        tdim = mesh.topological_dimension()
+        tdim = mesh.topological_dimension
         if tdim not in (2, 3):
             raise ValueError(f"Riemannian metric should be 2D or 3D, not {tdim}D.")
         if isinstance(fs.dof_count, Iterable):
@@ -303,7 +303,7 @@ def _set_plex_coordinates(self):
         consistent.
         """
         entity_dofs = np.zeros(self._tdim + 1, dtype=np.int32)
-        entity_dofs[0] = self._mesh.geometric_dimension()
+        entity_dofs[0] = self._mesh.geometric_dimension
         coord_section = self._mesh.create_section(entity_dofs)[0]
         # NOTE: section doesn't have any fields, but PETSc assumes it to have one
         coord_dm = self._plex.getCoordinateDM()
@@ -349,6 +349,14 @@ def compute_hessian(self, field, method="mixed_L2", **kwargs):
         for the gradient recovery. The target space for the Hessian recovery is
         inherited from the metric itself.
         """
+        if len(field.ufl_shape) > 0:
+            value_error = (
+                "RiemannianMetric.compute_hessian only accepts scalar fields. To "
+                "recover Hessians of higher rank fields, call the method on separate "
+                "RiemannianMetrics for each component and combine them with "
+                "RiemannianMetric.combine."
+            )
+            raise ValueError(value_error)
         if method == "L2":
             gradient = recover_gradient_l2(
                 field, target_space=kwargs.get("target_space")
@@ -518,7 +526,7 @@ def interp(f):
         if not np.isclose(_a_max, 1.0) and _a_max < 1.0:
             raise ValueError(f"Encountered a_max value smaller than unity: {_a_max}.")
 
-        dim = mesh.topological_dimension()
+        dim = mesh.topological_dimension
         boundary_tag = self._variable_parameters.get("dm_plex_metric_boundary_tag")
         if boundary_tag is None:
             node_set = self.function_space().node_set
@@ -621,25 +629,28 @@ def intersect(self, *metrics):
                     "Cannot intersect metrics with different function spaces."
                 )
 
-        # Intersect the metrics recursively one at a time
         if len(metrics) == 0:
-            pass
-        elif len(metrics) == 1:
-            v1 = self._create_from_array(self.dat.data_with_halos)
-            v2 = self._create_from_array(metrics[0].dat.data_ro_with_halos)
-            vout = self._create_from_array(np.zeros_like(self.dat.data_with_halos))
-
-            # Compute the intersection on the PETSc level
-            self._plex.metricIntersection2(v1, v2, vout)
-
-            # Assign to the output of the intersection
-            size = np.shape(self.dat.data_with_halos)
-            self.dat.data_with_halos[:] = np.reshape(vout.array, size)
-            v2.destroy()
-            v1.destroy()
-            vout.destroy()
-        else:
+            return self
+
+        # Intersect the metrics recursively, starting with metrics[0]
+        v1 = self._create_from_array(self.dat.data_with_halos)
+        v2 = self._create_from_array(metrics[0].dat.data_ro_with_halos)
+        vout = self._create_from_array(np.zeros_like(self.dat.data_with_halos))
+
+        # Compute the intersection on the PETSc level
+        self._plex.metricIntersection2(v1, v2, vout)
+
+        # Assign to the output of the intersection
+        size = np.shape(self.dat.data_with_halos)
+        self.dat.data_with_halos[:] = np.reshape(vout.array, size)
+        v2.destroy()
+        v1.destroy()
+        vout.destroy()
+
+        # Intersect with the remaining metrics
+        if len(metrics) > 1:
             self.intersect(*metrics[1:])
+
         return self
 
     @PETSc.Log.EventDecorator()
@@ -734,7 +745,7 @@ def compute_eigendecomposition(self, reorder=False):
         mesh = V_ten.mesh()
         fe = (V_ten.ufl_element().family(), V_ten.ufl_element().degree())
         V_vec = firedrake.VectorFunctionSpace(mesh, *fe)
-        dim = mesh.topological_dimension()
+        dim = mesh.topological_dimension
         evectors, evalues = firedrake.Function(V_ten), firedrake.Function(V_vec)
         if reorder:
             name = "get_reordered_eigendecomposition"
@@ -788,7 +799,7 @@ def assemble_eigendecomposition(self, evectors, evalues):
                 "Mismatching finite element space degrees:"
                 f" {fe_ten.degree()} vs. {fe_vec.degree()}."
             )
-        dim = V_ten.mesh().topological_dimension()
+        dim = V_ten.mesh().topological_dimension
         op2.par_loop(
             get_metric_kernel("set_eigendecomposition", dim),
             V_ten.node_set,
@@ -848,7 +859,7 @@ def density_and_quotients(self, reorder=False):
         fs_ten = self.function_space()
         mesh = fs_ten.mesh()
         fe = (fs_ten.ufl_element().family(), fs_ten.ufl_element().degree())
-        dim = mesh.topological_dimension()
+        dim = mesh.topological_dimension
         evectors, evalues = self.compute_eigendecomposition(reorder=reorder)
 
         # Extract density and quotients
@@ -886,7 +897,7 @@ def compute_isotropic_metric(
         mesh = ufl.domain.extract_unique_domain(error_indicator)
         if mesh != self.function_space().mesh():
             raise ValueError("Cannot use an error indicator from a different mesh.")
-        dim = mesh.topological_dimension()
+        dim = mesh.topological_dimension
 
         # Interpolate P0 indicators into P1 space
         if interpolant == "Clement":
@@ -980,7 +991,7 @@ def compute_anisotropic_dwr_metric(
         mesh = ufl.domain.extract_unique_domain(error_indicator)
         if mesh != self.function_space().mesh():
             raise ValueError("Cannot use an error indicator from a different mesh.")
-        dim = mesh.topological_dimension()
+        dim = mesh.topological_dimension
         if convergence_rate < 1.0:
             raise ValueError(
                 f"Convergence rate must be at least one, not {convergence_rate}."
@@ -1133,7 +1144,7 @@ def determine_metric_complexity(H_interior, H_boundary, target, p, **kwargs):
     :returns: unique solution of algebraic problem
     :rtype: :class:`float`
     """
-    d = H_interior.function_space().mesh().topological_dimension()
+    d = H_interior.function_space().mesh().topological_dimension
     if d not in (2, 3):
         raise ValueError(f"Spatial dimension {d} not supported.")
     if np.isinf(p):
@@ -1183,7 +1194,7 @@ def intersect_on_boundary(*metrics, boundary_tag="on_boundary"):
     n = len(metrics)
     assert n > 0, "Nothing to combine"
     fs = metrics[0].function_space()
-    dim = fs.mesh().topological_dimension()
+    dim = fs.mesh().topological_dimension
     if dim not in (2, 3):
         raise ValueError(
             f"Spatial dimension {dim} not supported. Must be either 2 or 3."

animate/quality.py

@@ -8,10 +8,10 @@
 import ufl
 from firedrake.__future__ import interpolate
 from firedrake.petsc import PETSc
+from petsctools import get_petsc_dirs
 from pyop2 import op2
-from pyop2.utils import get_petsc_dir
 
-petsc_dirs = get_petsc_dir()
+petsc_dirs = get_petsc_dirs()
 include_dir = ["%s/include/eigen3" % petsc_dirs[-1]]
 
 __all__ = ["QualityMeasure"]
@@ -60,7 +60,7 @@ def __init__(self, mesh, metric=None, python=False):
         self.mesh = mesh
         self.metric = metric
         self.python = python
-        self.dim = mesh.topological_dimension()
+        self.dim = mesh.topological_dimension
         self.coords = mesh.coordinates
         self.P0 = firedrake.FunctionSpace(mesh, "DG", 0)
         src_dir = os.path.join(os.path.dirname(__file__), "cxx")

animate/recovery.py

@@ -132,7 +132,7 @@ def recover_boundary_hessian(f, method="Clement", target_space=None, **kwargs):
     """
 
     mesh = ufl.domain.extract_unique_domain(f)
-    d = mesh.topological_dimension()
+    d = mesh.topological_dimension
     assert d in (2, 3)
 
     # Apply Gram-Schmidt to get tangent vectors

test/test_checkpoint.py

@@ -7,7 +7,6 @@
 import ufl
 from firedrake.function import Function
 from firedrake.functionspace import FunctionSpace, TensorFunctionSpace
-from firedrake.mesh import _generate_default_mesh_topology_name
 from test_setup import uniform_mesh
 
 from animate.checkpointing import get_checkpoint_dir, load_checkpoint, save_checkpoint
@@ -57,8 +56,8 @@ def test_save(self, filename="test_save.h5", metric=None):
         metric = metric or self.metric
         save_checkpoint(fpath, metric)
         self.assertTrue(os.path.exists(fpath))
-        mesh_name = metric._mesh.name
-        topology_name = _generate_default_mesh_topology_name(mesh_name)
+        mesh = ufl.domain.extract_unique_domain(metric)
+        topology_name = mesh.topology_dm.name
         with h5py.File(fpath, "r") as h5:
             self.assertTrue("topologies" in h5)
             self.assertTrue(topology_name in h5["topologies"].keys())

test/test_interpolation.py

@@ -66,7 +66,7 @@ def get_space(self, rank, family, degree):
         elif rank == 1:
             return VectorFunctionSpace(self.mesh, family, degree)
         else:
-            shape = tuple(rank * [self.mesh.topological_dimension()])
+            shape = tuple(rank * [self.mesh.topological_dimension])
             return TensorFunctionSpace(self.mesh, family, degree, shape=shape)
 
     def analytic(self, rank):