Avoid use of .vector().gather() (#202)

Closes #200

Alternative for #201.
Removes usage of .vector()
Taking this opportunity to also eliminate .vector().gather() (which now
would be .dat.global_data) as all-to-all gathers of function data should
generally be avoided. All places where this was used were actually doing
a reduction operation straight after, so now replaced
 with just an MPI allreduce. Introducing utility functions
function_data_min/max/sum for this as it gets a little wieldy.

Copies the other parts of #201:
It also drops use of MeshGeometry.init(), which no longer exists.

This PR also temporarily turns off other broken 3D tests reported in
#197.

Author: stephankramer

animate/metric.py

@@ -24,6 +24,7 @@
     recover_gradient_l2,
     recover_hessian_clement,
 )
+from .utility import function_data_min, function_data_sum
 
 __all__ = ["RiemannianMetric", "determine_metric_complexity", "intersect_on_boundary"]
 
@@ -496,13 +497,13 @@ def interp(f):
         a_max = interp(self._variable_parameters["dm_plex_metric_a_max"])
 
         # Check minimal h_min value is positive and smaller than minimal h_max value
-        _hmin = h_min.vector().gather().min()
+        _hmin = function_data_min(h_min)
         if _hmin <= 0.0:
             raise ValueError(f"Encountered non-positive h_min value: {_hmin}.")
-        if h_max.vector().gather().min() < _hmin:
+        if function_data_min(h_max) < _hmin:
             raise ValueError(
                 "Minimum h_max value is smaller than minimum h_min value:"
-                f"{h_max.vector().gather().min()} < {_hmin}."
+                f"{function_data_min(h_max)} < {_hmin}."
             )
 
         # Check h_max is always at least h_min
@@ -512,7 +513,7 @@ def interp(f):
             raise ValueError("Encountered regions where h_max < h_min.")
 
         # Check minimal a_max value is close to unity or larger
-        _a_max = a_max.vector().gather().min()
+        _a_max = function_data_min(a_max)
         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}.")
 
@@ -934,8 +935,8 @@ def compute_isotropic_dwr_metric(
         )
 
     def _any_inf(self, f):
