Fix spherical shell labels and add Stokes residual diagnostics

src/underworld3/cython/petsc_compat.h

@@ -45,6 +45,39 @@ PetscErrorCode DMSetAuxiliaryVec_UW(DM dm, DMLabel label, PetscInt value, PetscI
     return DMSetAuxiliaryVec(dm, label, value, part, aux);
 }
 
+PetscErrorCode UW_PetscDSGetBoundaryWeakForm(PetscDS ds, PetscInt bd, PetscWeakForm *wf)
+{
+    PetscCall(PetscDSGetBoundary(ds, bd, wf, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL));
+    return PETSC_SUCCESS;
+}
+
+PetscErrorCode UW_DMPlexComputeResidualByKeyVolumeOnly(DM dm, PetscFormKey key, IS cellIS, PetscReal time, Vec locX, Vec locX_t, PetscReal t, Vec locF, void *ctx)
+{
+    DM vdm = NULL;
+    PetscSection section = NULL;
+    PetscSection global_section = NULL;
+    PetscDS ds = NULL;
+    PetscDS vds = NULL;
+    Vec aux = NULL;
+
+    PetscCall(DMClone(dm, &vdm));
+    PetscCall(DMGetLocalSection(dm, &section));
+    PetscCall(DMSetLocalSection(vdm, section));
+    PetscCall(DMGetGlobalSection(dm, &global_section));
+    PetscCall(DMSetGlobalSection(vdm, global_section));
+    PetscCall(DMCopyFields(dm, PETSC_DETERMINE, PETSC_DETERMINE, vdm));
+    PetscCall(DMCreateDS(vdm));
+    PetscCall(DMGetDS(dm, &ds));
+    PetscCall(DMGetDS(vdm, &vds));
+    PetscCall(PetscDSCopyConstants(ds, vds));
+    PetscCall(PetscDSCopyEquations(ds, vds));
+    PetscCall(DMGetAuxiliaryVec(dm, key.label, key.value, key.part, &aux));
+    if (aux) PetscCall(DMSetAuxiliaryVec(vdm, key.label, key.value, key.part, aux));
+    PetscCall(DMPlexComputeResidualByKey(vdm, key, cellIS, time, locX, locX_t, t, locF, ctx));
+    PetscCall(DMDestroy(&vdm));
+    return PETSC_SUCCESS;
+}
+
 // copy paste function signitures from $PETSC_DIR/include/petscds.h - would be nice to automate this.
 #define UW_SIG_F0 PetscInt, PetscInt, PetscInt, const PetscInt[], const PetscInt[], const PetscScalar[], const PetscScalar[], const PetscScalar[], const PetscInt[], const PetscInt[], const PetscScalar[], const PetscScalar[], const PetscScalar[], PetscReal, const PetscReal[], const PetscReal[], PetscInt, const PetscScalar[], PetscScalar[]
 #define UW_SIG_G0 PetscInt, PetscInt, PetscInt, const PetscInt[], const PetscInt[], const PetscScalar[], const PetscScalar[], const PetscScalar[], const PetscInt[], const PetscInt[], const PetscScalar[], const PetscScalar[], const PetscScalar[], PetscReal, PetscReal, const PetscReal[], const PetscReal[], PetscInt, const PetscScalar[], PetscScalar[]

src/underworld3/cython/petsc_extras.pxi

@@ -31,10 +31,20 @@ cdef CHKERRQ(PetscErrorCode ierr):
     cdef int interr = <int>ierr
     if ierr != 0: raise RuntimeError(f"PETSc error code '{interr}' was encountered.\nhttps://www.mcs.anl.gov/petsc/petsc-current/include/petscerror.h.html")
 
