Merge branch 'main' into leo/refactor_interpolate

Author: leo-collins

.github/workflows/core.yml

@@ -335,7 +335,7 @@ jobs:
           matrix.arch == 'default'
         run: |
           . venv/bin/activate
-          git clone --depth 1 https://github.com/thetisproject/thetis.git thetis-repo
+          git clone --depth 1 https://github.com/thetisproject/thetis.git thetis-repo --branch ${{ inputs.target_branch }}
           pip install --verbose ./thetis-repo
           python -m pytest -n 8 --verbose thetis-repo/test_adjoint/test_swe_adjoint.py
         timeout-minutes: 10
@@ -359,7 +359,7 @@ jobs:
           matrix.arch == 'default'
         run: |
           . venv/bin/activate
-          git clone --depth 1 https://github.com/g-adopt/g-adopt.git g-adopt-repo
+          git clone --depth 1 https://github.com/g-adopt/g-adopt.git g-adopt-repo --branch ${{ inputs.target_branch }}
           pip install --verbose ./g-adopt-repo
           make -C g-adopt-repo/demos/mantle_convection/base_case check
         timeout-minutes: 5

demos/patch/hcurl_riesz_star.py.rst

@@ -5,11 +5,17 @@ Contributed by `Robert Kirby <https://sites.baylor.edu/robert_kirby/>`_
 and `Pablo Brubeck <https://www.maths.ox.ac.uk/people/pablo.brubeckmartinez/>`_.
 
 Multigrid in H(div) and H(curl) also requires relaxation based on topological patches.
-Here, we demonstrate how to do this in the latter case. ::
+Here, we demonstrate how to do this in the latter case.
+
+We start by importing firedrake and setting up a :func:`.MeshHierarchy` and the
+exact solution and forcing data. Crucially, the meshes must have an overlapping
+parallel domain decomposition that supports the vertex star patches. This is set
+via the `distribution_parameters` kwarg of the :func:`.Mesh` constructor. ::
 
   from firedrake import *
 
-  base = UnitCubeMesh(2, 2, 2)
+  dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)}
+  base = UnitCubeMesh(2, 2, 2, distribution_parameters=dparams)
   mh = MeshHierarchy(base, 3)
   mesh = mh[-1]
 

demos/patch/hdiv_riesz_star.py.rst

@@ -5,11 +5,17 @@ Contributed by `Robert Kirby <https://sites.baylor.edu/robert_kirby/>`_
 and `Pablo Brubeck <https://www.maths.ox.ac.uk/people/pablo.brubeckmartinez/>`_.
 
 Multigrid in H(div) and H(curl) also requires relaxation based on topological patches.
-Here, we demonstrate how to do this in the former case. ::
+Here, we demonstrate how to do this in the former case.
+
+We start by importing firedrake and setting up a :func:`.MeshHierarchy` and the
+exact solution and forcing data. Crucially, the meshes must have an overlapping
+parallel domain decomposition that supports the vertex star patches. This is set
+via the `distribution_parameters` kwarg of the :func:`.Mesh` constructor. ::
 
   from firedrake import *
 
-  base = UnitCubeMesh(2, 2, 2)
+  dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)}
+  base = UnitCubeMesh(2, 2, 2, distribution_parameters=dparams)
   mh = MeshHierarchy(base, 3)
   mesh = mh[-1]
 

demos/patch/poisson_mg_patches.py.rst

@@ -4,7 +4,7 @@ Using patch relaxation for multigrid
 Contributed by `Robert Kirby <https://sites.baylor.edu/robert_kirby/>`_
 and `Pablo Brubeck <https://www.maths.ox.ac.uk/people/pablo.brubeckmartinez/>`_.
 
-Simple relaxation like point Jacobi are not optimal or even suitable
+Simple multigrid relaxation methods like point Jacobi are not optimal or even suitable
 smoothers for all applications.  Firedrake supports additive Schwarz methods
 based on local patch-based decompositions through two different paths.
 
@@ -20,12 +20,15 @@ degree-independent iteration counts.  Here, all the degrees of freedom in the ce
 For many problems, point Jacobi is even worse, and patches are required even to
 get a convergent method.  We refer the reader to other demos.
 
