WIP: Calculate Jacobian coloring using PetscOperator matrices

CMakeLists.txt

@@ -179,6 +179,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
@@ -313,6 +314,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/mesh/tokamak_coordinates.cxx
     ./src/physics/gyro_average.cxx

include/bout/petsc_jacobian.hxx

@@ -0,0 +1,54 @@
+#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.
+///
+/// Limitation: this helper assumes that @p Jfd uses a uniform per-cell
+/// interleaving of the form <tt>global_row = cell * nvars + var</tt> over the
+/// evolving cells only. It is therefore not valid for solver layouts that mix
+/// Field2D and Field3D variables, include evolving boundary variables, or use
+/// any other non-uniform cell-to-row mapping. Callers must check that the
+/// solver Jacobian ordering matches this assumption before using this helper.
+///
+/// @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.
+///
+/// The same layout restriction applies as for the block-specific overload:
+/// @p Jfd must use the uniform <tt>cell * nvars + var</tt> ordering over
+/// evolving cells only.
+///
+/// @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

@@ -290,6 +290,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;
+
   friend PetscOperator<CellSpaceTag, CellSpaceTag>
   makeNeumannOperator(const PetscCellMappingPtr& mapping, BoundaryDirection direction);
 
@@ -1050,6 +1081,6 @@ PetscCellOperator makeNeumannOperator(const PetscCellMappingPtr& mapping,
 
 #warning PETSc not available. No PetscOperators.
 
-#endif
+#endif // BOUT_HAS_PETSC
 
 #endif // BOUT_PETSC_OPERATORS

src/mesh/mesh.cxx

@@ -189,10 +189,16 @@ int Mesh::get(bool& bval, const std::string& name, bool def) {
 }
 
 int Mesh::get(Array<int>& var, const std::string& name) {
+  if (source == nullptr) {
+    return 1;
+  }
   return source->get(var, name) ? 0 : 1;
 }
 
 int Mesh::get(Array<BoutReal>& var, const std::string& name) {
+  if (source == nullptr) {
+    return 1;
+  }
   return source->get(var, name) ? 0 : 1;
 }
 

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

@@ -130,6 +130,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()),
@@ -332,14 +359,14 @@ PetscOperators::Parallel PetscOperators::getParallel() const {
   auto* coords = mesh->getCoordinates();
 
   // Parallel spacing in cell space
-  Field3D dl = coords->dy * sqrt(coords->g_22);
+  Field3D dl = Coordinates::FieldMetric{coords->dy * sqrt(coords->g_22)};
   dl.splitParallelSlices();
   dl.yup() = 0.0;
   dl.ydown() = 0.0;
   dl.applyParallelBoundary("parallel_neumann_o1");
 
   // Cell volume
-  Field3D dV = coords->J * coords->dx * coords->dy * coords->dz;
+  Field3D dV = Coordinates::FieldMetric{coords->J * coords->dx * coords->dy * coords->dz};
   dV.splitParallelSlices();
   dV.yup() = 0.0;
   dV.ydown() = 0.0;

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-petsc-operators)
 add_subdirectory(test-restart-io)
 add_subdirectory(test-restarting)

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