Add PETSc pressure nullspace support to Stokes

src/underworld3/cython/petsc_generic_snes_solvers.pyx

@@ -1,5 +1,6 @@
 from xmlrpc.client import Boolean
 
+import numpy as np
 import sympy
 from sympy import sympify
 
@@ -447,6 +448,10 @@ class SolverBaseClass(uw_object):
                 self.dm = None  # Should be able to avoid nuking this if we
                             # can insert new functions in template (surface integrals problematic in
                             # the current implementation )
+                if hasattr(self, "_stokes_nullspace"):
+                    self._stokes_nullspace = None
+                if hasattr(self, "_stokes_nullspace_basis"):
+                    self._stokes_nullspace_basis = ()
 
         # This is a workaround for some problem in the PETSc machinery
         # where we need a surface integral term somewhere on every process
@@ -2829,6 +2834,10 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
         self.boundary_conditions = False
         # self._constitutive_model = None
         self._saddle_preconditioner = None
+        self._petsc_use_pressure_nullspace = False
+        self._petsc_velocity_nullspace_basis = ()
+        self._stokes_nullspace = None
+        self._stokes_nullspace_basis = ()
 
         # Construct strainrate tensor for future usage.
         # Grab gradients, and let's switch out to sympy.Matrix notation
@@ -3088,6 +3097,265 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
         self.is_setup = False
         self._saddle_preconditioner = function
 