-We start by importing firedrake and setting up a mesh hierarchy and the
-exact solution and forcing data. ::
+We start by importing firedrake and setting up a :func:`.MeshHierarchy` and the
+exact solution and forcing data. Crucially, the meshes must have an overlapping
+parallel domain decomposition that supports the vertex star patches. This is set
+via the `distribution_parameters` kwarg of the :func:`.Mesh` constructor. ::
 
   from firedrake import *
 
-  base = UnitSquareMesh(4, 4)
+  dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)}
+  base = UnitSquareMesh(4, 4, distribution_parameters=dparams)
   mh = MeshHierarchy(base, 1)
   mesh = mh[-1]
 

demos/patch/stokes_vanka_patches.py.rst

@@ -16,11 +16,17 @@ of the patch but not pressures.
 In practice, we arrive at mesh-independent multigrid convergence using these relaxation.
 We can construct Vanka patches either through :class:`~.PatchPC`, in which the bilinear form
 is assembled on each vertex patch, or through :class:`~.ASMVankaPC`, in which the patch
-operators are extracted from the globally assembled stiffness matrix. ::
+operators are extracted from the globally assembled stiffness matrix.
+
+We start by importing firedrake and setting up a :func:`.MeshHierarchy` and the
+exact solution and forcing data. Crucially, the meshes must have an overlapping
+parallel domain decomposition that supports the Vanka patches. This is set
+via the `distribution_parameters` kwarg of the :func:`.Mesh` constructor. ::
 
   from firedrake import *
 
-  base = UnitSquareMesh(4, 4)
+  dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 2)}
+  base = UnitSquareMesh(4, 4, distribution_parameters=dparams)
   mh = MeshHierarchy(base, 3)
   mesh = mh[-1]
 

docs/source/install.rst

@@ -170,6 +170,17 @@ For the default build, running ``firedrake-configure`` with
 .. literalinclude:: petsc_configure_options.txt
    :language: text
 
+.. note::
+   If you install MPI through PETSc by passing ``--download-openmpi`` or
+   ``--download-mpich`` it is helpful to run the command::
+
+      $ export PATH=$PETSC_DIR/$PETSC_ARCH:$PATH
+
+
+   where ``PETSC_DIR=/path/to/petsc`` and ``PETSC_ARCH=arch-firedrake-default``.
+   This will allow the MPI executables (``mpicc``, ``mpiexec``, etc) installed by
+   PETSc to be found before any other versions installed on your machine.
+
 
 .. _install_firedrake:
 
@@ -224,6 +235,11 @@ install Firedrake. To do this perform the following steps:
       you have exactly followed the instructions up to this point this should
       already be the case.
 
+   .. note::
+      If you are using a non-system MPI it may be necessary to set ``LD_LIBRARY_PATH``
+      so that it can be detected by mpi4py. See `here <https://mpi4py.readthedocs.io/en/stable/install.html#linux>`__
+      for more information.
+
 #. Install Firedrake::
 
       $ pip install --no-binary h5py 'firedrake[check]'

docs/source/preconditioning.rst

@@ -69,6 +69,19 @@ multiplicatively within an MPI rank and additively between ranks.
    entities into lines or planes (useful for advection-dominated
    problems).
 
+.. note::
+   The additive Schwarz preconditioners listed here construct patches around
+   mesh entities.  Crucially, the mesh must have an overlapping parallel domain
+   decomposition that supports the patches. This is set via the
+   `distribution_parameters` kwarg of the :func:`.Mesh` constructor.  For
+   instance, vertex-star patches require ::
+
+      distribution_parameters["overlap_type"] = (DistributedMeshOverlapType.VERTEX, 1)
+
+   while Vanka patches require ::
+
+      distribution_parameters["overlap_type"] = (DistributedMeshOverlapType.VERTEX, 2)
+
 Multigrid methods
 =================
 

firedrake/adjoint_utils/blocks/function.py

@@ -202,15 +202,12 @@ def __init__(self, func, idx, ad_block_tag=None):
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                                prepared=None):
         eval_adj = firedrake.Cofunction(block_variable.output.function_space().dual())
