WIP: Calculate Jacobian coloring using PetscOperator matrices

CMakeLists.txt

@@ -178,6 +178,7 @@ set(BOUT_SOURCES
     ./include/bout/paralleltransform.hxx
     ./include/bout/petsc_interface.hxx
     ./include/bout/petsc_operators.hxx
+    ./include/bout/petsc_jacobian.hxx
     ./include/bout/petsclib.hxx
     ./include/bout/physicsmodel.hxx
     ./include/bout/rajalib.hxx
@@ -310,6 +311,7 @@ set(BOUT_SOURCES
     ./src/mesh/parallel_boundary_op.cxx
     ./src/mesh/parallel_boundary_region.cxx
     ./src/mesh/petsc_operators.cxx
+    ./src/mesh/petsc_jacobian.cxx
     ./src/mesh/surfaceiter.cxx
     ./src/physics/gyro_average.cxx
     ./src/physics/physicsmodel.cxx

include/bout/petsc_jacobian.hxx

@@ -0,0 +1,43 @@
+#pragma once
+
+#ifndef BOUT_PETSC_JACOBIAN_H
+#define BOUT_PETSC_JACOBIAN_H
+
+#include "bout/build_defines.hxx"
+
+#if BOUT_HAS_PETSC
+
+#include <petscmat.h>
+
+/// Insert the nonzero pattern of @p sub into the variable block
+/// (@p out_var, @p in_var) of the Jacobian @p Jfd.
+///
+/// @p Jfd is a square matrix of size (n_evolving * nvars) where nvars is
+/// inferred as Jfd_global_size / sub_global_size.  Each nonzero (r, c) in
+/// @p sub produces an entry at (r * nvars + out_var, c * nvars + in_var)
+/// in @p Jfd.
+///
+/// @p Jfd must already be preallocated.  Entries are inserted with
+/// INSERT_VALUES; MatAssemblyBegin/End must be called by the caller after
+/// all insertions are complete.
+///
+/// @param Jfd     The Jacobian matrix to populate. Must be preallocated.
+/// @param sub     Evolving-cell submatrix providing the nonzero pattern.
+/// @param out_var Row variable index in [0, nvars).
+/// @param in_var  Column variable index in [0, nvars).
+void addOperatorSparsity(Mat Jfd, Mat sub, int out_var, int in_var);
+
+/// @brief Insert the nonzero pattern of @p sub into every variable block of
+///        @p Jfd.
+///
+/// Equivalent to calling addOperatorSparsity(Jfd, sub, out_var, in_var) for
+/// every combination of @p out_var and @p in_var in [0, nvars), where
+/// @c nvars is inferred as @c Jfd_global / @c sub_global.
+///
+/// @param Jfd     The Jacobian matrix to populate. Must be preallocated.
+/// @param sub     Evolving-cell submatrix providing the nonzero pattern.
+void addOperatorSparsity(Mat Jfd, Mat sub);
+
+#endif // BOUT_HAS_PETSC
+
+#endif // BOUT_PETSC_JACOBIAN_H

include/bout/petsc_operators.hxx

@@ -1,3 +1,8 @@
+#pragma once
+
+#ifndef BOUT_PETSC_OPERATORS_H
+#define BOUT_PETSC_OPERATORS_H
+
 #include "bout/build_defines.hxx"
 
 #if BOUT_HAS_PETSC
@@ -8,7 +13,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 +49,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 +283,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 +904,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 +953,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.
@@ -1021,4 +1060,6 @@ private:
 
 #warning PETSc not available. No PetscOperators.
 
-#endif
+#endif // BOUT_HAS_PETSC
+
+#endif //BOUT_PETSC_OPERATORS_H

src/mesh/petsc_jacobian.cxx

@@ -0,0 +1,65 @@
+#include "bout/build_defines.hxx"
+
+#if BOUT_HAS_PETSC
+
+#include "bout/assert.hxx"
+#include "bout/petsc_jacobian.hxx"
+#include "bout/petsclib.hxx"
+
+#include <petscmat.h>
+
+void addOperatorSparsity(Mat Jfd, Mat sub, int out_var, int in_var) {
+  // Infer nvars from global sizes
+  PetscInt jfd_global{0}, sub_global{0};
+  BOUT_DO_PETSC(MatGetSize(Jfd, &jfd_global, nullptr));
+  BOUT_DO_PETSC(MatGetSize(sub, &sub_global, nullptr));
+
+  ASSERT1(sub_global > 0);
+  ASSERT1(jfd_global % sub_global == 0);
+
+  const PetscInt nvars = jfd_global / sub_global;
+
+  ASSERT1(out_var >= 0 && out_var < static_cast<int>(nvars));
+  ASSERT1(in_var >= 0 && in_var < static_cast<int>(nvars));
+
+  // Iterate over locally owned rows of sub and insert into Jfd
+  PetscInt rstart{0}, rend{0};
+  BOUT_DO_PETSC(MatGetOwnershipRange(sub, &rstart, &rend));
+
+  const PetscScalar one = 1.0;
+
+  for (PetscInt sub_row = rstart; sub_row < rend; ++sub_row) {
+    PetscInt ncols{0};
+    const PetscInt* sub_cols{nullptr};
+    const PetscScalar* vals{nullptr};
+    BOUT_DO_PETSC(MatGetRow(sub, sub_row, &ncols, &sub_cols, &vals));
+
+    const PetscInt jfd_row = (sub_row * nvars) + out_var;
+
+    for (PetscInt k = 0; k < ncols; ++k) {
+      const PetscInt jfd_col = (sub_cols[k] * nvars) + in_var;
+      BOUT_DO_PETSC(MatSetValues(Jfd, 1, &jfd_row, 1, &jfd_col, &one, INSERT_VALUES));
+    }
+
+    BOUT_DO_PETSC(MatRestoreRow(sub, sub_row, &ncols, &sub_cols, &vals));
+  }
+}
+
+void addOperatorSparsity(Mat Jfd, Mat sub) {
+  PetscInt jfd_global{0}, sub_global{0};
+  MatGetSize(Jfd, &jfd_global, nullptr);
+  MatGetSize(sub, &sub_global, nullptr);
+
+  ASSERT1(sub_global > 0);
+  ASSERT1(jfd_global % sub_global == 0);
+
+  const int nvars = static_cast<int>(jfd_global / sub_global);
+
+  for (int out_var = 0; out_var < nvars; ++out_var) {
+    for (int in_var = 0; in_var < nvars; ++in_var) {
+      addOperatorSparsity(Jfd, sub, out_var, in_var);
+    }
+  }
+}
+
+#endif // BOUT_HAS_PETSC

src/mesh/petsc_operators.cxx

@@ -6,7 +6,6 @@
 #include "bout/bout_types.hxx"
 #include "bout/boutexception.hxx"
 #include "bout/field3d.hxx"
-#include "bout/output.hxx"
 #include "bout/output_bout_types.hxx"
 #include "bout/petsc_operators.hxx"
 #include "bout/region.hxx"
@@ -127,6 +126,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/integrated/CMakeLists.txt

@@ -34,6 +34,7 @@ add_subdirectory(test-naulin-laplace)
 add_subdirectory(test-options-netcdf)
 add_subdirectory(test-petsc_laplace)
 add_subdirectory(test-petsc_laplace_MAST-grid)
+add_subdirectory(test-petsc-ordering)
 add_subdirectory(test-restart-io)
 add_subdirectory(test-restarting)
 add_subdirectory(test-slepc-solver)

tests/integrated/test-petsc-ordering/CMakeLists.txt

@@ -0,0 +1,7 @@
+bout_add_integrated_test(
+  test-petsc-ordering
+  SOURCES test_petsc_ordering.cxx
+  REQUIRES BOUT_HAS_PETSC
+  USE_RUNTEST USE_DATA_BOUT_INP
+  PROCESSORS 4
+)

tests/integrated/test-petsc-ordering/data/BOUT.inp

@@ -0,0 +1,14 @@
+# Input file for test_petsc_ordering integrated test.
+# Uses a small grid large enough to exercise MPI decomposition on 4 ranks
+# (nx=10 gives 2 interior x-points per rank when NXPE=4; ny and nz similarly).
+
+[mesh]
+nx = 10     # 8 interior + 2 guard (mxg=1 each side)
+ny = 8
+nz = 4
+
+ixseps1 = nx
+ixseps2 = nx
+
+[output]
+enabled = false   # Suppress data file output; we only need stdout