+    @property
+    def petsc_use_pressure_nullspace(self):
+        """
+        Enable PETSc handling of the constant-pressure nullspace.
+
+        When enabled, the solver attaches the constant-pressure mode to
+        the coupled Stokes nullspace basis before solve. Additional
+        user-supplied velocity nullspace modes can be configured through
+        ``petsc_velocity_nullspace_basis``.
+
+        For free-slip shell problems this is typically used together with
+        the rigid-body rotation modes documented on
+        ``petsc_velocity_nullspace_basis``.
+
+        Examples
+        --------
+        2-D annulus with pressure gauge only:
+        >>> stokes.petsc_use_pressure_nullspace = True
+
+        2-D annulus with pressure plus rigid rotation:
+        >>> x, y = mesh.X
+        >>> stokes.petsc_use_pressure_nullspace = True
+        >>> stokes.petsc_velocity_nullspace_basis = [sympy.Matrix([-y, x])]
+
+        3-D spherical shell with pressure plus the three rigid rotations:
+        >>> x, y, z = mesh.X
+        >>> stokes.petsc_use_pressure_nullspace = True
+        >>> stokes.petsc_velocity_nullspace_basis = [
+        ...     sympy.Matrix([0, -z, y]),
+        ...     sympy.Matrix([z, 0, -x]),
+        ...     sympy.Matrix([-y, x, 0]),
+        ... ]
+        """
+        return self._petsc_use_pressure_nullspace
+
+    @petsc_use_pressure_nullspace.setter
+    def petsc_use_pressure_nullspace(self, value):
+        self._petsc_use_pressure_nullspace = bool(value)
+        self._reset_stokes_nullspace()
+        self.is_setup = False
+
+    @property
+    def petsc_velocity_nullspace_basis(self):
+        """
+        Optional exact velocity nullspace modes for the coupled Stokes solve.
+
+        Each entry must be a vector-valued SymPy expression defined in the
+        mesh coordinate system and representing an exact null mode of the
+        configured Stokes operator. Typical examples are rigid-body rotation
+        modes for annulus or spherical-shell free-slip problems.
+
+        For centered shell geometries with exact free-slip / no-penetration
+        boundary conditions, the rigid-body rotation modes are:
+
+        - 2-D annulus: one mode, ``(-y, x)``, equivalent to ``r e_theta``
+        - 3-D spherical shell: three modes,
+          ``(0, -z, y)``, ``(z, 0, -x)``, and ``(-y, x, 0)``
+
+        These are the velocity fields generated by rigid rotations
+        ``u = omega x x``. They are tangent to concentric circles / spheres
+        and have zero strain rate, so they are exact velocity null modes for
+        the free-slip shell Stokes operator.
+
+        They are not exact null modes when the boundary conditions select a
+        specific tangential velocity, for example:
+
+        - essential velocity boundary conditions
+        - penalty boundary conditions on the full velocity error
+          ``u - u_analytic``
+
+        To remove shell nullspaces in Stokes, set the pressure mode and then
+        provide the exact rotation basis:
+
+        Examples
+        --------
+        2-D annulus:
+        >>> x, y = mesh.X
+        >>> stokes.petsc_use_pressure_nullspace = True
+        >>> stokes.petsc_velocity_nullspace_basis = [sympy.Matrix([-y, x])]
+
+        3-D spherical shell:
+        >>> x, y, z = mesh.X
+        >>> stokes.petsc_use_pressure_nullspace = True
+        >>> stokes.petsc_velocity_nullspace_basis = [
+        ...     sympy.Matrix([0, -z, y]),
+        ...     sympy.Matrix([z, 0, -x]),
+        ...     sympy.Matrix([-y, x, 0]),
+        ... ]
+        """
+        return self._petsc_velocity_nullspace_basis
+
+    @petsc_velocity_nullspace_basis.setter
+    def petsc_velocity_nullspace_basis(self, modes):
+        if modes is None:
+            modes = ()
+
+        velocity_modes = []
+        for mode in modes:
+            matrix_mode = sympy.Matrix(mode)
+            if matrix_mode.shape == (1, self.mesh.dim):
+                matrix_mode = matrix_mode.T
+            if matrix_mode.shape != (self.mesh.dim, 1):
+                raise ValueError(
+                    "Each petsc_velocity_nullspace_basis mode must have shape "
+                    f"({self.mesh.dim}, 1) or (1, {self.mesh.dim}); got {matrix_mode.shape}."
+                )
+            velocity_modes.append(matrix_mode)
+
+        self._petsc_velocity_nullspace_basis = tuple(velocity_modes)
+        self._reset_stokes_nullspace()
+        self.is_setup = False
+
+    def _reset_stokes_nullspace(self):
+        self._stokes_nullspace = None
+        self._stokes_nullspace_basis = ()
+
+    def _pressure_dirichlet_bcs(self):
+        """Return essential boundary conditions applied to the pressure field."""
+
+        pressure_field_id = getattr(getattr(self, "p", None), "field_id", None)
+        if pressure_field_id is None:
+            raise RuntimeError("Pressure field is unavailable; cannot inspect pressure Dirichlet BCs.")
+
+        return [bc for bc in self.essential_bcs if bc.f_id == pressure_field_id]
+
+    def _build_pressure_nullspace_vector(self):
+        """Create the constant-pressure basis vector for the coupled Stokes DM."""
+
+        template_vec = self.dm.getGlobalVec()
+        try:
+            null_vec = template_vec.duplicate()
+        finally:
+            self.dm.restoreGlobalVec(template_vec)
+
+        null_vec.set(0.0)
+
+        pressure_is = self._subdict["pressure"][0]
+        pressure_subvec = null_vec.getSubVector(pressure_is)
+        pressure_subvec.set(1.0)
+        null_vec.restoreSubVector(pressure_is, pressure_subvec)
+
+        return null_vec
+
+    def _build_velocity_nullspace_vector(self, mode):
+        """Create a velocity nullspace basis vector from a user-supplied mode."""
+
+        mode_values = np.asarray(uw.function.evaluate(mode, self.u.coords_nd), dtype=np.float64)
+        if mode_values.ndim == 1:
+            mode_values = mode_values.reshape(-1, 1)
+        elif mode_values.ndim > 2:
+            mode_values = mode_values.reshape(mode_values.shape[0], -1)
+
+        if mode_values.shape != (self.u.coords_nd.shape[0], self.u.num_components):
+            raise ValueError(
+                "Velocity nullspace mode evaluation must return an array with shape "
+                f"({self.u.coords_nd.shape[0]}, {self.u.num_components}); got {mode_values.shape}."
+            )
+
+        template_vec = self.dm.getGlobalVec()
+        try:
+            null_vec = template_vec.duplicate()
+        finally:
+            self.dm.restoreGlobalVec(template_vec)
+
+        null_vec.set(0.0)
+
+        velocity_is, velocity_subdm = self._subdict["velocity"]
+        velocity_basis = self.u.vec.duplicate()
+        try:
+            velocity_basis.array[:] = mode_values.reshape(velocity_basis.array.shape)
+            velocity_subvec = null_vec.getSubVector(velocity_is)
+            velocity_subdm.localToGlobal(velocity_basis, velocity_subvec, addv=False)
+            null_vec.restoreSubVector(velocity_is, velocity_subvec)
+        finally:
+            velocity_basis.destroy()
+
+        return null_vec
+
+    def _build_stokes_nullspace(self):
+        """Create the configured coupled Stokes nullspace basis."""
+
+        basis_vectors = []
+
+        if self._petsc_use_pressure_nullspace:
+            basis_vectors.append(self._build_pressure_nullspace_vector())
+
+        for mode in self._petsc_velocity_nullspace_basis:
+            basis_vectors.append(self._build_velocity_nullspace_vector(mode))
+
+        if not basis_vectors:
+            return None
+
+        orthonormal_basis = []
+        for basis_vec in basis_vectors:
+            for orth_vec in orthonormal_basis:
+                basis_vec.axpy(-orth_vec.dot(basis_vec), orth_vec)
+
+            basis_norm = basis_vec.norm()
+            if np.isclose(basis_norm, 0.0):
+                raise ValueError(
+                    "Configured Stokes nullspace basis contains a dependent or zero mode."
+                )
+
+            basis_vec.scale(1.0 / basis_norm)
+            orthonormal_basis.append(basis_vec)
+
+        self._stokes_nullspace_basis = tuple(orthonormal_basis)
+        self._stokes_nullspace = PETSc.NullSpace().create(
+            constant=False,
+            vectors=self._stokes_nullspace_basis,
+            comm=self.dm.comm,
+        )
+
+        return self._stokes_nullspace
+
+    def _attach_stokes_nullspace(self):
+        """Attach the configured coupled Stokes nullspace to the solver matrices."""
+
+        if not self._petsc_use_pressure_nullspace and not self._petsc_velocity_nullspace_basis:
+            return
+
+        pressure_bcs = self._pressure_dirichlet_bcs()
+        if pressure_bcs:
+            boundaries = ", ".join(sorted({bc.boundary for bc in pressure_bcs}))
+            raise ValueError(
+                "PETSc Stokes nullspace support requires the pressure field to be "
+                f"free of Dirichlet boundary conditions. Found pressure Dirichlet BCs on: {boundaries}"
+            )
+
+        if "pressure" not in self._subdict or "velocity" not in self._subdict:
+            raise RuntimeError("Velocity/pressure field decomposition is unavailable; cannot attach nullspace.")
+
+        self.snes.setUp()
+
+        jacobian = self.snes.getJacobian()
+        operator_matrix = jacobian[0]
+        preconditioner_matrix = jacobian[1] if len(jacobian) > 1 else None
+
+        nullspace = self._stokes_nullspace
+        if nullspace is None:
+            nullspace = self._build_stokes_nullspace()
+
+        if nullspace is None:
+            return
+
+        operator_matrix.setNullSpace(nullspace)
+        operator_matrix.setTransposeNullSpace(nullspace)
+
+        if preconditioner_matrix is not None:
+            preconditioner_matrix.setNullSpace(nullspace)
+            preconditioner_matrix.setTransposeNullSpace(nullspace)
+
+        if self.verbose and uw.mpi.rank == 0:
+            print(
+                f"Stokes Saddle Pt ({self.name}): attached Stokes nullspace with "
+                f"{len(self._stokes_nullspace_basis)} basis mode(s)",
+                flush=True,
+            )
+
 
     ## F0, F1 should be f0 and F1, (pf0 for Saddles can be added here)
     ## don't add new ones uf0, uF1 are redundant