+cdef extern from "petscdstypes.h":
+    ctypedef void *PetscWeakForm "PetscWeakForm"
+    ctypedef struct PetscFormKey:
+        PetscDMLabel label
+        PetscInt value
+        PetscInt field
+        PetscInt part
+
 cdef extern from "petsc_compat.h":
 
     PetscErrorCode PetscDSAddBoundary_UW( PetscDM, DMBoundaryConditionType, const char[], const char[] , PetscInt, PetscInt, PetscInt *, void (*)(), void (*)(), PetscInt, const PetscInt *, void *)
     PetscErrorCode DMSetAuxiliaryVec_UW(PetscDM, PetscDMLabel, PetscInt, PetscInt, PetscVec)
+    PetscErrorCode UW_PetscDSGetBoundaryWeakForm(PetscDS, PetscInt, PetscWeakForm *)
+    PetscErrorCode UW_DMPlexComputeResidualByKeyVolumeOnly( PetscDM, PetscFormKey, PetscIS, PetscReal, PetscVec, PetscVec, PetscReal, PetscVec, void *)
     # PetscErrorCode UW_PetscDSSetBdResidual(PetscDS, PetscDMLabel, PetscInt, PetscInt, PetscInt, PetscInt, void*, PetscInt, void*)
 
     PetscErrorCode UW_PetscDSSetBdResidual(PetscDS, PetscDMLabel, PetscInt, PetscInt, PetscInt, PetscInt, void*, void*)
@@ -52,6 +62,9 @@ cdef extern from "petsc_compat.h":
 cdef extern from "petsc.h" nogil:
     PetscErrorCode PetscDSSetConstants(PetscDS, PetscInt, const PetscScalar[])
     PetscErrorCode DMPlexSNESComputeBoundaryFEM( PetscDM, void *, void *)
+    PetscErrorCode DMPlexSNESComputeResidualFEM( PetscDM, PetscVec, PetscVec, void *)
+    PetscErrorCode DMPlexComputeResidualByKey( PetscDM, PetscFormKey, PetscIS, PetscReal, PetscVec, PetscVec, PetscReal, PetscVec, void *)
+    PetscErrorCode DMPlexComputeBdResidualSingle( PetscDM, PetscWeakForm, PetscFormKey, PetscVec, PetscVec, PetscReal, PetscVec )
     # PetscErrorCode DMPlexSetSNESLocalFEM( PetscDM, void *, void *, void *)
     # PetscErrorCode DMPlexSetSNESLocalFEM( PetscDM, PetscBool, void *)
     PetscErrorCode DMPlexComputeGeometryFVM( PetscDM dm, PetscVec *cellgeom, PetscVec *facegeom)
@@ -62,6 +75,7 @@ cdef extern from "petsc.h" nogil:
     PetscErrorCode PetscDSSetJacobian( PetscDS, PetscInt, PetscInt, PetscDSJacobianFn, PetscDSJacobianFn, PetscDSJacobianFn, PetscDSJacobianFn)
     PetscErrorCode PetscDSSetJacobianPreconditioner( PetscDS, PetscInt, PetscInt, PetscDSJacobianFn, PetscDSJacobianFn, PetscDSJacobianFn, PetscDSJacobianFn)
     PetscErrorCode PetscDSSetResidual( PetscDS, PetscInt, PetscDSResidualFn, PetscDSResidualFn )
+    PetscErrorCode PetscDSGetWeakForm( PetscDS, PetscWeakForm * )
 
     PetscErrorCode PetscDSSetBdJacobian( PetscDS, PetscInt, PetscInt, PetscDSBdJacobianFn, PetscDSBdJacobianFn, PetscDSBdJacobianFn, PetscDSBdJacobianFn)
     PetscErrorCode PetscDSSetBdJacobianPreconditioner( PetscDS, PetscInt, PetscInt, PetscDSBdJacobianFn, PetscDSBdJacobianFn, PetscDSBdJacobianFn, PetscDSBdJacobianFn)