-        if type(adj_inputs[0]) is firedrake.Cofunction:
-            eval_adj.sub(self.sub_idx).assign(adj_inputs[0])
-        else:
-            eval_adj.sub(self.sub_idx).assign(adj_inputs[0].function)
+        eval_adj.sub(self.sub_idx).assign(adj_inputs[0])
         return eval_adj
 
     def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
                                prepared=None):
-        return firedrake.Function.sub(tlm_inputs[0], self.sub_idx)
+        return tlm_inputs[0].sub(self.sub_idx)
 
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx,
@@ -220,9 +217,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
         return eval_hessian
 
     def recompute_component(self, inputs, block_variable, idx, prepared):
-        return maybe_disk_checkpoint(
-            firedrake.Function.sub(inputs[0], self.sub_idx)
-        )
+        return maybe_disk_checkpoint(inputs[0].sub(self.sub_idx))
 
     def __str__(self):
         return f"{self.get_dependencies()[0]}[{self.sub_idx}]"
@@ -264,23 +259,31 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
             adj_inputs[0].subfunctions[self.sub_idx].zero()
             return adj_inputs[0]
 
-    def evaluate_tlm(self, markings=False):
-        tlm_input = self.get_dependencies()[0].tlm_value
-        if tlm_input is None:
-            return
-        output = self.get_outputs()[0]
-        if markings and not output.marked_in_path:
-            return
-        fs = output.output.function_space()
-        f = type(output.output)(fs)
-        output.add_tlm_output(
-            type(output.output).assign(f.sub(self.sub_idx), tlm_input)
-        )
+    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
+        sub_tlm = tlm_inputs[0]
+        parent_in = tlm_inputs[1]
+
+        if sub_tlm is None and parent_in is None:
+            return None
+
+        output = self.get_outputs()[0].output
+        parent_out = type(output)(output.function_space())
+
+        if parent_in is not None:
+            parent_out.assign(parent_in)
+        if sub_tlm is not None:
+            parent_out.sub(self.sub_idx).assign(sub_tlm)
+
+        return parent_out
 
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx,
                                    relevant_dependencies, prepared=None):
-        return hessian_inputs[0]
+        if idx == 0:
+            return hessian_inputs[0].subfunctions[self.sub_idx].copy(deepcopy=True)
+        else:
+            hessian_inputs[0].subfunctions[self.sub_idx].zero()
+            return hessian_inputs[0]
 
     def recompute_component(self, inputs, block_variable, idx, prepared):
         sub_func = inputs[0]

firedrake/preconditioners/asm.py

@@ -4,6 +4,7 @@
 from firedrake.preconditioners.base import PCBase
 from firedrake.petsc import PETSc
 from firedrake.dmhooks import get_function_space
+from firedrake.mesh import DistributedMeshOverlapType
 from firedrake.logging import warning
 from tinyasm import _tinyasm as tinyasm
 from mpi4py import MPI
@@ -164,6 +165,8 @@ def get_patches(self, V):
         opts = PETSc.Options(self.prefix)
         depth = opts.getInt("construct_dim", default=0)
         ordering = opts.getString("mat_ordering_type", default="natural")
+        validate_overlap(mesh_unique, depth, "star")
+
         # Accessing .indices causes the allocation of a global array,
         # so we need to cache these for efficiency
         V_local_ises_indices = tuple(iset.indices for iset in V.dof_dset.local_ises)
@@ -241,8 +244,11 @@ def get_patches(self, V):
         ises = []
         if depth != -1:
             (start, end) = mesh_dm.getDepthStratum(depth)
+            patch_dim = depth
         else:
             (start, end) = mesh_dm.getHeightStratum(height)
+            patch_dim = mesh_dm.getDimension() - height
+        validate_overlap(mesh_unique, patch_dim, "vanka")
 
         for seed in range(start, end):
             # Only build patches over owned DoFs
@@ -469,6 +475,7 @@ def get_patches(self, V):
             else:
                 continue
 
+            validate_overlap(mesh_unique, base_depth, "star")
             start, end = mesh_dm.getDepthStratum(base_depth)
             for seed in range(start, end):
                 # Only build patches over owned DoFs
