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);
 
@@ -678,6 +709,16 @@ public:
   /// low-level callers that need direct PETSc access.
   Mat raw() const { return *this->mat_operator; }
 
+  /// @brief Return the shared mapping for the operator's output (row) space.
+  const std::shared_ptr<const PetscIndexMapping>& getOutMapping() const {
+    return out_mapping;
+  }
+
+  /// @brief Return the shared mapping for the operator's input (column) space.
+  const std::shared_ptr<const PetscIndexMapping>& getInMapping() const {
+    return in_mapping;
+  }
+
   /// @brief Construct a diagonal operator from a vector of diagonal entries.
   ///
   /// Only available when @p OutSpaceTag == @p InSpaceTag (square, same-space
@@ -1077,6 +1118,6 @@ PetscCellOperator makeNeumannOperator(const PetscCellMappingPtr& mapping,
 
 #warning PETSc not available. No PetscOperators.
 
-#endif
+#endif // BOUT_HAS_PETSC
 
 #endif // BOUT_PETSC_OPERATORS

include/bout/solver.hxx

@@ -1,20 +1,10 @@
 /**************************************************************************
  * Base class for all solvers. Specifies required interface functions
  *
- * Changelog:
- *
- * 2009-08 Ben Dudson, Sean Farley
- *    * Major overhaul, and changed API. Trying to make consistent
- *      interface to PETSc and SUNDIALS solvers
- *
- * 2013-08 Ben Dudson
- *    * Added OO-style API, to allow multiple physics models to coexist
- *      For now both APIs are supported
- *
  **************************************************************************
- * Copyright 2010 B.D.Dudson, S.Farley, M.V.Umansky, X.Q.Xu
+ * Copyright 2010 - 2026 BOUT++ contributors
  *
- * Contact: Ben Dudson, bd512@york.ac.uk
+ * Contact: Ben Dudson, dudson2@llnl.gov
  *
  * This file is part of BOUT++.
  *
@@ -45,7 +35,6 @@
 #include "bout/monitor.hxx"
 #include "bout/options.hxx"
 #include "bout/region.hxx"
-#include "bout/unused.hxx"
 
 #include <cstdint>
 #include <iterator>
@@ -72,6 +61,10 @@ using TimestepMonitorFunc = int (*)(Solver* solver, BoutReal simtime, BoutReal l
 #include "bout/field2d.hxx"
 #include "bout/field3d.hxx"
 #include "bout/generic_factory.hxx"
+#if BOUT_HAS_PETSC
+#include "bout/petsc_interface.hxx"
+#include "bout/petsc_operators.hxx"
+#endif
 #include "bout/vector2d.hxx"
 #include "bout/vector3d.hxx"
 
@@ -81,6 +74,8 @@ using TimestepMonitorFunc = int (*)(Solver* solver, BoutReal simtime, BoutReal l
 
 #include <list>
 #include <string>
+#include <string_view>
+#include <vector>
 
 using SolverType = std::string;
 constexpr auto SOLVERCVODE = "cvode";
@@ -210,6 +205,30 @@ using RegisterUnavailableSolver = SolverFactory::RegisterUnavailableInFactory;
  */
 class Solver {
 public:
+  /// Variable reference handle
+  class VarRef {
+  public:
+    static constexpr int AllValue = -1;
+    static constexpr int InvalidValue = -2;
+
+    constexpr VarRef() = default;
+
+    static constexpr VarRef All() { return VarRef(AllValue); }
+    static constexpr VarRef Invalid() { return VarRef(InvalidValue); }
+
+    constexpr bool isAll() const { return value == AllValue; }
+    constexpr bool isInvalid() const { return value == InvalidValue; }
+    constexpr bool isConcrete() const { return value >= 0; }
+    constexpr bool isValid() const { return isAll() || isConcrete(); }
+    constexpr int index() const { return value; }
+
+  private:
+    explicit constexpr VarRef(int value) : value(value) {}
+
+    int value{InvalidValue};
+    friend class Solver;
+  };
+
   Solver(Options* opts = nullptr);
   virtual ~Solver() = default;
 
@@ -265,6 +284,21 @@ public:
   virtual void add(Vector3D& v, const std::string& name,
                    const std::string& description = "");
 
+  /// Get the solver-variable reference for a scalar field/component by name.
+  VarRef getVarRef(std::string_view name) const;
+
+#if BOUT_HAS_PETSC
+  /// Register a Jacobian sparsity contribution for all variable blocks.
+  bool addJacobianPattern(const PetscCellOperator& op);
+
+  /// Register a Jacobian sparsity contribution for one or more variable blocks.
+  ///
+  /// Either variable reference may be VarRef::All(), which expands when the
+  /// Jacobian matrix is created during solver initialisation.
+  virtual bool addJacobianPattern(const PetscCellOperator& op, VarRef out_var,
+                                  VarRef in_var);
+#endif
+
   /// Returns true if constraints available
   virtual bool constraints() { return has_constraints; }
 
@@ -394,7 +428,7 @@ protected:
     bool covariant{false};               /// For vectors
     bool evolve_bndry{false};            /// Are the boundary regions being evolved?
     std::string name;                    /// Name of the variable
-    std::string description{""};         /// Description of what the variable is
+    std::string description;             /// Description of what the variable is
   };
 
   /// A structure for iterating over fields
@@ -688,6 +722,26 @@ private:
   /// Physics model being evolved
   PhysicsModel* model{nullptr};
 
+protected:
+#if BOUT_HAS_PETSC
+  struct DeferredJacobianPattern {
+    bout::petsc::UniqueMat submatrix{new Mat{nullptr}};
+    VarRef out_var;
+    VarRef in_var;
+  };
+
+  /// Queue a Jacobian-pattern contribution for PETSc-preconditioner-based solvers.
+  bool queueJacobianPattern(const PetscCellOperator& op, VarRef out_var, VarRef in_var);
+
+  /// Insert any queued Jacobian-pattern contributions into an existing Jacobian matrix.
+  void applyQueuedJacobianPatterns(Mat Jfd) const;
+
+  /// Check whether the current solver variable layout is compatible with
+  /// addOperatorSparsity().
+  bool canApplyQueuedJacobianPatterns() const;
+#endif
+
+private:
   /// Should non-split physics models be treated as diffusive?
   bool is_nonsplit_model_diffusive{true};
 
@@ -704,6 +758,10 @@ private:
   /// List of timestep monitor functions
   std::list<TimestepMonitorFunc> timestep_monitors;
 
+#if BOUT_HAS_PETSC
+  std::vector<DeferredJacobianPattern> deferred_jacobian_patterns;
+#endif
+
   /// Should be run before user RHS is called
   void pre_rhs(BoutReal t);
   /// Should be run after user RHS is called

manual/sphinx/index.rst

@@ -67,6 +67,7 @@ The documentation is divided into the following sections:
    :name: bout-interfaces
 
    user_docs/time_integration
+   user_docs/preconditioning
    user_docs/parallel-transforms
    user_docs/differential_operators
    user_docs/parallel_operators
@@ -92,7 +93,6 @@ The documentation is divided into the following sections:
    :name: coordinates
 
    user_docs/coordinates
-   user_docs/preconditioning
    user_docs/BOUT_Gradperp_op
 
 .. toctree::

manual/sphinx/user_docs/parallel_operators.rst

@@ -292,11 +292,93 @@ Current limitations include:
 - PETSc and 3D metrics are required
 - The grid metadata must be generated ahead of time
 
+Using operator stencils in solver Jacobians
+-------------------------------------------
+
+The ``PetscCellOperator`` objects can also be used to augment the Jacobian
+pattern of PETSc-backed implicit solvers. This is useful when the default
+solver stencil is too small for the operators used in the model, but the
+coupling still has a sparse cell-to-cell structure.
+
+At a high level:
+
+- Build or reuse a ``PetscCellOperator``
+- Get solver variable references with ``solver->getVarRef(...)``
+- Register the operator with ``solver->addJacobianPattern(...)`` before
+  ``solver->init()``
+
+For example:
+
+.. code-block:: C++
+
+   #if BOUT_HAS_PETSC
+   PetscOperators ops(mesh);
+   auto parallel = ops.getParallel();
+
+   auto n = solver->getVarRef("n");
+   auto T = solver->getVarRef("T");
+
+   // Apply the operator stencil to every Jacobian block
+   solver->addJacobianPattern(parallel.Div_par_Grad_par);
+
+   // Apply the same cell stencil only to dF_n / dT
+   solver->addJacobianPattern(parallel.Grad_par, n, T);
+   #endif
+
+The one-argument overload is equivalent to:
+
+.. code-block:: C++
+
+   solver->addJacobianPattern(op, Solver::VarRef::All(), Solver::VarRef::All());
+
+The ``VarRef::All()`` sentinel expands when the solver Jacobian is created:
+
+- ``(All, All)`` inserts the stencil into every variable block
+- ``(out, All)`` inserts it into one output-variable row block against all inputs
+- ``(All, in)`` inserts it into all output-variable row blocks against one input
+- ``(out, in)`` inserts it into a single block
+
+This interface is only available when BOUT++ is built with PETSc. Even then,
+``addJacobianPattern(...)`` may return ``false`` if the chosen solver does not
+use the PETSc preconditioner/Jacobian path. See :ref:`sec-time-integration` and
+:ref:`sec-preconditioning` for the solver-side behavior.
+
+
+How solver registration works
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+``PetscOperator`` tracks its row and column mappings, and exposes them through
+the operator API. For Jacobian-pattern registration, BOUT++ uses the operator's
+cell mapping to restrict the full cell-space matrix to the evolving subset of
+cells that actually appear in the solver state vector.
+
+The resulting evolving-cell submatrix is then inserted into one or more
+Jacobian variable blocks using the helper documented in
+``bout/petsc_jacobian.hxx``. This means the operator contributes only sparsity
+information here; the Jacobian entries are still computed later by the solver's
+finite-difference machinery.
+
+Current limitations
+~~~~~~~~~~~~~~~~~~~
+
+This solver-Jacobian path currently assumes the Jacobian uses a uniform
+per-cell interleaving of the evolving variables. In practice that means it is
+currently intended for the same cases documented in
+``bout/petsc_jacobian.hxx``:
+
+- Evolving ``Field3D`` variables only
+- No evolving boundary cells in the solver state
+- No solver layouts with mixed or non-uniform cell-to-row mappings
+
+If those assumptions do not hold, use the standard solver coloring stencil or a
+different Jacobian/preconditioning strategy instead.
 
 See also
 --------
 
 - :ref:`sec-diffops`
 - :ref:`sec-fci`
 - :ref:`sec-parallel-bc-fci`
+- :ref:`sec-time-integration`
+- :ref:`sec-preconditioning`
 - :ref:`sec-petsc`