@@ -3744,6 +4012,8 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
         for index,name in enumerate(names):
             self._subdict[name] = (isets[index],dms[index])
 
+        self._attach_stokes_nullspace()
+
         self.is_setup = True
         self.constitutive_model._solver_is_setup = True
 
@@ -3846,6 +4116,9 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
             self.petsc_options.setValue("snes_max_it", 0)
             self.snes.setType("nrichardson")
             self.snes.setFromOptions()
+            # PETSc may rebuild operator state after setFromOptions(), so reattach
+            # the configured Stokes nullspace before each solve path.
+            self._attach_stokes_nullspace()
             self.snes.solve(None, gvec)
 
             # with self.mesh.access():
@@ -3871,6 +4144,7 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
             self.snes.atol = self.atol
             self.snes.setType("nrichardson")
             self.snes.setFromOptions()
+            self._attach_stokes_nullspace()
             self.snes.solve(None, gvec)
             self._warn_on_divergence(phase="picard")
 
@@ -3880,6 +4154,7 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
             self.snes.atol = self.atol
             self.petsc_options.setValue("snes_max_it", snes_max_it)
             self.snes.setFromOptions()
+            self._attach_stokes_nullspace()
             self.snes.solve(None, gvec)
 
         else:
@@ -3889,6 +4164,7 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
             self.snes.atol = self.atol
             self.petsc_options.setValue("snes_max_it", snes_max_it)
             self.snes.setFromOptions()
+            self._attach_stokes_nullspace()
             self.snes.solve(None, gvec)
 
         cdef DM dm = self.dm