@@ -524,3 +531,26 @@ def get_patches(self, V):
                     iset = PETSc.IS().createGeneral(indices, comm=PETSc.COMM_SELF)
                     ises.append(iset)
         return ises
+
+
+def validate_overlap(mesh, patch_dim, patch_type):
+    if patch_type == "python":
+        return
+    patch_depth = {"pardecomp": 0, "star": 1, "vanka": 2}[patch_type]
+
+    tdim = mesh.topology_dm.getDimension()
+    overlap_entity, overlap_depth = mesh._distribution_parameters["overlap_type"]
+    overlap_dim = {
+        DistributedMeshOverlapType.VERTEX: 0,
+        DistributedMeshOverlapType.FACET: tdim-1,
+        DistributedMeshOverlapType.NONE: tdim,
+    }[overlap_entity]
+
+    if mesh.comm.size > 1:
+        if overlap_dim > patch_dim:
+            patch_entity = {0: "vertex", 1: "edge", 2: "face", tdim: "cell"}[patch_dim]
+            warning(f"{overlap_entity} does not support {patch_entity}-patches. "
+                    "Did you forget to set overlap_type in your mesh's distribution_parameters?")
+        if overlap_depth < patch_depth:
+            warning(f"Mesh overlap depth of {overlap_depth} does not support {patch_type}-patches. "
+                    "Did you forget to set overlap_type in your mesh's distribution_parameters?")

firedrake/preconditioners/patch.py

@@ -1,4 +1,5 @@
 from firedrake.preconditioners.base import PCBase, SNESBase, PCSNESBase
+from firedrake.preconditioners.asm import validate_overlap
 from firedrake.petsc import PETSc
 from firedrake.cython.patchimpl import set_patch_residual, set_patch_jacobian
 from firedrake.solving_utils import _SNESContext
@@ -784,15 +785,24 @@ def initialize(self, obj):
         if mesh_unique.cell_set._extruded:
             raise NotImplementedError("Not implemented on extruded meshes")
 
-        if "overlap_type" not in mesh_unique._distribution_parameters:
-            if mesh.comm.size > 1:
-                # Want to do
-                # warnings.warn("You almost surely want to set an overlap_type in your mesh's distribution_parameters.")
-                # but doesn't warn!
-                PETSc.Sys.Print("Warning: you almost surely want to set an overlap_type in your mesh's distribution_parameters.")
+        # Validate the mesh overlap
+        prefix = (obj.getOptionsPrefix() or "") + "patch_"
+        opts = PETSc.Options(prefix)
+        petsc_prefix = self._petsc_prefix
+        patch_type = opts.getString(f"{petsc_prefix}construct_type")
+        patch_dim = opts.getInt(f"{petsc_prefix}construct_dim", -1)
+        patch_codim = opts.getInt(f"{petsc_prefix}construct_codim", -1)
+        if patch_dim != -1:
+            assert patch_codim == -1, "Cannot set both dim and codim"
+        elif patch_codim != -1:
+            assert patch_dim == -1, "Cannot set both dim and codim"
+            patch_dim = self.plex.getDimension() - patch_codim
+        else:
+            patch_dim = 0
+        validate_overlap(mesh_unique, patch_dim, patch_type)
 
         patch = obj.__class__().create(comm=mesh.comm)
-        patch.setOptionsPrefix((obj.getOptionsPrefix() or "") + "patch_")
+        patch.setOptionsPrefix(prefix)
         self.configure_patch(patch, obj)
         patch.setType("patch")
 
@@ -948,6 +958,8 @@ def view(self, pc, viewer=None):
 
 class PatchPC(PCBase, PatchBase):
 
+    _petsc_prefix = "pc_patch_"
+
     def configure_patch(self, patch, pc):
         (A, P) = pc.getOperators()
         patch.setOperators(A, P)
@@ -960,6 +972,9 @@ def applyTranspose(self, pc, x, y):
 
 
 class PatchSNES(SNESBase, PatchBase):
+
+    _petsc_prefix = "snes_patch_"
+
     def configure_patch(self, patch, snes):
         patch.setTolerances(max_it=1)
         patch.setConvergenceTest("skip")