@@ -83,6 +97,7 @@ cdef extern from "petsc.h" nogil:
     PetscErrorCode DMGetRegionDS(PetscDM dm, PetscDMLabel label, PetscIS *fields, PetscDS *ds, PetscDS *dsIn)
     PetscErrorCode DMGetRegionNumDS(PetscDM dm, PetscInt num, PetscDMLabel *label, PetscIS *fields, PetscDS *ds, PetscDS *dsIn)
     PetscErrorCode DMSetRegionNumDS(PetscDM dm, PetscInt num, PetscDMLabel label, PetscIS fields, PetscDS ds, PetscDS dsIn)
+    PetscErrorCode DMGetDS(PetscDM dm, PetscDS *ds)
     PetscErrorCode DMGetNumDS(PetscDM dm, PetscInt *num)
     PetscErrorCode DMGetCellDS(PetscDM dm, PetscInt point, PetscDS *ds, PetscDS *dsIn)
     PetscErrorCode PetscDSSetCoordinateDimension(PetscDS ds, PetscInt dim)

src/underworld3/cython/petsc_generic_snes_solvers.pyx

@@ -6682,6 +6682,270 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
         if uw.mpi.rank == 0 and self.verbose:
             print(f"Region DS: inactive region '{label_name}' gets trivial DS", flush=True)
 
