WIP: Calculate Jacobian coloring using PetscOperator matrices

include/bout/petsc_operators.hxx

@@ -8,7 +8,6 @@
 #include "bout/boutexception.hxx"
 #include "bout/field3d.hxx"
 #include "bout/mesh.hxx"
-#include "bout/output.hxx"
 #include "bout/output_bout_types.hxx"
 #include "bout/petsc_interface.hxx"
 #include "bout/petsclib.hxx"
@@ -45,6 +44,10 @@ struct ForwardLegSpaceTag {};
 /// leg space. See @ref ForwardLegSpaceTag and @ref CellSpaceTag.
 struct BackwardLegSpaceTag {};
 
+// Forward declare
+template <typename OutSpaceTag, typename InSpaceTag>
+class PetscOperator;
+
 /// @brief Bidirectional index mapping between mesh-file stored numbering and PETSc
 ///        row ownership.
 ///
@@ -275,6 +278,37 @@ public:
     }
   }
 
+  /// @brief Create a PETSc IS containing the global PETSc indices of all
+  ///        locally owned evolving interior cells, in the order that
+  ///        mapOwnedInteriorCells visits them.
+  ///
+  /// The returned IS selects the evolving subset of the full cell space C,
+  /// excluding inner/outer X-boundary cells and forward/backward parallel
+  /// boundary virtual cells.  It is the correct IS to pass to
+  /// MatCreateSubMatrix when restricting a PetscCellOperator to the degrees
+  /// of freedom that the SNES solver actually integrates.
+  ///
+  /// The caller owns the returned IS and must call ISDestroy when finished.
+  ///
+  /// @returns A PETSc IS of local size equal to the number of locally owned
+  ///          evolving cells, holding global PETSc row indices.
+  IS makeEvolvingIS() const;
+
+  /// @brief Extract the evolving-cell submatrix from a cell-to-cell operator.
+  ///
+  /// Restricts @p op to the rows and columns that correspond to evolving
+  /// interior cells, discarding any rows or columns that belong to inner/outer
+  /// X-boundary cells or forward/backward parallel boundary virtual cells.
+  ///
+  /// The returned Mat is an independent copy (MAT_INITIAL_MATRIX): subsequent
+  /// modifications to @p op do not affect it.  The caller owns the returned
+  /// Mat and must call MatDestroy when finished.
+  ///
+  /// @param op A cell-to-cell operator whose row and column space is the full
+  ///           cell space C managed by this mapping.
+  /// @returns  A square Mat of global size n_evolving × n_evolving.
+  Mat extractEvolvingSubmatrix(const PetscOperator<CellSpaceTag, CellSpaceTag>& op) const;
+
 private:
   Field3D cell_number; ///< Stored cell numbers for interior/X-boundary cells.
   Field3D
@@ -865,8 +899,8 @@ public:
   ///
   /// Naming convention: the subscript on Grad/Div/Restrict refers to which side of
   /// the cell face the operator acts on, not the leg space it lives in.
-  ///   - @c Grad_plus / @c Div_plus: forward-side half-step operators (in L-).
-  ///   - @c Grad_minus / @c Div_minus: backward-side half-step operators (in L+).
+  ///   - @c Grad_plus / @c Div_plus: forward-side half-step operators (C to L+).
+  ///   - @c Grad_minus / @c Div_minus: backward-side half-step operators (C to L-).
   ///   - @c Restrict_minus = I+^T * W+: weighted restriction from L+ back to C,
   ///     paired with the minus-side gradient.
   ///   - @c Restrict_plus  = I-^T * W-: weighted restriction from L- back to C,
@@ -914,7 +948,7 @@ public:
     ///
     /// Evaluates the half-step fluxes separately on each side, multiplies by the
     /// interpolated diffusion coefficient K, and averages the two divergence
-    /// contributions. This is the primary user-facing SOM operator.
+    /// contributions.
     ///
     /// @param K Diffusion coefficient field; must have parallel slices allocated.
     /// @param f Field to differentiate; must have parallel slices allocated.

src/mesh/petsc_operators.cxx

@@ -127,6 +127,33 @@ PetscCellMapping::PetscCellMapping(const Field3D& cell_number,
                    local_indices);
 }
 
+IS PetscCellMapping::makeEvolvingIS() const {
+  // Collect global PETSc indices in mapOwnedInteriorCells order.
+  // Reserve the known count up front to avoid reallocation.
+  std::vector<PetscInt> indices;
+  indices.reserve(static_cast<std::size_t>(evolving_region.size()));
+
+  mapOwnedInteriorCells(
+      [&](PetscInt row, const Ind3D& /*i*/, int /*stored*/) { indices.push_back(row); });
+
+  IS is;
+  BOUT_DO_PETSC(ISCreateGeneral(BoutComm::get(), static_cast<PetscInt>(indices.size()),
+                                indices.data(), PETSC_COPY_VALUES, &is));
+  return is;
+}
+
+Mat PetscCellMapping::extractEvolvingSubmatrix(
+    const PetscOperator<CellSpaceTag, CellSpaceTag>& op) const {
+  IS is = makeEvolvingIS();
+
+  Mat sub;
+  BOUT_DO_PETSC(MatCreateSubMatrix(op.raw(), is, is, MAT_INITIAL_MATRIX, &sub));
+
+  BOUT_DO_PETSC(ISDestroy(&is));
+
+  return sub;
+}
+
 PetscLegMapping::PetscLegMapping(int total_legs, std::vector<int> local_leg_indices) {
   std::sort(local_leg_indices.begin(), local_leg_indices.end());
   local_leg_indices.erase(std::unique(local_leg_indices.begin(), local_leg_indices.end()),

tests/unit/CMakeLists.txt

@@ -98,6 +98,8 @@ set(serial_tests_source
     ./mesh/test_mesh.cxx
     ./mesh/test_paralleltransform.cxx
     ./mesh/test_petsc_operators.cxx
+    ./mesh/test_petsc_operators_make_evolving_is.cxx
+    ./mesh/test_petsc_operators_extract_evolving_submatrix.cxx
     ./solver/test_fakesolver.cxx
     ./solver/test_fakesolver.hxx
     ./solver/test_solver.cxx