-        arr = f.vector().gather()
-        return np.isinf(arr).any() or np.isnan(arr).any()
+        arr_sum = function_data_sum(f)
+        return not np.isfinite(arr_sum)
 
     @PETSc.Log.EventDecorator()
     def compute_anisotropic_dwr_metric(
@@ -1002,7 +1003,7 @@ def compute_anisotropic_dwr_metric(
         P0 = firedrake.FunctionSpace(mesh, "DG", 0)
         K_opt = pow(error_indicator, 1 / (convergence_rate + 1))
         K_opt_av = (
-            K_opt / firedrake.assemble(interpolate(K_opt, P0)).vector().gather().sum()
+            K_opt / function_data_sum(firedrake.assemble(interpolate(K_opt, P0)))
         )
         K_ratio = target_complexity * pow(abs(K_opt_av * K_hat / K), 2 / dim)
 

animate/recovery.py

@@ -13,6 +13,7 @@
 from .interpolation import clement_interpolant
 from .math import construct_basis
 from .quality import QualityMeasure, include_dir
+from .utility import function_data_max
 
 __all__ = ["recover_gradient_l2", "recover_hessian_clement", "recover_boundary_hessian"]
 
@@ -152,7 +153,7 @@ def recover_boundary_hessian(f, method="Clement", target_space=None, **kwargs):
     h = firedrake.assemble(
         interpolate(ufl.CellDiameter(mesh), firedrake.FunctionSpace(mesh, "DG", 0))
     )
-    h = firedrake.Constant(1 / h.vector().gather().max() ** 2)
+    h = firedrake.Constant(1 / function_data_max(h) ** 2)
     sp = {
         "ksp_type": "gmres",
         "ksp_gmres_restart": 20,

animate/utility.py

@@ -10,8 +10,10 @@
 import ufl
 from firedrake.__future__ import interpolate
 from firedrake.petsc import PETSc
+from mpi4py import MPI
 
-__all__ = ["Mesh", "VTKFile", "norm", "errornorm"]
+__all__ = ["Mesh", "VTKFile", "norm", "errornorm",
+           "function_data_min", "function_data_max", "function_data_sum"]
 
 
 @PETSc.Log.EventDecorator()
@@ -309,3 +311,21 @@ def function2cofunction(func):
     else:
         cofunc.dat.data_with_halos[:] = func.dat.data_with_halos
     return cofunc
+
+
+def function_data_min(f):
+    """Compute node-wise global minimum of Firedrake function"""
+    mesh = ufl.domain.extract_unique_domain(f)
+    return mesh.comm.allreduce(f.dat.data_ro.min(), MPI.MIN)
+
+
+def function_data_max(f):
+    """Compute node-wise global maximum of Firedrake function"""
+    mesh = ufl.domain.extract_unique_domain(f)
+    return mesh.comm.allreduce(f.dat.data_ro.max(), MPI.MAX)
+
+
+def function_data_sum(f):
+    """Compute global sum of nodal values of Firedrake function"""
+    mesh = ufl.domain.extract_unique_domain(f)
+    return mesh.comm.allreduce(f.dat.data_ro.sum(), MPI.SUM)

demos/ping_pong.py

@@ -64,8 +64,8 @@
 
 niter = 50
 initial_integral = assemble(source * dx)
-initial_min = source.vector().gather().min()
-initial_max = source.vector().gather().max()
+initial_min = function_data_min(source)
+initial_max = function_data_max(source)
 quantities = {
     "integral": {"interpolate": [initial_integral]},
     "minimum": {"interpolate": [initial_min]},
@@ -77,8 +77,8 @@
     interpolate(f_interp, tmp)
     interpolate(tmp, f_interp)
     quantities["integral"]["interpolate"].append(assemble(f_interp * dx))
-    quantities["minimum"]["interpolate"].append(f_interp.vector().gather().min())
-    quantities["maximum"]["interpolate"].append(f_interp.vector().gather().max())
+    quantities["minimum"]["interpolate"].append(function_data_min(f_interp))
+    quantities["maximum"]["interpolate"].append(function_data_max(f_interp))
 f_interp.rename("Interpolate")
 
 fig, axes = plt.subplots(ncols=3, figsize=(18, 5))
@@ -139,8 +139,8 @@
     project(f_proj, tmp)
     project(tmp, f_proj)
     quantities["integral"]["project"].append(assemble(f_proj * dx))
-    quantities["minimum"]["project"].append(f_proj.vector().gather().min())
-    quantities["maximum"]["project"].append(f_proj.vector().gather().max())
+    quantities["minimum"]["project"].append(function_data_min(f_proj))
+    quantities["maximum"]["project"].append(function_data_max(f_proj))
 f_proj.rename("Project")
 
 fig, axes = plt.subplots(ncols=3, figsize=(18, 5))
@@ -189,8 +189,8 @@
     project(f_bounded, tmp, bounded=True)
     project(tmp, f_bounded, bounded=True)
     quantities["integral"]["bounded"].append(assemble(f_bounded * dx))
-    quantities["minimum"]["bounded"].append(f_bounded.vector().gather().min())
-    quantities["maximum"]["bounded"].append(f_bounded.vector().gather().max())
+    quantities["minimum"]["bounded"].append(function_data_min(f_bounded))
+    quantities["maximum"]["bounded"].append(function_data_max(f_bounded))
 f_bounded.rename("Bounded project")
 
 fig, axes = plt.subplots(ncols=3, figsize=(18, 5))

test/test_adapt.py

@@ -76,7 +76,7 @@ def test_no_adapt(dim, serialise):
     Ensure mesh adaptation operations can be turned off.
     """
     mesh = uniform_mesh(dim)
-    dofs = mesh.coordinates.vector().gather().shape
+    dofs = mesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
     mp = {
         "dm_plex_metric": {
             "no_insert": None,
@@ -87,7 +87,8 @@ def test_no_adapt(dim, serialise):
     }
     metric = uniform_metric(mesh, metric_parameters=mp)
     newmesh = try_adapt(mesh, metric, serialise=serialise)
-    assert newmesh.coordinates.vector().gather().shape == dofs
+    newdofs = newmesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
+    assert newdofs == dofs
 
 
 @pytest.mark.parallel(nprocs=2)
@@ -138,7 +139,6 @@ def test_preserve_facet_tags_2d(meshname):
     metric = uniform_metric(mesh)
     newmesh = try_adapt(mesh, metric)
 
-    newmesh.init()
     tags = set(mesh.exterior_facets.unique_markers)
     newtags = set(newmesh.exterior_facets.unique_markers)
     assert tags == newtags, "Facet tags do not match"
@@ -152,15 +152,17 @@ def test_preserve_facet_tags_2d(meshname):
 
 @pytest.mark.parametrize(
     "dim,serialise",
-    [(2, True), (3, True)],
-    ids=["mmg2d", "mmg3d"],
+    # [(2, True), (3, True)],  # FIXME: Broken test (#197)
+    [(2, True)],
+    # ids=["mmg2d", "mmg3d"],  # FIXME: Broken test (#197)
+    ids=["mmg2d"],
 )
 def test_adapt(dim, serialise):
     """
     Test that we can successfully invoke Mmg and that it changes the DoF count.
     """
     mesh = uniform_mesh(dim)
-    dofs = mesh.coordinates.vector().gather().shape
+    dofs = mesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
     mp = {
         "dm_plex_metric": {
             "target_complexity": 100.0,
@@ -169,14 +171,15 @@ def test_adapt(dim, serialise):
     }
     metric = uniform_metric(mesh, metric_parameters=mp)
     newmesh = try_adapt(mesh, metric, serialise=serialise)
-    assert newmesh.coordinates.vector().gather().shape != dofs
+    newdofs = newmesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
+    assert newdofs != dofs
 
 
 @pytest.mark.parallel(nprocs=2)
 @pytest.mark.parametrize(
     "dim,serialise",
-    [(2, True), (3, True)],  # [(2, True), (3, True), (3, False)], # FIXME: hang (#136)
-    ids=["mmg2d", "mmg3d"],  # ["mmg2d", "mmg3d", "ParMmg"],
+    [(2, True)],  # [(2, True), (3, True), (3, False)],  # FIXME: broken (#136,#197)
+    ids=["mmg2d"],  # ["mmg2d", "mmg3d", "ParMmg"],  # FIXME: broken (#136,#197)
 )
 def test_adapt_parallel_np2(dim, serialise):
     """
@@ -190,8 +193,9 @@ def test_adapt_parallel_np2(dim, serialise):
 @pytest.mark.parallel(nprocs=3)
 @pytest.mark.parametrize(
     "dim,serialise",
-    [(2, True), (3, True)],  # [(2, True), (3, True), (3, False)], # FIXME: hang (#136)
-    ids=["mmg2d", "mmg3d"], # ["mmg2d", "mmg3d", "ParMmg"],
+    # [(2, True), (3, True), (3, False)],  # FIXME: broken tests (#136, #197)
+    [(2, True)],
+    ids=["mmg2d"],  # ["mmg2d", "mmg3d", "ParMmg"],  # FIXME: broken tests (#136, #197)
 )
 def test_adapt_parallel_np3(dim, serialise):
     """
@@ -202,7 +206,13 @@ def test_adapt_parallel_np3(dim, serialise):
     test_adapt(dim, serialise=serialise)
 
 
-@pytest.mark.parametrize("dim", [2, 3], ids=["mmg2d", "mmg3d"])
+@pytest.mark.parametrize(
+    "dim",
+    # [2, 3],  # FIXME: Broken test (#197)
+    [2],
+    # ids=["mmg2d", "mmg3d"]  # FIXME: Broken test (#197)
+    ids=["mmg2d"],
+)
 def test_enforce_spd_h_min(dim):
     """
     Tests that :meth:`animate.metric.RiemannianMetric.enforce_spd` applies minimum
@@ -212,14 +222,21 @@ def test_enforce_spd_h_min(dim):
     h = 0.1
     metric = uniform_metric(mesh, a=1 / h**2)
     newmesh = try_adapt(mesh, metric)
-    num_vertices = newmesh.coordinates.vector().gather().shape[0]
+    dofs = newmesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
     metric.set_parameters({"dm_plex_metric_h_min": 0.2})  # h_min > h => h := h_min
     metric.enforce_spd(restrict_sizes=True)
     newmesh = try_adapt(mesh, metric)
-    assert newmesh.coordinates.vector().gather().shape[0] < num_vertices
+    newdofs = newmesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
+    assert newdofs < dofs
 
 
-@pytest.mark.parametrize("dim", [2, 3], ids=["mmg2d", "mmg3d"])
+@pytest.mark.parametrize(
+    "dim",
+    # [2, 3],  # FIXME: Broken test (#197)
+    [2],
+    # ids=["mmg2d", "mmg3d"]  # FIXME: Broken test (#197)
+    ids=["mmg2d"],
+)
 def test_enforce_spd_h_max(dim):
     """
     Tests that :meth:`animate.metric.RiemannianMetric.enforce_spd` applies maximum
@@ -229,11 +246,12 @@ def test_enforce_spd_h_max(dim):
     h = 0.1
     metric = uniform_metric(mesh, a=1 / h**2)
     newmesh = try_adapt(mesh, metric)
-    num_vertices = newmesh.coordinates.vector().gather().shape[0]
+    dofs = newmesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
     metric.set_parameters({"dm_plex_metric_h_max": 0.05})  # h_max < h => h := h_max
     metric.enforce_spd(restrict_sizes=True)
     newmesh = try_adapt(mesh, metric)
-    assert newmesh.coordinates.vector().gather().shape[0] > num_vertices
+    newdofs = newmesh.coordinates.dat.dataset.layout_vec.getSizes()[-1]
+    assert newdofs > dofs
 
 
 # Debugging

test/test_metric.py

@@ -16,6 +16,7 @@
 from test_setup import uniform_mesh, uniform_metric
 
 from animate.metric import P0Metric, RiemannianMetric
+from animate.utility import function_data_min
 
 
 class MetricTestCase(unittest.TestCase):
@@ -707,7 +708,7 @@ def test_eigendecomposition(self, dim, reorder):
             for i in range(dim - 1):
                 f = Function(P1).interpolate(evalues[i])
                 f -= Function(P1).interpolate(evalues[i + 1])
-                if f.vector().gather().min() < 0.0:
+                if function_data_min(f) < 0.0:
                     raise ValueError(
                         f"Eigenvalues are not in descending order: {evalues.dat.data}"
                     )

test/test_quality.py

@@ -10,6 +10,7 @@
 from test_setup import uniform_mesh
 
 from animate.quality import QualityMeasure
+from animate.utility import function_data_sum
 
 
 @pytest.fixture(params=[2, 3])
@@ -60,7 +61,7 @@ def test_uniform_quality_2d(self, measure, expected):
         truth = Function(q.function_space()).assign(expected)
         self.assertAlmostEqual(errornorm(truth, q), 0.0, places=6)
         if measure == "area":
-            s = q.vector().gather().sum()
+            s = function_data_sum(q)
             self.assertAlmostEqual(s, 1.0)
 
     @parameterized.expand(
@@ -79,7 +80,7 @@ def test_uniform_quality_3d(self, measure, expected):
         truth = Function(q.function_space()).assign(expected)
         self.assertAlmostEqual(errornorm(truth, q), 0.0)
         if measure == "volume":
-            s = q.vector().gather().sum()
+            s = function_data_sum(q)
             self.assertAlmostEqual(s, 1.0)
 
     @parameterized.expand(