+    def compute_volume_residual_fields(
+        self,
+        time=None,
+        verbose=False,
+        cell_indices=None,
+        residual_field_id=None,
+        include_boundary_terms=False,
+    ):
+        """Return PETSc FEM residual fields in each solver field's local layout.
+
+        This is a low-level diagnostic hook for post-processing derived
+        boundary quantities such as consistent-boundary-flux traction. By
+        default it calls PETSc's ``DMPlexSNESComputeResidualFEM`` directly. If
+        ``cell_indices`` is supplied, it instead calls
+        a UW wrapper around ``DMPlexComputeResidualByKey`` on a cloned DM with
+        a copied ``PetscDS`` that has no registered boundary objects, so the
+        selected-cell path returns volume terms only. Set
+        ``include_boundary_terms=True`` to call PETSc's original keyed
+        residual behavior, which appends registered boundary residuals. The
+        returned arrays are local to each rank and have the same flat layout as
+        the corresponding MeshVariable PETSc vector.
+        """
+        cdef DM _time_dm_residual
+        cdef DM dm
+        cdef Vec xvec
+        cdef Vec fvec
+        cdef PetscFormKey key
+        cdef IS ccell_is
+        cdef PetscReal residual_time = 0.0
+        cdef PetscReal implicit_form_time = <PetscReal>-1.7976931348623157e308
+
+        self._build(verbose, False, None)
+
+        if time is not None:
+            if hasattr(time, 'magnitude') or hasattr(time, '_pint_qty'):
+                t_nd = float(uw.non_dimensionalise(time))
+            else:
+                t_nd = float(time)
+            _time_dm_residual = self.dm
+            UW_DMSetTime(_time_dm_residual.dm, t_nd)
+            residual_time = <PetscReal>t_nd
+
+        self.mesh.update_lvec()
+        self.dm.setAuxiliaryVec(self.mesh.lvec, None)
+        self._update_constants()
+
+        gvec = self.dm.getGlobalVec()
+        xlocal = self.dm.getLocalVec()
+        flocal = self.dm.getLocalVec()
+        gvec.setArray(0.0)
+        xlocal.setArray(0.0)
+        flocal.setArray(0.0)
+
+        try:
+            for name, var in self.fields.items():
+                sgvec = gvec.getSubVector(self._subdict[name][0])
+                subdm = self._subdict[name][1]
+                subdm.localToGlobal(var.vec, sgvec)
+                gvec.restoreSubVector(self._subdict[name][0], sgvec)
+
+            self.dm.globalToLocal(gvec, xlocal)
+
+            dm = self.dm
+            xvec = xlocal
+            fvec = flocal
+            if cell_indices is None:
+                CHKERRQ(DMPlexSNESComputeResidualFEM(dm.dm, xvec.vec, fvec.vec, NULL))
+            else:
+                if residual_field_id is None:
+                    residual_field_id = 0
+                cell_is = PETSc.IS().createGeneral(
+                    list(cell_indices), comm=PETSc.COMM_SELF
+                )
+                try:
+                    ccell_is = cell_is
+                    key.label = NULL
+                    key.value = 0
+                    key.field = <PetscInt>residual_field_id
+                    key.part = 0
+                    if include_boundary_terms:
+                        CHKERRQ(DMPlexComputeResidualByKey(
+                            dm.dm, key, ccell_is.iset, implicit_form_time,
+                            xvec.vec, NULL, residual_time, fvec.vec, NULL,
+                        ))
+                    else:
+                        CHKERRQ(UW_DMPlexComputeResidualByKeyVolumeOnly(
+                            dm.dm, key, ccell_is.iset, implicit_form_time,
+                            xvec.vec, NULL, residual_time, fvec.vec, NULL,
+                        ))
+                finally:
+                    cell_is.destroy()
+
+            local_section = self.dm.getLocalSection()
+            pStart, pEnd = local_section.getChart()
+            out = {}
+
+            for name, var in self.fields.items():
+                field_id = getattr(var, "_solver_field_id", None)
+                if field_id is None:
+                    field_id = getattr(var, "field_id", None)
+                if field_id is None:
+                    continue
+
+                is_field = None
+                created_is_field = False
+                if name == "velocity" and getattr(self, "_velocity_is", None) is not None:
+                    is_field = self._velocity_is
+                elif name == "pressure" and getattr(self, "_pressure_is", None) is not None:
+                    is_field = self._pressure_is
+                elif getattr(self, "_multiplier_is", None) is not None and name in self._multiplier_is:
+                    is_field = self._multiplier_is[name]
+                else:
+                    indices = []
+                    for point in range(pStart, pEnd):
+                        dof = local_section.getFieldDof(point, field_id)
+                        if dof > 0:
+                            offset = local_section.getFieldOffset(point, field_id)
+                            for i in range(dof):
+                                indices.append(offset + i)
+
+                    is_field = PETSc.IS().createGeneral(indices, comm=PETSc.COMM_SELF)
+                    created_is_field = True
+
+                try:
+                    subvec = flocal.getSubVector(is_field)
+                    try:
+                        out[name] = np.array(subvec.array, copy=True)
+                    finally:
+                        flocal.restoreSubVector(is_field, subvec)
+                finally:
+                    if created_is_field:
+                        is_field.destroy()
+
+            return out
+        finally:
+            self.dm.restoreLocalVec(flocal)
+            self.dm.restoreLocalVec(xlocal)
+            self.dm.restoreGlobalVec(gvec)
+
+    def compute_boundary_residual_fields(self, boundary, time=None, verbose=False, residual_field_id=0):
+        """Return the registered FEM boundary residual for one named boundary.
+
+        This is a low-level diagnostic hook for weak-boundary-condition
+        debugging. It assembles PETSc's boundary residual terms registered on
+        ``boundary`` through ``DMPlexComputeBdResidualSingle``. For Nitsche
+        free slip, this includes the full registered weak boundary residual,
+        not only the scalar penalty term. The returned arrays are local to each
+        rank and have the same flat layout as the corresponding MeshVariable
+        PETSc vector.
+        """
+        cdef DM _time_dm_boundary_residual
+        cdef DM dm
+        cdef Vec xvec
+        cdef Vec fvec
+        cdef PetscFormKey key
+        cdef PetscDS ds
+        cdef PetscWeakForm wf
+        cdef DMLabel c_label
+
+        self._build(verbose, False, None)
+
+        boundary_bc = None
+        for bc in self.natural_bcs:
+            if bc.boundary == boundary and bc.f_id == residual_field_id:
+                boundary_bc = bc
+                break
+        if boundary_bc is None:
+            raise ValueError(
+                f"No natural/Nitsche boundary residual is registered for "
+                f"boundary '{boundary}' and field {residual_field_id}."
+            )
+
+        if time is not None:
+            if hasattr(time, 'magnitude') or hasattr(time, '_pint_qty'):
+                t_nd = float(uw.non_dimensionalise(time))
+            else:
+                t_nd = float(time)
+            _time_dm_boundary_residual = self.dm
+            UW_DMSetTime(_time_dm_boundary_residual.dm, t_nd)
+
+        self.mesh.update_lvec()
+        self.dm.setAuxiliaryVec(self.mesh.lvec, None)
+        self._update_constants()
+
+        gvec = self.dm.getGlobalVec()
+        xlocal = self.dm.getLocalVec()
+        flocal = self.dm.getLocalVec()
+        gvec.setArray(0.0)
+        xlocal.setArray(0.0)
+        flocal.setArray(0.0)
+
+        try:
+            for name, var in self.fields.items():
+                sgvec = gvec.getSubVector(self._subdict[name][0])
+                subdm = self._subdict[name][1]
+                subdm.localToGlobal(var.vec, sgvec)
+                gvec.restoreSubVector(self._subdict[name][0], sgvec)
+
+            self.dm.globalToLocal(gvec, xlocal)
+
+            dm = self.dm
+            xvec = xlocal
+            fvec = flocal
+            CHKERRQ(DMGetDS(dm.dm, &ds))
+            CHKERRQ(UW_PetscDSGetBoundaryWeakForm(
+                ds, <PetscInt>boundary_bc.PETScID, &wf,
+            ))
+
+            c_label = self.dm.getLabel("UW_Boundaries")
+            key.label = c_label.dmlabel
+            key.value = <PetscInt>boundary_bc.boundary_label_val
+            key.field = <PetscInt>residual_field_id
+            key.part = 0
+            CHKERRQ(DMPlexComputeBdResidualSingle(
+                dm.dm, wf, key, xvec.vec, NULL, 0.0, fvec.vec,
+            ))
+
+            local_section = self.dm.getLocalSection()
+            pStart, pEnd = local_section.getChart()
+            out = {}
+
+            for name, var in self.fields.items():
+                field_id = getattr(var, "_solver_field_id", None)
+                if field_id is None:
+                    field_id = getattr(var, "field_id", None)
+                if field_id is None:
+                    continue
+
+                is_field = None
+                created_is_field = False
+                if name == "velocity" and getattr(self, "_velocity_is", None) is not None:
+                    is_field = self._velocity_is
+                elif name == "pressure" and getattr(self, "_pressure_is", None) is not None:
+                    is_field = self._pressure_is
+                elif getattr(self, "_multiplier_is", None) is not None and name in self._multiplier_is:
+                    is_field = self._multiplier_is[name]
+                else:
+                    indices = []
+                    for point in range(pStart, pEnd):
+                        dof = local_section.getFieldDof(point, field_id)
+                        if dof > 0:
+                            offset = local_section.getFieldOffset(point, field_id)
+                            for i in range(dof):
+                                indices.append(offset + i)
+
+                    is_field = PETSc.IS().createGeneral(indices, comm=PETSc.COMM_SELF)
+                    created_is_field = True
+
+                try:
+                    subvec = flocal.getSubVector(is_field)
+                    try:
+                        out[name] = np.array(subvec.array, copy=True)
+                    finally:
+                        flocal.restoreSubVector(is_field, subvec)
+                finally:
+                    if created_is_field:
+                        is_field.destroy()
+
+            return out
+        finally:
+            self.dm.restoreLocalVec(flocal)
+            self.dm.restoreLocalVec(xlocal)
+            self.dm.restoreGlobalVec(gvec)
+
     def _ensure_local_field_index_sets(self, clvec, local_section):
         """Build (once) and cache the LOCAL index sets that decompose a parent-DM
         local vector into the per-field MeshVariable storage: velocity, pressure