Rework PMGBase.coarsen() and add NonlinearVariationalProblem.rediscretise()

firedrake/mg/ufl_utils.py

@@ -314,7 +314,7 @@ def coarsen_callback(fine, coarse):
     F = self(problem.F, self, coefficient_mapping=coefficient_mapping)
     J = self(problem.J, self, coefficient_mapping=coefficient_mapping)
     Jp = self(problem.Jp, self, coefficient_mapping=coefficient_mapping)
-    u = coefficient_mapping[problem.u_restrict]
+    u = self(problem.u_restrict, self, coefficient_mapping=coefficient_mapping)
 
     new_problem = firedrake.NonlinearVariationalProblem(
         F, u, bcs=bcs, J=J, Jp=Jp, objective=E, is_linear=problem.is_linear,

firedrake/nullspace.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import numpy
 
 from pyop2.mpi import COMM_WORLD
 
+import firedrake
 from firedrake import function
 from firedrake.logging import warning
 from firedrake.matrix import MatrixBase
@@ -97,6 +100,24 @@ def orthonormalize(self):
         self.check_orthogonality()
         self._ad_orthogonalized = True
 
+    def rediscretise(self, function_space: firedrake.functionspaceimpl.WithGeometryBase) -> VectorSpaceBasis:
+        r"""Reconstruct this basis on a new function space.
+
+        Parameters
+        ----------
+        function_space
+            the new :class:`~.FunctionSpace`.
+
+        Returns
+        -------
+        VectorSpaceBasis
+            The basis vectors interpolated onto `function_space` and re-orthonormalized.
+        """
+        vecs = [function.Function(function_space).interpolate(vec) for vec in self._vecs]
+        new_basis = VectorSpaceBasis(vecs, constant=self._constant, comm=self.comm)
+        new_basis.orthonormalize()
+        return new_basis
+
     @PETSc.Log.EventDecorator()
     def orthogonalize(self, b):
         r"""Orthogonalize ``b`` with respect to this :class:`.VectorSpaceBasis`.
@@ -239,6 +260,27 @@ def __init__(self, function_space, bases):
         self._bases = bases
         self._nullspace = None
 
+    def rediscretise(self, function_space: firedrake.functionspaceimpl.WithGeometryBase) -> MixedVectorSpaceBasis:
+        r"""Reconstruct this basis on a new mixed function space.
+
+        Parameters
+        ----------
+        function_space
+            the new :class:`~.FunctionSpace`.
+
+        Returns
+        -------
+        MixedVectorSpaceBasis
+            The bases reconstructed on the sub-spaces of `function_space`.
+        """
+        bases = []
+        for V_, basis in zip(function_space, self._bases):
+            if isinstance(basis, VectorSpaceBasis):
+                bases.append(basis.rediscretise(V_))
+            else:
+                bases.append(function_space.sub(basis.index))
+        return MixedVectorSpaceBasis(function_space, bases)
+
     def _build_monolithic_basis(self):
         r"""Build a basis for the complete mixed space.
 

firedrake/preconditioners/low_order.py

@@ -8,15 +8,15 @@ class P1PC(PMGPC):
     def coarsen_element(self, ele):
         if super().max_degree(ele) <= self.coarse_degree:
             raise ValueError
-        return super().reconstruct_degree(ele, self.coarse_degree)
+        return ele.reconstruct(degree=self.coarse_degree)
 
 
 class P1SNES(PMGSNES):
     """A two-level nonlinear solver with agressive p-coarsening."""
     def coarsen_element(self, ele):
         if super().max_degree(ele) <= self.coarse_degree:
             raise ValueError
-        return super().reconstruct_degree(ele, self.coarse_degree)
+        return ele.reconstruct(degree=self.coarse_degree)
 
 
 class LORPC(PMGPC):
@@ -33,6 +33,6 @@ def coarsen_element(self, ele):
             iso_variant = f"iso({degree})"
         else:
             iso_variant = f"{variant},iso({degree})"
-        cele = super().reconstruct_degree(ele, self.coarse_degree)
+        cele = ele.reconstruct(degree=self.coarse_degree)
         cele = cele.reconstruct(variant=iso_variant)
         return cele