From d43cea2ee40a654898a4b4e4801aeb4934bb800d Mon Sep 17 00:00:00 2001
From: Seimon Powell <seimon.powell@ukaea.uk>
Date: Thu, 24 Apr 2025 14:54:20 +0100
Subject: [PATCH 1/9] Enable alternative stencil shapes in jocobian colouring

4 new options added

solver:stencil:<cross,square,taxi> = N where <option> is either cross, square, or taxi:
cross - all points (|x| == 0 and |y| <= N) or (|x| <= N and |y| == 0)
square - all points |x| < N or |y| < N
taxi - all points |x| + |y| < N
can be combined will just create a union of every stencil
New default is solver:stencil:taxi = 2
Also option to force colouring matrix to be symmetric
solver:force_symmetric_coloring=true
---
 src/solver/impls/snes/snes.cxx | 905 ++++++++++-----------------------
 1 file changed, 270 insertions(+), 635 deletions(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index e382bbd3f8..aec986511f 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -1,4 +1,4 @@
-#include "bout/build_defines.hxx"
+#include "bout/build_config.hxx"
 
 #if BOUT_HAS_PETSC
 
@@ -9,15 +9,52 @@
 #include <bout/msg_stack.hxx>
 #include <bout/utils.hxx>
 
-#include <array>
 #include <cmath>
 #include <vector>
 
 #include <bout/output.hxx>
 
-#include "petscmat.h"
 #include "petscsnes.h"
 
+class ColoringStencil {
+private:
+  bool static in_compact(int const i, int const j, int const n_compact) {
+    return std::abs(i) <= n_compact && std::abs(j) <= n_compact;
+  }
+  bool static in_cross(int const i, int const j, int const n_cross) {
+    if (i == 0) {
+      return std::abs(j) <= n_cross;
+    }
+    if (j == 0) {
+      return std::abs(i) <= n_cross;
+    }
+    return false;
+  }
+  bool static in_taxi(int const i, int const j, int const n_taxi) {
+    return std::abs(i) + std::abs(j) <= n_taxi;
+  }
+
+public:
+  auto static get_offsets(int n_compact, int n_taxi, int n_cross) {
+    ASSERT2(n_compact >= 0 && n_cross >= 0 && n_taxi >= 0
+            && n_compact + n_cross + n_taxi > 0);
+    auto inside = [&](int i, int j) {
+      return in_compact(i, j, n_compact) || in_taxi(i, j, n_taxi)
+             || in_cross(i, j, n_cross);
+    };
+    std::vector<std::pair<int, int>> xyoffset;
+    auto loop_bound = std::max({n_compact, n_taxi, n_cross});
+    for (int i = -loop_bound; i <= loop_bound; ++i) {
+      for (int j = -loop_bound; j <= loop_bound; ++j) {
+        if (inside(i, j)) {
+          xyoffset.emplace_back(i, j);
+        }
+      }
+    }
+    return xyoffset;
+  }
+};
+
 /*
  * PETSc callback function, which evaluates the nonlinear
  * function to be solved by SNES.
@@ -120,31 +157,12 @@ SNESSolver::SNESSolver(Options* opts)
       matrix_free((*options)["matrix_free"]
                       .doc("Use matrix free Jacobian?")
                       .withDefault<bool>(false)),
-      matrix_free_operator((*options)["matrix_free_operator"]
-                               .doc("Use matrix free Jacobian-vector operator?")
-                               .withDefault<bool>(true)),
       lag_jacobian((*options)["lag_jacobian"]
                        .doc("Re-use the Jacobian this number of SNES iterations")
                        .withDefault(50)),
       use_coloring((*options)["use_coloring"]
                        .doc("Use matrix coloring to calculate Jacobian?")
-                       .withDefault<bool>(true)),
-      jacobian_recalculated(false),
-      prune_jacobian((*options)["prune_jacobian"]
-                         .doc("Remove small elements in the Jacobian?")
-                         .withDefault<bool>(false)),
-      prune_abstol((*options)["prune_abstol"]
-                       .doc("Prune values with absolute values smaller than this")
-                       .withDefault<BoutReal>(1e-16)),
-      prune_fraction((*options)["prune_fraction"]
-                         .doc("Prune if fraction of small elements is larger than this")
-                         .withDefault<BoutReal>(0.2)),
-      scale_rhs((*options)["scale_rhs"]
-                    .doc("Scale time derivatives (Jacobian row scaling)?")
-                    .withDefault<bool>(false)),
-      scale_vars((*options)["scale_vars"]
-                     .doc("Scale variables (Jacobian column scaling)?")
-                     .withDefault<bool>(false)) {}
+                       .withDefault<bool>(true)) {}
 
 int SNESSolver::init() {
 
@@ -184,38 +202,12 @@ int SNESSolver::init() {
 
   if (equation_form == BoutSnesEquationForm::rearranged_backward_euler) {
     // Need an intermediate vector for rearranged Backward Euler
-    ierr = VecDuplicate(snes_x, &delta_x);
-    CHKERRQ(ierr);
+    VecDuplicate(snes_x, &delta_x);
   }
 
   if (predictor) {
     // Storage for previous solution
-    ierr = VecDuplicate(snes_x, &x1);
-    CHKERRQ(ierr);
-  }
-
-  if (scale_rhs) {
-    // Storage for rhs factors, one per evolving variable
-    ierr = VecDuplicate(snes_x, &rhs_scaling_factors);
-    CHKERRQ(ierr);
-    // Set all factors to 1 to start with
-    ierr = VecSet(rhs_scaling_factors, 1.0);
-    CHKERRQ(ierr);
-    // Array to store inverse Jacobian row norms
-    ierr = VecDuplicate(snes_x, &jac_row_inv_norms);
-    CHKERRQ(ierr);
-  }
-
-  if (scale_vars) {
-    // Storage for var factors, one per evolving variable
-    ierr = VecDuplicate(snes_x, &var_scaling_factors);
-    CHKERRQ(ierr);
-    // Set all factors to 1 to start with
-    ierr = VecSet(var_scaling_factors, 1.0);
-    CHKERRQ(ierr);
-    // Storage for scaled 'x' state vectors
-    ierr = VecDuplicate(snes_x, &scaled_x);
-    CHKERRQ(ierr);
+    VecDuplicate(snes_x, &x1);
   }
 
   // Nonlinear solver interface (SNES)
@@ -235,7 +227,7 @@ int SNESSolver::init() {
   }
 
   // Set up the Jacobian
-  if (matrix_free or matrix_free_operator) {
+  if (matrix_free) {
     /*
       PETSc SNES matrix free Jacobian, using a different
       operator for differencing.
@@ -251,17 +243,12 @@ int SNESSolver::init() {
     // Set a function to be called for differencing
     // This can be a linearised form of the SNES function
     MatMFFDSetFunction(Jmf, FormFunctionForDifferencing, this);
-  }
 
-  if (matrix_free) {
-    // Use matrix free for both operator and preconditioner
     // Calculate Jacobian matrix free using FormFunctionForDifferencing
     SNESSetJacobian(snes, Jmf, Jmf, MatMFFDComputeJacobian, this);
 
   } else {
-    // Calculate the Jacobian using finite differences.
-    // The finite difference Jacobian (Jfd) may be used for both operator
-    // and preconditioner or, if matrix_free_operator, in only the preconditioner.
+    // Calculate the Jacobian using finite differences
     if (use_coloring) {
       // Use matrix coloring
       // This greatly reduces the number of times the rhs() function needs
@@ -279,263 +266,203 @@ int SNESSolver::init() {
 
       output_progress.write("Setting Jacobian matrix sizes\n");
 
+      int localN = getLocalN(); // Number of rows on this processor
       int n2d = f2d.size();
       int n3d = f3d.size();
 
-      // Set size of Matrix on each processor to nlocal x nlocal
-      MatCreate(BoutComm::get(), &Jfd);
-      MatSetSizes(Jfd, nlocal, nlocal, PETSC_DETERMINE, PETSC_DETERMINE);
-      MatSetFromOptions(Jfd);
-
-      std::vector<PetscInt> d_nnz(nlocal);
-      std::vector<PetscInt> o_nnz(nlocal);
-
-      // Set values for most points
-      const int ncells_x = (mesh->LocalNx > 1) ? 2 : 0;
-      const int ncells_y = (mesh->LocalNy > 1) ? 2 : 0;
-      const int ncells_z = (mesh->LocalNz > 1) ? 2 : 0;
-
-      const auto star_pattern = (1 + ncells_x + ncells_y) * (n3d + n2d) + ncells_z * n3d;
-
-      // Offsets. Start with the central cell
-      std::vector<std::pair<int, int>> xyoffsets{{0, 0}};
-      if (ncells_x != 0) {
-        // Stencil includes points in X
-        xyoffsets.push_back({-1, 0});
-        xyoffsets.push_back({1, 0});
-      }
-      if (ncells_y != 0) {
-        // Stencil includes points in Y
-        xyoffsets.push_back({0, -1});
-        xyoffsets.push_back({0, 1});
+      // Set size of Matrix on each processor to localN x localN
+      MatCreate(BoutComm::get(), &Jmf);
+      MatSetSizes(Jmf, localN, localN, PETSC_DETERMINE, PETSC_DETERMINE);
+      MatSetFromOptions(Jmf);
+      // Determine which row/columns of the matrix are locally owned
+      int Istart, Iend;
+      MatGetOwnershipRange(Jmf, &Istart, &Iend);
+      // Convert local into global indices
+      // Note: Not in the boundary cells, to keep -1 values
+      for (const auto& i : mesh->getRegion3D("RGN_NOBNDRY")) {
+        index[i] += Istart;
       }
+      // Now communicate to fill guard cells
+      mesh->communicate(index);
 
-      output_info.write("Star pattern: {} non-zero entries\n", star_pattern);
-      for (int i = 0; i < nlocal; i++) {
-        // Non-zero elements on this processor
-        d_nnz[i] = star_pattern;
-        // Non-zero elements on neighboring processor
-        o_nnz[i] = 0;
+      // Non-zero elements on this processor
+      std::vector<PetscInt> d_nnz;
+      std::vector<PetscInt> o_nnz;
+      auto n_square = (*options)["stencil:square"]
+                          .doc("Extent of stencil (square)")
+                          .withDefault<int>(0);
+      auto n_taxi = (*options)["stencil:taxi"]
+                        .doc("Extent of stencil (taxi-cab norm)")
+                        .withDefault<int>(0);
+      auto n_cross = (*options)["stencil:cross"]
+                         .doc("Extent of stencil (cross)")
+                         .withDefault<int>(0);
+      //Set n_taxi 2 if nothing else is set
+      //TODO: Probably a better way to do this
+      if (n_square == 0 && n_taxi == 0 && n_cross == 0) {
+        output_info.write("Setting beuler:stencil:taxi = 2\n");
+        n_taxi = 2;
       }
 
-      // X boundaries
-      if (ncells_x != 0) {
-        if (mesh->firstX()) {
-          // Lower X boundary
-          for (int y = mesh->ystart; y <= mesh->yend; y++) {
-            for (int z = 0; z < mesh->LocalNz; z++) {
-              const int localIndex = ROUND(index(mesh->xstart, y, z));
-              ASSERT2((localIndex >= 0) && (localIndex < nlocal));
-              const int num_fields = (z == 0) ? n2d + n3d : n3d;
-              for (int i = 0; i < num_fields; i++) {
-                d_nnz[localIndex + i] -= (n3d + n2d);
-              }
-            }
-          }
-        } else {
-          // On another processor
-          for (int y = mesh->ystart; y <= mesh->yend; y++) {
-            for (int z = 0; z < mesh->LocalNz; z++) {
-              const int localIndex = ROUND(index(mesh->xstart, y, z));
-              ASSERT2((localIndex >= 0) && (localIndex < nlocal));
-              const int num_fields = (z == 0) ? n2d + n3d : n3d;
-              for (int i = 0; i < num_fields; i++) {
-                d_nnz[localIndex + i] -= (n3d + n2d);
-                o_nnz[localIndex + i] += (n3d + n2d);
-              }
-            }
-          }
-        }
-        if (mesh->lastX()) {
-          // Upper X boundary
-          for (int y = mesh->ystart; y <= mesh->yend; y++) {
-            for (int z = 0; z < mesh->LocalNz; z++) {
-              const int localIndex = ROUND(index(mesh->xend, y, z));
-              ASSERT2((localIndex >= 0) && (localIndex < nlocal));
-              const int num_fields = (z == 0) ? n2d + n3d : n3d;
-              for (int i = 0; i < num_fields; i++) {
-                d_nnz[localIndex + i] -= (n3d + n2d);
-              }
-            }
-          }
-        } else {
-          // On another processor
+      auto const xyoffsets = ColoringStencil::get_offsets(n_square, n_taxi, n_cross);
+      {
+        //This is nasty but can't think of a better and robust way to
+        //count the non-zeros for some arbitery stencil
+        std::vector<std::set<int>> d_nnz_map2d(localN);
+        std::vector<std::set<int>> o_nnz_map2d(localN);
+        std::vector<std::set<int>> d_nnz_map3d(localN);
+        std::vector<std::set<int>> o_nnz_map3d(localN);
+        //Loop over every element in 2D to count the *unique* non-zeros
+        for (int x = mesh->xstart; x <= mesh->xend; x++) {
           for (int y = mesh->ystart; y <= mesh->yend; y++) {
-            for (int z = 0; z < mesh->LocalNz; z++) {
-              const int localIndex = ROUND(index(mesh->xend, y, z));
-              ASSERT2((localIndex >= 0) && (localIndex < nlocal));
-              const int num_fields = (z == 0) ? n2d + n3d : n3d;
-              for (int i = 0; i < num_fields; i++) {
-                d_nnz[localIndex + i] -= (n3d + n2d);
-                o_nnz[localIndex + i] += (n3d + n2d);
-              }
-            }
-          }
-        }
-      }
 
-      // Y boundaries
-      if (ncells_y != 0) {
-        for (int x = mesh->xstart; x <= mesh->xend; x++) {
-          // Default to no boundary
-          // NOTE: This assumes that communications in Y are to other
-          //   processors. If Y is communicated with this processor (e.g. NYPE=1)
-          //   then this will result in PETSc warnings about out of range allocations
-
-          // z = 0 case
-          int localIndex = ROUND(index(x, mesh->ystart, 0));
-          ASSERT2(localIndex >= 0);
-
-          // All 2D and 3D fields
-          for (int i = 0; i < n2d + n3d; i++) {
-            o_nnz[localIndex + i] += (n3d + n2d);
-            d_nnz[localIndex + i] -= (n3d + n2d);
-          }
+            int ind0 = ROUND(index(x, y, 0)) - Istart;
 
-          for (int z = 1; z < mesh->LocalNz; z++) {
-            localIndex = ROUND(index(x, mesh->ystart, z));
+            // 2D fields
+            for (int i = 0; i < n2d; i++) {
+              PetscInt row = ind0 + i;
+              //if (row < Istart || row >= Iend) continue;
+              // Loop through each point in the stencil
+              for (const auto& [x_off, y_off] : xyoffsets) {
+                int xi = x + x_off;
+                int yi = y + y_off;
+                if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
+                    || (yi >= mesh->LocalNy)) {
+                  continue;
+                }
 
-            // Only 3D fields
-            for (int i = 0; i < n3d; i++) {
-              o_nnz[localIndex + i] += (n3d + n2d);
-              d_nnz[localIndex + i] -= (n3d + n2d);
+                int ind2 = ROUND(index(xi, yi, 0));
+                if (ind2 < 0) {
+                  continue; // A boundary point
+                }
+
+                // Depends on all variables on this cell
+                for (int j = 0; j < n2d; j++) {
+                  PetscInt col = ind2 + j;
+                  if (col >= Istart && col < Iend) {
+                    d_nnz_map2d[row].insert(col);
+                  } else {
+                    o_nnz_map2d[row].insert(col);
+                  }
+                }
+              }
             }
-          }
+            // 3D fields
+            for (int z = 0; z < mesh->LocalNz; z++) {
+              int ind = ROUND(index(x, y, z)) - Istart;
 
-          // z = 0 case
-          localIndex = ROUND(index(x, mesh->yend, 0));
-          // All 2D and 3D fields
-          for (int i = 0; i < n2d + n3d; i++) {
-            o_nnz[localIndex + i] += (n3d + n2d);
-            d_nnz[localIndex + i] -= (n3d + n2d);
-          }
+              for (int i = 0; i < n3d; i++) {
+                PetscInt row = ind + i;
+                //if (row < Istart || row >= Iend) continue;
+                if (z == 0) {
+                  row += n2d;
+                }
 
-          for (int z = 1; z < mesh->LocalNz; z++) {
-            localIndex = ROUND(index(x, mesh->yend, z));
+                // Depends on 2D fields
+                for (int j = 0; j < n2d; j++) {
+                  PetscInt col = ind0 + j;
+                  if (col >= Istart && col < Iend) {
+                    d_nnz_map2d[row].insert(col);
+                  } else {
+                    o_nnz_map2d[row].insert(col);
+                  }
+                }
 
-            // Only 3D fields
-            for (int i = 0; i < n3d; i++) {
-              o_nnz[localIndex + i] += (n3d + n2d);
-              d_nnz[localIndex + i] -= (n3d + n2d);
-            }
-          }
-        }
+                // Star pattern
+                for (const auto& [x_off, y_off] : xyoffsets) {
+                  int xi = x + x_off;
+                  int yi = y + y_off;
 
-        for (RangeIterator it = mesh->iterateBndryLowerY(); !it.isDone(); it++) {
-          // A boundary, so no communication
+                  if ((xi < 0) || (yi < 0) || (xi > mesh->LocalNx)
+                      || (yi > mesh->LocalNy)) {
+                    continue;
+                  }
 
-          // z = 0 case
-          int localIndex = ROUND(index(it.ind, mesh->ystart, 0));
-          if (localIndex < 0) {
-            // This can occur because it.ind includes values in x boundary e.g. x=0
-            continue;
-          }
-          // All 2D and 3D fields
-          for (int i = 0; i < n2d + n3d; i++) {
-            o_nnz[localIndex + i] -= (n3d + n2d);
-          }
+                  int ind2 = ROUND(index(xi, yi, 0));
+                  if (ind2 < 0) {
+                    continue; // Boundary point
+                  }
 
-          for (int z = 1; z < mesh->LocalNz; z++) {
-            int localIndex = ROUND(index(it.ind, mesh->ystart, z));
+                  if (z == 0) {
+                    ind2 += n2d;
+                  }
 
-            // Only 3D fields
-            for (int i = 0; i < n3d; i++) {
-              o_nnz[localIndex + i] -= (n3d + n2d);
+                  // 3D fields on this cell
+                  for (int j = 0; j < n3d; j++) {
+                    PetscInt col = ind2 + j;
+                    if (col >= Istart && col < Iend) {
+                      d_nnz_map3d[row].insert(col);
+                    } else {
+                      o_nnz_map3d[row].insert(col);
+                    }
+                  }
+                }
+              }
             }
           }
         }
 
-        for (RangeIterator it = mesh->iterateBndryUpperY(); !it.isDone(); it++) {
-          // A boundary, so no communication
-
-          // z = 0 case
-          const int localIndex = ROUND(index(it.ind, mesh->yend, 0));
-          if (localIndex < 0) {
-            continue; // Out of domain
-          }
-
-          // All 2D and 3D fields
-          for (int i = 0; i < n2d + n3d; i++) {
-            o_nnz[localIndex + i] -= (n3d + n2d);
-          }
-
-          for (int z = 1; z < mesh->LocalNz; z++) {
-            const int localIndex = ROUND(index(it.ind, mesh->yend, z));
+        d_nnz.reserve(localN);
+        d_nnz.reserve(localN);
 
-            // Only 3D fields
-            for (int i = 0; i < n3d; i++) {
-              o_nnz[localIndex + i] -= (n3d + n2d);
-            }
-          }
+        for (int i = 0; i < localN; ++i) {
+          //Assume all elements in the z direction are potentially coupled
+          d_nnz.emplace_back(d_nnz_map3d[i].size() * mesh->LocalNz
+                             + d_nnz_map2d[i].size());
+          o_nnz.emplace_back(o_nnz_map3d[i].size() * mesh->LocalNz
+                             + o_nnz_map2d[i].size());
         }
       }
 
       output_progress.write("Pre-allocating Jacobian\n");
-
       // Pre-allocate
-      MatMPIAIJSetPreallocation(Jfd, 0, d_nnz.data(), 0, o_nnz.data());
-      MatSeqAIJSetPreallocation(Jfd, 0, d_nnz.data());
-      MatSetUp(Jfd);
-      MatSetOption(Jfd, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
+      MatMPIAIJSetPreallocation(Jmf, 0, d_nnz.data(), 0, o_nnz.data());
+      MatSeqAIJSetPreallocation(Jmf, 0, d_nnz.data());
+      MatSetUp(Jmf);
+      MatSetOption(Jmf, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
 
-      // Determine which row/columns of the matrix are locally owned
-      int Istart, Iend;
-      MatGetOwnershipRange(Jfd, &Istart, &Iend);
-
-      // Convert local into global indices
-      // Note: Not in the boundary cells, to keep -1 values
-      for (const auto& i : mesh->getRegion3D("RGN_NOBNDRY")) {
-        index[i] += Istart;
-      }
-
-      // Now communicate to fill guard cells
-      mesh->communicate(index);
+      //  // Determine which row/columns of the matrix are locally owned
+      //  int Istart, Iend;
+      //  MatGetOwnershipRange(Jmf, &Istart, &Iend);
 
       //////////////////////////////////////////////////
       // Mark non-zero entries
 
       output_progress.write("Marking non-zero Jacobian entries\n");
-
-      const PetscScalar val = 1.0;
-
+      PetscScalar val = 1.0;
       for (int x = mesh->xstart; x <= mesh->xend; x++) {
         for (int y = mesh->ystart; y <= mesh->yend; y++) {
 
-          const int ind0 = ROUND(index(x, y, 0));
+          int ind0 = ROUND(index(x, y, 0));
 
           // 2D fields
           for (int i = 0; i < n2d; i++) {
-            const PetscInt row = ind0 + i;
-
-            // Loop through each point in the 5-point stencil
-            for (const auto& xyoffset : xyoffsets) {
-              const int xi = x + xyoffset.first;
-              const int yi = y + xyoffset.second;
+            PetscInt row = ind0 + i;
 
-              if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
-                  || (yi >= mesh->LocalNy)) {
+            // Loop through each point in the stencil
+            for (const auto& [x_off, y_off] : xyoffsets) {
+              int xi = x + x_off;
+              int yi = y + y_off;
+              if ((xi < 0) || (yi < 0) || (xi > mesh->LocalNx) || (yi > mesh->LocalNy)) {
                 continue;
               }
 
-              const int ind2 = ROUND(index(xi, yi, 0));
-
+              int ind2 = ROUND(index(xi, yi, 0));
               if (ind2 < 0) {
                 continue; // A boundary point
               }
 
               // Depends on all variables on this cell
               for (int j = 0; j < n2d; j++) {
-                const PetscInt col = ind2 + j;
-                ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
+                PetscInt col = ind2 + j;
+                ierr = MatSetValues(Jmf, 1, &row, 1, &col, &val, INSERT_VALUES);
                 CHKERRQ(ierr);
               }
             }
           }
-
           // 3D fields
           for (int z = 0; z < mesh->LocalNz; z++) {
-
-            const int ind = ROUND(index(x, y, z));
+            int ind = ROUND(index(x, y, z));
 
             for (int i = 0; i < n3d; i++) {
               PetscInt row = ind + i;
@@ -545,94 +472,93 @@ int SNESSolver::init() {
 
               // Depends on 2D fields
               for (int j = 0; j < n2d; j++) {
-                const PetscInt col = ind0 + j;
-                ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
+                PetscInt col = ind0 + j;
+                ierr = MatSetValues(Jmf, 1, &row, 1, &col, &val, INSERT_VALUES);
                 CHKERRQ(ierr);
               }
 
               // Star pattern
-              for (const auto& xyoffset : xyoffsets) {
-                const int xi = x + xyoffset.first;
-                const int yi = y + xyoffset.second;
+              for (const auto& [x_off, y_off] : xyoffsets) {
+                int xi = x + x_off;
+                int yi = y + y_off;
 
                 if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
                     || (yi >= mesh->LocalNy)) {
                   continue;
                 }
-
-                int ind2 = ROUND(index(xi, yi, z));
-                if (ind2 < 0) {
-                  continue; // Boundary point
-                }
-
-                if (z == 0) {
-                  ind2 += n2d;
-                }
-
-                // 3D fields on this cell
-                for (int j = 0; j < n3d; j++) {
-                  const PetscInt col = ind2 + j;
-                  ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
-                  if (ierr != 0) {
-                    output.write("ERROR: {} : ({}, {}) -> ({}, {}) : {} -> {}\n", row, x,
-                                 y, xi, yi, ind2, ind2 + n3d - 1);
+                for (int zi = 0; zi < mesh->LocalNz; ++zi) {
+                  int ind2 = ROUND(index(xi, yi, zi));
+                  if (ind2 < 0) {
+                    continue; // Boundary point
                   }
-                  CHKERRQ(ierr);
-                }
-              }
-
-              int nz = mesh->LocalNz;
-              if (nz > 1) {
-                // Multiple points in z
 
-                int zp = (z + 1) % nz;
-
-                int ind2 = ROUND(index(x, y, zp));
-                if (zp == 0) {
-                  ind2 += n2d;
-                }
-                for (int j = 0; j < n3d; j++) {
-                  const PetscInt col = ind2 + j;
-                  ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
-                  CHKERRQ(ierr);
-                }
+                  if (z == 0) {
+                    ind2 += n2d;
+                  }
 
-                int zm = (z - 1 + nz) % nz;
-                ind2 = ROUND(index(x, y, zm));
-                if (zm == 0) {
-                  ind2 += n2d;
-                }
-                for (int j = 0; j < n3d; j++) {
-                  const PetscInt col = ind2 + j;
-                  ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
-                  CHKERRQ(ierr);
+                  // 3D fields on this cell
+                  for (int j = 0; j < n3d; j++) {
+                    PetscInt col = ind2 + j;
+                    //printf("%d %d\n",row,col);
+                    ierr = MatSetValues(Jmf, 1, &row, 1, &col, &val, INSERT_VALUES);
+
+                    if (ierr != 0) {
+                      output.write("ERROR: {} {} : ({}, {}) -> ({}, {}) : {} -> {}\n",
+                                   row, x, y, xi, yi, ind2, ind2 + n3d - 1);
+                    }
+                    CHKERRQ(ierr);
+                  }
                 }
               }
             }
           }
         }
       }
+
       // Finished marking non-zero entries
 
       output_progress.write("Assembling Jacobian matrix\n");
 
       // Assemble Matrix
-      MatAssemblyBegin(Jfd, MAT_FINAL_ASSEMBLY);
-      MatAssemblyEnd(Jfd, MAT_FINAL_ASSEMBLY);
+      MatAssemblyBegin(Jmf, MAT_FINAL_ASSEMBLY);
+      MatAssemblyEnd(Jmf, MAT_FINAL_ASSEMBLY);
+
+      //The above will probably miss some non-zero entries at process boundaries
+      //Making sure the colouring matrix is symmetric will in some/all(?)
+      //of the missing non-zeros
+      if ((*options)["force_symmetric_coloring"]
+              .doc("Modifies coloring matrix to force it to be symmetric")
+              .withDefault<bool>(false)) {
+        Mat Jmf_T;
+        MatCreateTranspose(Jmf, &Jmf_T);
+        MatAXPY(Jmf, 1, Jmf_T, DIFFERENT_NONZERO_PATTERN);
+      }
 
       output_progress.write("Creating Jacobian coloring\n");
-      updateColoring();
 
-      if (prune_jacobian) {
-        // Will remove small elements from the Jacobian.
-        // Save a copy to recover from over-pruning
-        ierr = MatDuplicate(Jfd, MAT_SHARE_NONZERO_PATTERN, &Jfd_original);
-        CHKERRQ(ierr);
-      }
+      ISColoring iscoloring;
+
+      MatColoring coloring; // This new in PETSc 3.5
+      MatColoringCreate(Jmf, &coloring);
+      MatColoringSetType(coloring, MATCOLORINGSL);
+      MatColoringSetFromOptions(coloring);
+      // Calculate index sets
+      MatColoringApply(coloring, &iscoloring);
+      MatColoringDestroy(&coloring);
+
+      // Create data structure for SNESComputeJacobianDefaultColor
+      MatFDColoringCreate(Jmf, iscoloring, &fdcoloring);
+      // Set the function to difference
+      MatFDColoringSetFunction(
+          fdcoloring, reinterpret_cast<PetscErrorCode (*)()>(FormFunctionForColoring),
+          this);
+      MatFDColoringSetFromOptions(fdcoloring);
+      MatFDColoringSetUp(Jmf, iscoloring, fdcoloring);
+      ISColoringDestroy(&iscoloring);
+
+      SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianDefaultColor, fdcoloring);
     } else {
       // Brute force calculation
-      // There is usually no reason to use this, except as a check of
-      // the coloring calculation.
 
       MatCreateAIJ(
           BoutComm::get(), nlocal, nlocal,  // Local sizes
@@ -640,20 +566,17 @@ int SNESSolver::init() {
           3, // Number of nonzero entries in diagonal portion of local submatrix
           nullptr,
           0, // Number of nonzeros per row in off-diagonal portion of local submatrix
-          nullptr, &Jfd);
-
-      if (matrix_free_operator) {
-        SNESSetJacobian(snes, Jmf, Jfd, SNESComputeJacobianDefault, this);
-      } else {
-        SNESSetJacobian(snes, Jfd, Jfd, SNESComputeJacobianDefault, this);
-      }
-
-      MatSetOption(Jfd, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);
+          nullptr, &Jmf);
+#if PETSC_VERSION_GE(3, 4, 0)
+      SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianDefault, this);
+#else
+      // Before 3.4
+      SNESSetJacobian(snes, Jmf, Jmf, SNESDefaultComputeJacobian, this);
+#endif
+      MatSetOption(Jmf, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);
     }
 
     // Re-use Jacobian
-    // Note: If the 'Amat' Jacobian is matrix free, SNESComputeJacobian
-    //       always updates its reference 'u' vector every nonlinear iteration
     SNESSetLagJacobian(snes, lag_jacobian);
     // Set Jacobian and preconditioner to persist across time steps
     SNESSetLagJacobianPersists(snes, PETSC_TRUE);
@@ -664,8 +587,8 @@ int SNESSolver::init() {
   // Set tolerances
   SNESSetTolerances(snes, atol, rtol, stol, maxits, PETSC_DEFAULT);
 
-  // Force SNES to take at least one nonlinear iteration.
-  // This may prevent the solver from getting stuck in false steady state conditions
+// Force SNES to take at least one nonlinear iteration.
+// This may prevent the solver from getting stuck in false steady state conditions
 #if PETSC_VERSION_GE(3, 8, 0)
   SNESSetForceIteration(snes, PETSC_TRUE);
 #endif
@@ -747,7 +670,6 @@ int SNESSolver::init() {
 
 int SNESSolver::run() {
   TRACE("SNESSolver::run()");
-  int ierr;
   // Set initial guess at the solution from variables
   {
     BoutReal* xdata = nullptr;
@@ -764,63 +686,7 @@ int SNESSolver::run() {
     bool looping = true;
     int snes_failures = 0; // Count SNES convergence failures
     int saved_jacobian_lag = 0;
-    int loop_count = 0;
     do {
-      if (scale_vars) {
-        // Individual variable scaling
-        // Note: If variables are rescaled then the Jacobian columns
-        //       need to be scaled or recalculated
-
-        if (loop_count % 100 == 0) {
-          // Rescale state (snes_x) so that all quantities are around 1
-          // If quantities are near zero then RTOL is used
-          int istart, iend;
-          VecGetOwnershipRange(snes_x, &istart, &iend);
-
-          // Take ownership of snes_x and var_scaling_factors data
-          PetscScalar* snes_x_data = nullptr;
-          ierr = VecGetArray(snes_x, &snes_x_data);
-          CHKERRQ(ierr);
-          PetscScalar* x1_data;
-          ierr = VecGetArray(x1, &x1_data);
-          CHKERRQ(ierr);
-          PetscScalar* var_scaling_factors_data;
-          ierr = VecGetArray(var_scaling_factors, &var_scaling_factors_data);
-          CHKERRQ(ierr);
-
-          // Normalise each value in the state
-          // Limit normalisation so scaling factor is never smaller than rtol
-          for (int i = 0; i < iend - istart; ++i) {
-            const PetscScalar norm =
-                BOUTMAX(std::abs(snes_x_data[i]), rtol / var_scaling_factors_data[i]);
-            snes_x_data[i] /= norm;
-            x1_data[i] /= norm; // Update history for predictor
-            var_scaling_factors_data[i] *= norm;
-          }
-
-          // Restore vector underlying data
-          ierr = VecRestoreArray(var_scaling_factors, &var_scaling_factors_data);
-          CHKERRQ(ierr);
-          ierr = VecRestoreArray(x1, &x1_data);
-          CHKERRQ(ierr);
-          ierr = VecRestoreArray(snes_x, &snes_x_data);
-          CHKERRQ(ierr);
-
-          if (diagnose) {
-            // Print maximum and minimum scaling factors
-            PetscReal max_scale, min_scale;
-            VecMax(var_scaling_factors, nullptr, &max_scale);
-            VecMin(var_scaling_factors, nullptr, &min_scale);
-            output.write("Var scaling: {} -> {}\n", min_scale, max_scale);
-          }
-
-          // Force recalculation of the Jacobian
-          SNESGetLagJacobian(snes, &saved_jacobian_lag);
-          SNESSetLagJacobian(snes, 1);
-        }
-      }
-      ++loop_count;
-
       // Copy the state (snes_x) into initial values (x0)
       VecCopy(snes_x, x0);
 
@@ -862,13 +728,6 @@ int SNESSolver::run() {
       // Find out if converged
       SNESConvergedReason reason;
       SNESGetConvergedReason(snes, &reason);
-
-      // Get number of iterations
-      int nl_its;
-      SNESGetIterationNumber(snes, &nl_its);
-      int lin_its;
-      SNESGetLinearSolveIterations(snes, &lin_its);
-
       if ((ierr != 0) or (reason < 0)) {
         // Diverged or SNES failed
 
@@ -899,19 +758,6 @@ int SNESSolver::run() {
         VecCopy(x0, snes_x);
 
         // Recalculate the Jacobian
-        if (jacobian_pruned and (snes_failures > 2) and (4 * lin_its > 3 * maxl)) {
-          // Taking 3/4 of maximum linear iterations on average per linear step
-          // May indicate a preconditioner problem.
-          // Restore pruned non-zero elements
-          if (diagnose) {
-            output.write("\nRestoring Jacobian\n");
-          }
-          ierr = MatCopy(Jfd_original, Jfd, DIFFERENT_NONZERO_PATTERN);
-          CHKERRQ(ierr);
-          // The non-zero pattern has changed, so update coloring
-          updateColoring();
-          jacobian_pruned = false; // Reset flag. Will be set after pruning.
-        }
         if (saved_jacobian_lag == 0) {
           SNESGetLagJacobian(snes, &saved_jacobian_lag);
           SNESSetLagJacobian(snes, 1);
@@ -944,24 +790,16 @@ int SNESSolver::run() {
         time1 = simtime;
       }
 
+      int nl_its;
+      SNESGetIterationNumber(snes, &nl_its);
+
       if (nl_its == 0) {
         // This can occur even with SNESSetForceIteration
         // Results in simulation state freezing and rapidly going to the end
 
-        if (scale_vars) {
-          // scaled_x <- snes_x * var_scaling_factors
-          ierr = VecPointwiseMult(scaled_x, snes_x, var_scaling_factors);
-          CHKERRQ(ierr);
-
-          const BoutReal* xdata = nullptr;
-          ierr = VecGetArrayRead(scaled_x, &xdata);
-          CHKERRQ(ierr);
-          load_vars(const_cast<BoutReal*>(xdata));
-          ierr = VecRestoreArrayRead(scaled_x, &xdata);
-          CHKERRQ(ierr);
-        } else {
+        {
           const BoutReal* xdata = nullptr;
-          ierr = VecGetArrayRead(snes_x, &xdata);
+          int ierr = VecGetArrayRead(snes_x, &xdata);
           CHKERRQ(ierr);
           load_vars(const_cast<BoutReal*>(xdata));
           ierr = VecRestoreArrayRead(snes_x, &xdata);
@@ -988,6 +826,9 @@ int SNESSolver::run() {
       if (diagnose) {
         // Gather and print diagnostic information
 
+        int lin_its;
+        SNESGetLinearSolveIterations(snes, &lin_its);
+
         output.print("\r"); // Carriage return for printing to screen
         output.write("Time: {}, timestep: {}, nl iter: {}, lin iter: {}, reason: {}",
                      simtime, timestep, nl_its, lin_its, static_cast<int>(reason));
@@ -998,51 +839,6 @@ int SNESSolver::run() {
         output.write("\n");
       }
 
-#if PETSC_VERSION_GE(3, 20, 0)
-      // MatFilter and MatEliminateZeros(Mat, bool) require PETSc >= 3.20
-      if (jacobian_recalculated and prune_jacobian) {
-        jacobian_recalculated = false; // Reset flag
-
-        // Remove small elements from the Jacobian and recompute the coloring
-        // Only do this if there are a significant number of small elements.
-        int small_elements = 0;
-        int total_elements = 0;
-
-        // Get index of rows owned by this processor
-        int rstart, rend;
-        MatGetOwnershipRange(Jfd, &rstart, &rend);
-
-        PetscInt ncols;
-        const PetscScalar* vals;
-        for (int row = rstart; row < rend; row++) {
-          MatGetRow(Jfd, row, &ncols, nullptr, &vals);
-          for (int col = 0; col < ncols; col++) {
-            if (std::abs(vals[col]) < prune_abstol) {
-              ++small_elements;
-            }
-            ++total_elements;
-          }
-          MatRestoreRow(Jfd, row, &ncols, nullptr, &vals);
-        }
-
-        if (small_elements > prune_fraction * total_elements) {
-          if (diagnose) {
-            output.write("\nPruning Jacobian elements: {} / {}\n", small_elements,
-                         total_elements);
-          }
-
-          // Prune Jacobian, keeping diagonal elements
-          ierr = MatFilter(Jfd, prune_abstol, PETSC_TRUE, PETSC_TRUE);
-
-          // Update the coloring from Jfd matrix
-          updateColoring();
-
-          // Mark the Jacobian as pruned. This is so that it is only restored if pruned.
-          jacobian_pruned = true;
-        }
-      }
-#endif // PETSC_VERSION_GE(3,20,0)
-
       if (looping) {
         if (nl_its <= lower_its) {
           // Increase timestep slightly
@@ -1059,18 +855,7 @@ int SNESSolver::run() {
     } while (looping);
 
     // Put the result into variables
-    if (scale_vars) {
-      // scaled_x <- snes_x * var_scaling_factors
-      int ierr = VecPointwiseMult(scaled_x, snes_x, var_scaling_factors);
-      CHKERRQ(ierr);
-
-      const BoutReal* xdata = nullptr;
-      ierr = VecGetArrayRead(scaled_x, &xdata);
-      CHKERRQ(ierr);
-      load_vars(const_cast<BoutReal*>(xdata));
-      ierr = VecRestoreArrayRead(scaled_x, &xdata);
-      CHKERRQ(ierr);
-    } else {
+    {
       const BoutReal* xdata = nullptr;
       int ierr = VecGetArrayRead(snes_x, &xdata);
       CHKERRQ(ierr);
@@ -1091,27 +876,12 @@ int SNESSolver::run() {
 // f = rhs
 PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {
   // Get data from PETSc into BOUT++ fields
-  if (scale_vars) {
-    // scaled_x <- x * var_scaling_factors
-    int ierr = VecPointwiseMult(scaled_x, x, var_scaling_factors);
-    CHKERRQ(ierr);
-
-    const BoutReal* xdata = nullptr;
-    ierr = VecGetArrayRead(scaled_x, &xdata);
-    CHKERRQ(ierr);
-    load_vars(const_cast<BoutReal*>(
-        xdata)); // const_cast needed due to load_vars API. Not writing to xdata.
-    ierr = VecRestoreArrayRead(scaled_x, &xdata);
-    CHKERRQ(ierr);
-  } else {
-    const BoutReal* xdata = nullptr;
-    int ierr = VecGetArrayRead(x, &xdata);
-    CHKERRQ(ierr);
-    load_vars(const_cast<BoutReal*>(
-        xdata)); // const_cast needed due to load_vars API. Not writing to xdata.
-    ierr = VecRestoreArrayRead(x, &xdata);
-    CHKERRQ(ierr);
-  }
+  const BoutReal* xdata = nullptr;
+  int ierr = VecGetArrayRead(x, &xdata);
+  CHKERRQ(ierr);
+  load_vars(const_cast<BoutReal*>(xdata));
+  ierr = VecRestoreArrayRead(x, &xdata);
+  CHKERRQ(ierr);
 
   try {
     // Call RHS function
@@ -1129,7 +899,7 @@ PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {
 
   // Copy derivatives back
   BoutReal* fdata = nullptr;
-  int ierr = VecGetArray(f, &fdata);
+  ierr = VecGetArray(f, &fdata);
   CHKERRQ(ierr);
   save_derivs(fdata);
   ierr = VecRestoreArray(f, &fdata);
@@ -1166,12 +936,6 @@ PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {
   }
   };
 
-  if (scale_rhs) {
-    // f <- f * rhs_scaling_factors
-    ierr = VecPointwiseMult(f, f, rhs_scaling_factors);
-    CHKERRQ(ierr);
-  }
-
   return 0;
 }
 
@@ -1218,133 +982,4 @@ PetscErrorCode SNESSolver::precon(Vec x, Vec f) {
   return 0;
 }
 
-PetscErrorCode SNESSolver::scaleJacobian(Mat Jac_new) {
-  jacobian_recalculated = true;
-
-  if (!scale_rhs) {
-    return 0; // Not scaling the RHS values
-  }
-
-  int ierr;
-
-  // Get index of rows owned by this processor
-  int rstart, rend;
-  MatGetOwnershipRange(Jac_new, &rstart, &rend);
-
-  // Check that the vector has the same ownership range
-  int istart, iend;
-  VecGetOwnershipRange(jac_row_inv_norms, &istart, &iend);
-  if ((rstart != istart) or (rend != iend)) {
-    throw BoutException("Ownership ranges different: [{}, {}) and [{}, {})\n", rstart,
-                        rend, istart, iend);
-  }
-
-  // Calculate the norm of each row of the Jacobian
-  PetscScalar* row_inv_norm_data;
-  ierr = VecGetArray(jac_row_inv_norms, &row_inv_norm_data);
-  CHKERRQ(ierr);
-
-  PetscInt ncols;
-  const PetscScalar* vals;
-  for (int row = rstart; row < rend; ++row) {
-    MatGetRow(Jac_new, row, &ncols, nullptr, &vals);
-
-    // Calculate a norm of this row of the Jacobian
-    PetscScalar norm = 0.0;
-    for (int col = 0; col < ncols; col++) {
-      PetscScalar absval = std::abs(vals[col]);
-      if (absval > norm) {
-        norm = absval;
-      }
-      // Can we identify small elements and remove them?
-      // so we don't need to calculate them next time
-    }
-
-    // Store in the vector as 1 / norm
-    row_inv_norm_data[row - rstart] = 1. / norm;
-
-    MatRestoreRow(Jac_new, row, &ncols, nullptr, &vals);
-  }
-
-  ierr = VecRestoreArray(jac_row_inv_norms, &row_inv_norm_data);
-  CHKERRQ(ierr);
-
-  // Modify the RHS scaling: factor = factor / norm
-  ierr = VecPointwiseMult(rhs_scaling_factors, rhs_scaling_factors, jac_row_inv_norms);
-  CHKERRQ(ierr);
-
-  if (diagnose) {
-    // Print maximum and minimum scaling factors
-    PetscReal max_scale, min_scale;
-    VecMax(rhs_scaling_factors, nullptr, &max_scale);
-    VecMin(rhs_scaling_factors, nullptr, &min_scale);
-    output.write("RHS scaling: {} -> {}\n", min_scale, max_scale);
-  }
-
-  // Scale the Jacobian rows by multiplying on the left by 1/norm
-  ierr = MatDiagonalScale(Jac_new, jac_row_inv_norms, nullptr);
-  CHKERRQ(ierr);
-
-  return 0;
-}
-
-///
-/// Input Parameters:
-///   snes - nonlinear solver object
-///   x1 - location at which to evaluate Jacobian
-///   ctx - MatFDColoring context or NULL
-///
-/// Output Parameters:
-///   Jac - Jacobian matrix (not altered in this routine)
-///   Jac_new - newly computed Jacobian matrix to use with preconditioner (generally the same as
-///   Jac)
-static PetscErrorCode ComputeJacobianScaledColor(SNES snes, Vec x1, Mat Jac, Mat Jac_new,
-                                                 void* ctx) {
-  PetscErrorCode err = SNESComputeJacobianDefaultColor(snes, x1, Jac, Jac_new, ctx);
-  CHKERRQ(err);
-
-  if ((err != 0) or (ctx == nullptr)) {
-    return err;
-  }
-
-  // Get the the SNESSolver pointer from the function call context
-  SNESSolver* fctx = nullptr;
-  err = MatFDColoringGetFunction(static_cast<MatFDColoring>(ctx), nullptr,
-                                 reinterpret_cast<void**>(&fctx));
-  CHKERRQ(err);
-
-  // Call the SNESSolver function
-  return fctx->scaleJacobian(Jac_new);
-}
-
-void SNESSolver::updateColoring() {
-  // Re-calculate the coloring
-  MatColoring coloring = NULL;
-  MatColoringCreate(Jfd, &coloring);
-  MatColoringSetType(coloring, MATCOLORINGSL);
-  MatColoringSetFromOptions(coloring);
-
-  // Calculate new index sets
-  ISColoring iscoloring = NULL;
-  MatColoringApply(coloring, &iscoloring);
-  MatColoringDestroy(&coloring);
-
-  // Replace the old coloring with the new one
-  MatFDColoringDestroy(&fdcoloring);
-  MatFDColoringCreate(Jfd, iscoloring, &fdcoloring);
-  MatFDColoringSetFunction(
-      fdcoloring, reinterpret_cast<PetscErrorCode (*)()>(FormFunctionForColoring), this);
-  MatFDColoringSetFromOptions(fdcoloring);
-  MatFDColoringSetUp(Jfd, iscoloring, fdcoloring);
-  ISColoringDestroy(&iscoloring);
-
-  // Replace the CTX pointer in SNES Jacobian
-  if (matrix_free_operator) {
-    // Use matrix-free calculation for operator, finite difference for preconditioner
-    SNESSetJacobian(snes, Jmf, Jfd, ComputeJacobianScaledColor, fdcoloring);
-  } else {
-    SNESSetJacobian(snes, Jfd, Jfd, ComputeJacobianScaledColor, fdcoloring);
-  }
-}
-
 #endif // BOUT_HAS_PETSC

From 19ec182e965ecbea66d731a8c2218c148d84a8a8 Mon Sep 17 00:00:00 2001
From: Seimon Powell <seimon.powell@ukaea.uk>
Date: Thu, 24 Apr 2025 15:51:58 +0100
Subject: [PATCH 2/9] Change message to match rest of output

---
 src/solver/impls/snes/snes.cxx | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index aec986511f..176bcd9342 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -298,9 +298,9 @@ int SNESSolver::init() {
                          .doc("Extent of stencil (cross)")
                          .withDefault<int>(0);
       //Set n_taxi 2 if nothing else is set
-      //TODO: Probably a better way to do this
+      //Probably a better way to do this
       if (n_square == 0 && n_taxi == 0 && n_cross == 0) {
-        output_info.write("Setting beuler:stencil:taxi = 2\n");
+        output_info.write("Setting solver:stencil:taxi = 2\n");
         n_taxi = 2;
       }
 

From 697f729c723c4f7eec5fdb8336b0c5d9df6cae28 Mon Sep 17 00:00:00 2001
From: Seimon Powell <seimon.powell@ukaea.uk>
Date: Thu, 24 Apr 2025 16:01:01 +0100
Subject: [PATCH 3/9] Obey the naming conventions

---
 src/solver/impls/snes/snes.cxx | 39 +++++++++++++++++-----------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index 176bcd9342..67b6d75d54 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -18,10 +18,10 @@
 
 class ColoringStencil {
 private:
-  bool static in_compact(int const i, int const j, int const n_compact) {
-    return std::abs(i) <= n_compact && std::abs(j) <= n_compact;
+  bool static isInSquare(int const i, int const j, int const n_square) {
+    return std::abs(i) <= n_square && std::abs(j) <= n_square;
   }
-  bool static in_cross(int const i, int const j, int const n_cross) {
+  bool static isInCross(int const i, int const j, int const n_cross) {
     if (i == 0) {
       return std::abs(j) <= n_cross;
     }
@@ -30,28 +30,28 @@ class ColoringStencil {
     }
     return false;
   }
-  bool static in_taxi(int const i, int const j, int const n_taxi) {
+  bool static isInTaxi(int const i, int const j, int const n_taxi) {
     return std::abs(i) + std::abs(j) <= n_taxi;
   }
 
 public:
-  auto static get_offsets(int n_compact, int n_taxi, int n_cross) {
-    ASSERT2(n_compact >= 0 && n_cross >= 0 && n_taxi >= 0
-            && n_compact + n_cross + n_taxi > 0);
+  auto static getOffsets(int n_square, int n_taxi, int n_cross) {
+    ASSERT2(n_square >= 0 && n_cross >= 0 && n_taxi >= 0
+            && n_square + n_cross + n_taxi > 0);
     auto inside = [&](int i, int j) {
-      return in_compact(i, j, n_compact) || in_taxi(i, j, n_taxi)
-             || in_cross(i, j, n_cross);
+      return isInSquare(i, j, n_square) || isInTaxi(i, j, n_taxi)
+             || isInCross(i, j, n_cross);
     };
-    std::vector<std::pair<int, int>> xyoffset;
-    auto loop_bound = std::max({n_compact, n_taxi, n_cross});
+    std::vector<std::pair<int, int>> xy_offsets;
+    auto loop_bound = std::max({n_square, n_taxi, n_cross});
     for (int i = -loop_bound; i <= loop_bound; ++i) {
       for (int j = -loop_bound; j <= loop_bound; ++j) {
         if (inside(i, j)) {
-          xyoffset.emplace_back(i, j);
+          xy_offsets.emplace_back(i, j);
         }
       }
     }
-    return xyoffset;
+    return xy_offsets;
   }
 };
 
@@ -304,10 +304,11 @@ int SNESSolver::init() {
         n_taxi = 2;
       }
 
-      auto const xyoffsets = ColoringStencil::get_offsets(n_square, n_taxi, n_cross);
+      auto const xy_offsets = ColoringStencil::getOffsets(n_square, n_taxi, n_cross);
       {
-        //This is nasty but can't think of a better and robust way to
+        //This is ugly but can't think of a better and robust way to
         //count the non-zeros for some arbitery stencil
+        //effectivly the same loop as the one that sets the non-zeros below
         std::vector<std::set<int>> d_nnz_map2d(localN);
         std::vector<std::set<int>> o_nnz_map2d(localN);
         std::vector<std::set<int>> d_nnz_map3d(localN);
@@ -323,7 +324,7 @@ int SNESSolver::init() {
               PetscInt row = ind0 + i;
               //if (row < Istart || row >= Iend) continue;
               // Loop through each point in the stencil
-              for (const auto& [x_off, y_off] : xyoffsets) {
+              for (const auto& [x_off, y_off] : xy_offsets) {
                 int xi = x + x_off;
                 int yi = y + y_off;
                 if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
@@ -369,7 +370,7 @@ int SNESSolver::init() {
                 }
 
                 // Star pattern
-                for (const auto& [x_off, y_off] : xyoffsets) {
+                for (const auto& [x_off, y_off] : xy_offsets) {
                   int xi = x + x_off;
                   int yi = y + y_off;
 
@@ -440,7 +441,7 @@ int SNESSolver::init() {
             PetscInt row = ind0 + i;
 
             // Loop through each point in the stencil
-            for (const auto& [x_off, y_off] : xyoffsets) {
+            for (const auto& [x_off, y_off] : xy_offsets) {
               int xi = x + x_off;
               int yi = y + y_off;
               if ((xi < 0) || (yi < 0) || (xi > mesh->LocalNx) || (yi > mesh->LocalNy)) {
@@ -478,7 +479,7 @@ int SNESSolver::init() {
               }
 
               // Star pattern
-              for (const auto& [x_off, y_off] : xyoffsets) {
+              for (const auto& [x_off, y_off] : xy_offsets) {
                 int xi = x + x_off;
                 int yi = y + y_off;
 

From 540ab7d39c0693fb9d5eebe9198a1d7b91ebdcac Mon Sep 17 00:00:00 2001
From: Ben Dudson <dudson2@llnl.gov>
Date: Thu, 24 Apr 2025 13:45:27 -0700
Subject: [PATCH 4/9] beuler: Merge improvements

PR #3009 included changes to allow matrix-free Jacobian-vector
products, while using the finite difference Jacobian for the
preconditioner. That gave significant speedups that will hopefully
also help now that the coloring stencil is fixed.
---
 src/solver/impls/snes/snes.cxx | 487 +++++++++++++++++++++++++++------
 1 file changed, 405 insertions(+), 82 deletions(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index 67b6d75d54..61d98adafc 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -157,12 +157,31 @@ SNESSolver::SNESSolver(Options* opts)
       matrix_free((*options)["matrix_free"]
                       .doc("Use matrix free Jacobian?")
                       .withDefault<bool>(false)),
+      matrix_free_operator((*options)["matrix_free_operator"]
+                               .doc("Use matrix free Jacobian-vector operator?")
+                               .withDefault<bool>(true)),
       lag_jacobian((*options)["lag_jacobian"]
                        .doc("Re-use the Jacobian this number of SNES iterations")
                        .withDefault(50)),
       use_coloring((*options)["use_coloring"]
                        .doc("Use matrix coloring to calculate Jacobian?")
-                       .withDefault<bool>(true)) {}
+                       .withDefault<bool>(true)),
+      jacobian_recalculated(false),
+      prune_jacobian((*options)["prune_jacobian"]
+                         .doc("Remove small elements in the Jacobian?")
+                         .withDefault<bool>(false)),
+      prune_abstol((*options)["prune_abstol"]
+                       .doc("Prune values with absolute values smaller than this")
+                       .withDefault<BoutReal>(1e-16)),
+      prune_fraction((*options)["prune_fraction"]
+                         .doc("Prune if fraction of small elements is larger than this")
+                         .withDefault<BoutReal>(0.2)),
+      scale_rhs((*options)["scale_rhs"]
+                    .doc("Scale time derivatives (Jacobian row scaling)?")
+                    .withDefault<bool>(false)),
+      scale_vars((*options)["scale_vars"]
+                     .doc("Scale variables (Jacobian column scaling)?")
+                     .withDefault<bool>(false)) {}
 
 int SNESSolver::init() {
 
@@ -202,12 +221,38 @@ int SNESSolver::init() {
 
   if (equation_form == BoutSnesEquationForm::rearranged_backward_euler) {
     // Need an intermediate vector for rearranged Backward Euler
-    VecDuplicate(snes_x, &delta_x);
+    ierr = VecDuplicate(snes_x, &delta_x);
+    CHKERRQ(ierr);
   }
 
   if (predictor) {
     // Storage for previous solution
-    VecDuplicate(snes_x, &x1);
+    ierr = VecDuplicate(snes_x, &x1);
+    CHKERRQ(ierr);
+  }
+
+  if (scale_rhs) {
+    // Storage for rhs factors, one per evolving variable
+    ierr = VecDuplicate(snes_x, &rhs_scaling_factors);
+    CHKERRQ(ierr);
+    // Set all factors to 1 to start with
+    ierr = VecSet(rhs_scaling_factors, 1.0);
+    CHKERRQ(ierr);
+    // Array to store inverse Jacobian row norms
+    ierr = VecDuplicate(snes_x, &jac_row_inv_norms);
+    CHKERRQ(ierr);
+  }
+
+  if (scale_vars) {
+    // Storage for var factors, one per evolving variable
+    ierr = VecDuplicate(snes_x, &var_scaling_factors);
+    CHKERRQ(ierr);
+    // Set all factors to 1 to start with
+    ierr = VecSet(var_scaling_factors, 1.0);
+    CHKERRQ(ierr);
+    // Storage for scaled 'x' state vectors
+    ierr = VecDuplicate(snes_x, &scaled_x);
+    CHKERRQ(ierr);
   }
 
   // Nonlinear solver interface (SNES)
@@ -227,7 +272,7 @@ int SNESSolver::init() {
   }
 
   // Set up the Jacobian
-  if (matrix_free) {
+  if (matrix_free or matrix_free_operator) {
     /*
       PETSc SNES matrix free Jacobian, using a different
       operator for differencing.
@@ -243,12 +288,17 @@ int SNESSolver::init() {
     // Set a function to be called for differencing
     // This can be a linearised form of the SNES function
     MatMFFDSetFunction(Jmf, FormFunctionForDifferencing, this);
+  }
 
+  if (matrix_free) {
+    // Use matrix free for both operator and preconditioner
     // Calculate Jacobian matrix free using FormFunctionForDifferencing
     SNESSetJacobian(snes, Jmf, Jmf, MatMFFDComputeJacobian, this);
 
   } else {
-    // Calculate the Jacobian using finite differences
+    // Calculate the Jacobian using finite differences.
+    // The finite difference Jacobian (Jfd) may be used for both operator
+    // and preconditioner or, if matrix_free_operator, in only the preconditioner.
     if (use_coloring) {
       // Use matrix coloring
       // This greatly reduces the number of times the rhs() function needs
@@ -266,17 +316,16 @@ int SNESSolver::init() {
 
       output_progress.write("Setting Jacobian matrix sizes\n");
 
-      int localN = getLocalN(); // Number of rows on this processor
       int n2d = f2d.size();
       int n3d = f3d.size();
 
-      // Set size of Matrix on each processor to localN x localN
-      MatCreate(BoutComm::get(), &Jmf);
-      MatSetSizes(Jmf, localN, localN, PETSC_DETERMINE, PETSC_DETERMINE);
-      MatSetFromOptions(Jmf);
+      // Set size of Matrix on each processor to nlocal x nlocal
+      MatCreate(BoutComm::get(), &Jfd);
+      MatSetSizes(Jfd, nlocal, nlocal, PETSC_DETERMINE, PETSC_DETERMINE);
+      MatSetFromOptions(Jfd);
       // Determine which row/columns of the matrix are locally owned
       int Istart, Iend;
-      MatGetOwnershipRange(Jmf, &Istart, &Iend);
+      MatGetOwnershipRange(Jfd, &Istart, &Iend);
       // Convert local into global indices
       // Note: Not in the boundary cells, to keep -1 values
       for (const auto& i : mesh->getRegion3D("RGN_NOBNDRY")) {
@@ -309,10 +358,10 @@ int SNESSolver::init() {
         //This is ugly but can't think of a better and robust way to
         //count the non-zeros for some arbitery stencil
         //effectivly the same loop as the one that sets the non-zeros below
-        std::vector<std::set<int>> d_nnz_map2d(localN);
-        std::vector<std::set<int>> o_nnz_map2d(localN);
-        std::vector<std::set<int>> d_nnz_map3d(localN);
-        std::vector<std::set<int>> o_nnz_map3d(localN);
+        std::vector<std::set<int>> d_nnz_map2d(nlocal);
+        std::vector<std::set<int>> o_nnz_map2d(nlocal);
+        std::vector<std::set<int>> d_nnz_map3d(nlocal);
+        std::vector<std::set<int>> o_nnz_map3d(nlocal);
         //Loop over every element in 2D to count the *unique* non-zeros
         for (int x = mesh->xstart; x <= mesh->xend; x++) {
           for (int y = mesh->ystart; y <= mesh->yend; y++) {
@@ -403,10 +452,10 @@ int SNESSolver::init() {
           }
         }
 
-        d_nnz.reserve(localN);
-        d_nnz.reserve(localN);
+        d_nnz.reserve(nlocal);
+        d_nnz.reserve(nlocal);
 
-        for (int i = 0; i < localN; ++i) {
+        for (int i = 0; i < nlocal; ++i) {
           //Assume all elements in the z direction are potentially coupled
           d_nnz.emplace_back(d_nnz_map3d[i].size() * mesh->LocalNz
                              + d_nnz_map2d[i].size());
@@ -417,14 +466,10 @@ int SNESSolver::init() {
 
       output_progress.write("Pre-allocating Jacobian\n");
       // Pre-allocate
-      MatMPIAIJSetPreallocation(Jmf, 0, d_nnz.data(), 0, o_nnz.data());
-      MatSeqAIJSetPreallocation(Jmf, 0, d_nnz.data());
-      MatSetUp(Jmf);
-      MatSetOption(Jmf, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
-
-      //  // Determine which row/columns of the matrix are locally owned
-      //  int Istart, Iend;
-      //  MatGetOwnershipRange(Jmf, &Istart, &Iend);
+      MatMPIAIJSetPreallocation(Jfd, 0, d_nnz.data(), 0, o_nnz.data());
+      MatSeqAIJSetPreallocation(Jfd, 0, d_nnz.data());
+      MatSetUp(Jfd);
+      MatSetOption(Jfd, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
 
       //////////////////////////////////////////////////
       // Mark non-zero entries
@@ -456,7 +501,7 @@ int SNESSolver::init() {
               // Depends on all variables on this cell
               for (int j = 0; j < n2d; j++) {
                 PetscInt col = ind2 + j;
-                ierr = MatSetValues(Jmf, 1, &row, 1, &col, &val, INSERT_VALUES);
+                ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
                 CHKERRQ(ierr);
               }
             }
@@ -474,7 +519,7 @@ int SNESSolver::init() {
               // Depends on 2D fields
               for (int j = 0; j < n2d; j++) {
                 PetscInt col = ind0 + j;
-                ierr = MatSetValues(Jmf, 1, &row, 1, &col, &val, INSERT_VALUES);
+                ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
                 CHKERRQ(ierr);
               }
 
@@ -500,8 +545,7 @@ int SNESSolver::init() {
                   // 3D fields on this cell
                   for (int j = 0; j < n3d; j++) {
                     PetscInt col = ind2 + j;
-                    //printf("%d %d\n",row,col);
-                    ierr = MatSetValues(Jmf, 1, &row, 1, &col, &val, INSERT_VALUES);
+                    ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES);
 
                     if (ierr != 0) {
                       output.write("ERROR: {} {} : ({}, {}) -> ({}, {}) : {} -> {}\n",
@@ -521,8 +565,8 @@ int SNESSolver::init() {
       output_progress.write("Assembling Jacobian matrix\n");
 
       // Assemble Matrix
-      MatAssemblyBegin(Jmf, MAT_FINAL_ASSEMBLY);
-      MatAssemblyEnd(Jmf, MAT_FINAL_ASSEMBLY);
+      MatAssemblyBegin(Jfd, MAT_FINAL_ASSEMBLY);
+      MatAssemblyEnd(Jfd, MAT_FINAL_ASSEMBLY);
 
       //The above will probably miss some non-zero entries at process boundaries
       //Making sure the colouring matrix is symmetric will in some/all(?)
@@ -530,36 +574,24 @@ int SNESSolver::init() {
       if ((*options)["force_symmetric_coloring"]
               .doc("Modifies coloring matrix to force it to be symmetric")
               .withDefault<bool>(false)) {
-        Mat Jmf_T;
-        MatCreateTranspose(Jmf, &Jmf_T);
-        MatAXPY(Jmf, 1, Jmf_T, DIFFERENT_NONZERO_PATTERN);
+        Mat Jfd_T;
+        MatCreateTranspose(Jfd, &Jfd_T);
+        MatAXPY(Jfd, 1, Jfd_T, DIFFERENT_NONZERO_PATTERN);
       }
 
       output_progress.write("Creating Jacobian coloring\n");
+      updateColoring();
 
-      ISColoring iscoloring;
-
-      MatColoring coloring; // This new in PETSc 3.5
-      MatColoringCreate(Jmf, &coloring);
-      MatColoringSetType(coloring, MATCOLORINGSL);
-      MatColoringSetFromOptions(coloring);
-      // Calculate index sets
-      MatColoringApply(coloring, &iscoloring);
-      MatColoringDestroy(&coloring);
-
-      // Create data structure for SNESComputeJacobianDefaultColor
-      MatFDColoringCreate(Jmf, iscoloring, &fdcoloring);
-      // Set the function to difference
-      MatFDColoringSetFunction(
-          fdcoloring, reinterpret_cast<PetscErrorCode (*)()>(FormFunctionForColoring),
-          this);
-      MatFDColoringSetFromOptions(fdcoloring);
-      MatFDColoringSetUp(Jmf, iscoloring, fdcoloring);
-      ISColoringDestroy(&iscoloring);
-
-      SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianDefaultColor, fdcoloring);
+      if (prune_jacobian) {
+        // Will remove small elements from the Jacobian.
+        // Save a copy to recover from over-pruning
+        ierr = MatDuplicate(Jfd, MAT_SHARE_NONZERO_PATTERN, &Jfd_original);
+        CHKERRQ(ierr);
+      }
     } else {
       // Brute force calculation
+      // There is usually no reason to use this, except as a check of
+      // the coloring calculation.
 
       MatCreateAIJ(
           BoutComm::get(), nlocal, nlocal,  // Local sizes
@@ -567,17 +599,20 @@ int SNESSolver::init() {
           3, // Number of nonzero entries in diagonal portion of local submatrix
           nullptr,
           0, // Number of nonzeros per row in off-diagonal portion of local submatrix
-          nullptr, &Jmf);
-#if PETSC_VERSION_GE(3, 4, 0)
-      SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianDefault, this);
-#else
-      // Before 3.4
-      SNESSetJacobian(snes, Jmf, Jmf, SNESDefaultComputeJacobian, this);
-#endif
-      MatSetOption(Jmf, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);
+          nullptr, &Jfd);
+
+      if (matrix_free_operator) {
+        SNESSetJacobian(snes, Jmf, Jfd, SNESComputeJacobianDefault, this);
+      } else {
+        SNESSetJacobian(snes, Jfd, Jfd, SNESComputeJacobianDefault, this);
+      }
+
+      MatSetOption(Jfd, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);
     }
 
     // Re-use Jacobian
+    // Note: If the 'Amat' Jacobian is matrix free, SNESComputeJacobian
+    //       always updates its reference 'u' vector every nonlinear iteration
     SNESSetLagJacobian(snes, lag_jacobian);
     // Set Jacobian and preconditioner to persist across time steps
     SNESSetLagJacobianPersists(snes, PETSC_TRUE);
@@ -588,8 +623,8 @@ int SNESSolver::init() {
   // Set tolerances
   SNESSetTolerances(snes, atol, rtol, stol, maxits, PETSC_DEFAULT);
 
-// Force SNES to take at least one nonlinear iteration.
-// This may prevent the solver from getting stuck in false steady state conditions
+  // Force SNES to take at least one nonlinear iteration.
+  // This may prevent the solver from getting stuck in false steady state conditions
 #if PETSC_VERSION_GE(3, 8, 0)
   SNESSetForceIteration(snes, PETSC_TRUE);
 #endif
@@ -671,6 +706,7 @@ int SNESSolver::init() {
 
 int SNESSolver::run() {
   TRACE("SNESSolver::run()");
+  int ierr;
   // Set initial guess at the solution from variables
   {
     BoutReal* xdata = nullptr;
@@ -687,7 +723,63 @@ int SNESSolver::run() {
     bool looping = true;
     int snes_failures = 0; // Count SNES convergence failures
     int saved_jacobian_lag = 0;
+    int loop_count = 0;
     do {
+      if (scale_vars) {
+        // Individual variable scaling
+        // Note: If variables are rescaled then the Jacobian columns
+        //       need to be scaled or recalculated
+
+        if (loop_count % 100 == 0) {
+          // Rescale state (snes_x) so that all quantities are around 1
+          // If quantities are near zero then RTOL is used
+          int istart, iend;
+          VecGetOwnershipRange(snes_x, &istart, &iend);
+
+          // Take ownership of snes_x and var_scaling_factors data
+          PetscScalar* snes_x_data = nullptr;
+          ierr = VecGetArray(snes_x, &snes_x_data);
+          CHKERRQ(ierr);
+          PetscScalar* x1_data;
+          ierr = VecGetArray(x1, &x1_data);
+          CHKERRQ(ierr);
+          PetscScalar* var_scaling_factors_data;
+          ierr = VecGetArray(var_scaling_factors, &var_scaling_factors_data);
+          CHKERRQ(ierr);
+
+          // Normalise each value in the state
+          // Limit normalisation so scaling factor is never smaller than rtol
+          for (int i = 0; i < iend - istart; ++i) {
+            const PetscScalar norm =
+                BOUTMAX(std::abs(snes_x_data[i]), rtol / var_scaling_factors_data[i]);
+            snes_x_data[i] /= norm;
+            x1_data[i] /= norm; // Update history for predictor
+            var_scaling_factors_data[i] *= norm;
+          }
+
+          // Restore vector underlying data
+          ierr = VecRestoreArray(var_scaling_factors, &var_scaling_factors_data);
+          CHKERRQ(ierr);
+          ierr = VecRestoreArray(x1, &x1_data);
+          CHKERRQ(ierr);
+          ierr = VecRestoreArray(snes_x, &snes_x_data);
+          CHKERRQ(ierr);
+
+          if (diagnose) {
+            // Print maximum and minimum scaling factors
+            PetscReal max_scale, min_scale;
+            VecMax(var_scaling_factors, nullptr, &max_scale);
+            VecMin(var_scaling_factors, nullptr, &min_scale);
+            output.write("Var scaling: {} -> {}\n", min_scale, max_scale);
+          }
+
+          // Force recalculation of the Jacobian
+          SNESGetLagJacobian(snes, &saved_jacobian_lag);
+          SNESSetLagJacobian(snes, 1);
+        }
+      }
+      ++loop_count;
+
       // Copy the state (snes_x) into initial values (x0)
       VecCopy(snes_x, x0);
 
@@ -729,6 +821,13 @@ int SNESSolver::run() {
       // Find out if converged
       SNESConvergedReason reason;
       SNESGetConvergedReason(snes, &reason);
+
+      // Get number of iterations
+      int nl_its;
+      SNESGetIterationNumber(snes, &nl_its);
+      int lin_its;
+      SNESGetLinearSolveIterations(snes, &lin_its);
+
       if ((ierr != 0) or (reason < 0)) {
         // Diverged or SNES failed
 
@@ -759,6 +858,19 @@ int SNESSolver::run() {
         VecCopy(x0, snes_x);
 
         // Recalculate the Jacobian
+        if (jacobian_pruned and (snes_failures > 2) and (4 * lin_its > 3 * maxl)) {
+          // Taking 3/4 of maximum linear iterations on average per linear step
+          // May indicate a preconditioner problem.
+          // Restore pruned non-zero elements
+          if (diagnose) {
+            output.write("\nRestoring Jacobian\n");
+          }
+          ierr = MatCopy(Jfd_original, Jfd, DIFFERENT_NONZERO_PATTERN);
+          CHKERRQ(ierr);
+          // The non-zero pattern has changed, so update coloring
+          updateColoring();
+          jacobian_pruned = false; // Reset flag. Will be set after pruning.
+        }
         if (saved_jacobian_lag == 0) {
           SNESGetLagJacobian(snes, &saved_jacobian_lag);
           SNESSetLagJacobian(snes, 1);
@@ -791,16 +903,24 @@ int SNESSolver::run() {
         time1 = simtime;
       }
 
-      int nl_its;
-      SNESGetIterationNumber(snes, &nl_its);
-
       if (nl_its == 0) {
         // This can occur even with SNESSetForceIteration
         // Results in simulation state freezing and rapidly going to the end
 
-        {
+        if (scale_vars) {
+          // scaled_x <- snes_x * var_scaling_factors
+          ierr = VecPointwiseMult(scaled_x, snes_x, var_scaling_factors);
+          CHKERRQ(ierr);
+
+          const BoutReal* xdata = nullptr;
+          ierr = VecGetArrayRead(scaled_x, &xdata);
+          CHKERRQ(ierr);
+          load_vars(const_cast<BoutReal*>(xdata));
+          ierr = VecRestoreArrayRead(scaled_x, &xdata);
+          CHKERRQ(ierr);
+        } else {
           const BoutReal* xdata = nullptr;
-          int ierr = VecGetArrayRead(snes_x, &xdata);
+          ierr = VecGetArrayRead(snes_x, &xdata);
           CHKERRQ(ierr);
           load_vars(const_cast<BoutReal*>(xdata));
           ierr = VecRestoreArrayRead(snes_x, &xdata);
@@ -827,9 +947,6 @@ int SNESSolver::run() {
       if (diagnose) {
         // Gather and print diagnostic information
 
-        int lin_its;
-        SNESGetLinearSolveIterations(snes, &lin_its);
-
         output.print("\r"); // Carriage return for printing to screen
         output.write("Time: {}, timestep: {}, nl iter: {}, lin iter: {}, reason: {}",
                      simtime, timestep, nl_its, lin_its, static_cast<int>(reason));
@@ -840,6 +957,51 @@ int SNESSolver::run() {
         output.write("\n");
       }
 
+#if PETSC_VERSION_GE(3, 20, 0)
+      // MatFilter and MatEliminateZeros(Mat, bool) require PETSc >= 3.20
+      if (jacobian_recalculated and prune_jacobian) {
+        jacobian_recalculated = false; // Reset flag
+
+        // Remove small elements from the Jacobian and recompute the coloring
+        // Only do this if there are a significant number of small elements.
+        int small_elements = 0;
+        int total_elements = 0;
+
+        // Get index of rows owned by this processor
+        int rstart, rend;
+        MatGetOwnershipRange(Jfd, &rstart, &rend);
+
+        PetscInt ncols;
+        const PetscScalar* vals;
+        for (int row = rstart; row < rend; row++) {
+          MatGetRow(Jfd, row, &ncols, nullptr, &vals);
+          for (int col = 0; col < ncols; col++) {
+            if (std::abs(vals[col]) < prune_abstol) {
+              ++small_elements;
+            }
+            ++total_elements;
+          }
+          MatRestoreRow(Jfd, row, &ncols, nullptr, &vals);
+        }
+
+        if (small_elements > prune_fraction * total_elements) {
+          if (diagnose) {
+            output.write("\nPruning Jacobian elements: {} / {}\n", small_elements,
+                         total_elements);
+          }
+
+          // Prune Jacobian, keeping diagonal elements
+          ierr = MatFilter(Jfd, prune_abstol, PETSC_TRUE, PETSC_TRUE);
+
+          // Update the coloring from Jfd matrix
+          updateColoring();
+
+          // Mark the Jacobian as pruned. This is so that it is only restored if pruned.
+          jacobian_pruned = true;
+        }
+      }
+#endif // PETSC_VERSION_GE(3,20,0)
+
       if (looping) {
         if (nl_its <= lower_its) {
           // Increase timestep slightly
@@ -856,7 +1018,18 @@ int SNESSolver::run() {
     } while (looping);
 
     // Put the result into variables
-    {
+    if (scale_vars) {
+      // scaled_x <- snes_x * var_scaling_factors
+      int ierr = VecPointwiseMult(scaled_x, snes_x, var_scaling_factors);
+      CHKERRQ(ierr);
+
+      const BoutReal* xdata = nullptr;
+      ierr = VecGetArrayRead(scaled_x, &xdata);
+      CHKERRQ(ierr);
+      load_vars(const_cast<BoutReal*>(xdata));
+      ierr = VecRestoreArrayRead(scaled_x, &xdata);
+      CHKERRQ(ierr);
+    } else {
       const BoutReal* xdata = nullptr;
       int ierr = VecGetArrayRead(snes_x, &xdata);
       CHKERRQ(ierr);
@@ -877,12 +1050,27 @@ int SNESSolver::run() {
 // f = rhs
 PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {
   // Get data from PETSc into BOUT++ fields
-  const BoutReal* xdata = nullptr;
-  int ierr = VecGetArrayRead(x, &xdata);
-  CHKERRQ(ierr);
-  load_vars(const_cast<BoutReal*>(xdata));
-  ierr = VecRestoreArrayRead(x, &xdata);
-  CHKERRQ(ierr);
+  if (scale_vars) {
+    // scaled_x <- x * var_scaling_factors
+    int ierr = VecPointwiseMult(scaled_x, x, var_scaling_factors);
+    CHKERRQ(ierr);
+
+    const BoutReal* xdata = nullptr;
+    ierr = VecGetArrayRead(scaled_x, &xdata);
+    CHKERRQ(ierr);
+    load_vars(const_cast<BoutReal*>(
+        xdata)); // const_cast needed due to load_vars API. Not writing to xdata.
+    ierr = VecRestoreArrayRead(scaled_x, &xdata);
+    CHKERRQ(ierr);
+  } else {
+    const BoutReal* xdata = nullptr;
+    int ierr = VecGetArrayRead(x, &xdata);
+    CHKERRQ(ierr);
+    load_vars(const_cast<BoutReal*>(
+        xdata)); // const_cast needed due to load_vars API. Not writing to xdata.
+    ierr = VecRestoreArrayRead(x, &xdata);
+    CHKERRQ(ierr);
+  }
 
   try {
     // Call RHS function
@@ -900,7 +1088,7 @@ PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {
 
   // Copy derivatives back
   BoutReal* fdata = nullptr;
-  ierr = VecGetArray(f, &fdata);
+  int ierr = VecGetArray(f, &fdata);
   CHKERRQ(ierr);
   save_derivs(fdata);
   ierr = VecRestoreArray(f, &fdata);
@@ -937,6 +1125,12 @@ PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {
   }
   };
 
+  if (scale_rhs) {
+    // f <- f * rhs_scaling_factors
+    ierr = VecPointwiseMult(f, f, rhs_scaling_factors);
+    CHKERRQ(ierr);
+  }
+
   return 0;
 }
 
@@ -983,4 +1177,133 @@ PetscErrorCode SNESSolver::precon(Vec x, Vec f) {
   return 0;
 }
 
+PetscErrorCode SNESSolver::scaleJacobian(Mat Jac_new) {
+  jacobian_recalculated = true;
+
+  if (!scale_rhs) {
+    return 0; // Not scaling the RHS values
+  }
+
+  int ierr;
+
+  // Get index of rows owned by this processor
+  int rstart, rend;
+  MatGetOwnershipRange(Jac_new, &rstart, &rend);
+
+  // Check that the vector has the same ownership range
+  int istart, iend;
+  VecGetOwnershipRange(jac_row_inv_norms, &istart, &iend);
+  if ((rstart != istart) or (rend != iend)) {
+    throw BoutException("Ownership ranges different: [{}, {}) and [{}, {})\n", rstart,
+                        rend, istart, iend);
+  }
+
+  // Calculate the norm of each row of the Jacobian
+  PetscScalar* row_inv_norm_data;
+  ierr = VecGetArray(jac_row_inv_norms, &row_inv_norm_data);
+  CHKERRQ(ierr);
+
+  PetscInt ncols;
+  const PetscScalar* vals;
+  for (int row = rstart; row < rend; ++row) {
+    MatGetRow(Jac_new, row, &ncols, nullptr, &vals);
+
+    // Calculate a norm of this row of the Jacobian
+    PetscScalar norm = 0.0;
+    for (int col = 0; col < ncols; col++) {
+      PetscScalar absval = std::abs(vals[col]);
+      if (absval > norm) {
+        norm = absval;
+      }
+      // Can we identify small elements and remove them?
+      // so we don't need to calculate them next time
+    }
+
+    // Store in the vector as 1 / norm
+    row_inv_norm_data[row - rstart] = 1. / norm;
+
+    MatRestoreRow(Jac_new, row, &ncols, nullptr, &vals);
+  }
+
+  ierr = VecRestoreArray(jac_row_inv_norms, &row_inv_norm_data);
+  CHKERRQ(ierr);
+
+  // Modify the RHS scaling: factor = factor / norm
+  ierr = VecPointwiseMult(rhs_scaling_factors, rhs_scaling_factors, jac_row_inv_norms);
+  CHKERRQ(ierr);
+
+  if (diagnose) {
+    // Print maximum and minimum scaling factors
+    PetscReal max_scale, min_scale;
+    VecMax(rhs_scaling_factors, nullptr, &max_scale);
+    VecMin(rhs_scaling_factors, nullptr, &min_scale);
+    output.write("RHS scaling: {} -> {}\n", min_scale, max_scale);
+  }
+
+  // Scale the Jacobian rows by multiplying on the left by 1/norm
+  ierr = MatDiagonalScale(Jac_new, jac_row_inv_norms, nullptr);
+  CHKERRQ(ierr);
+
+  return 0;
+}
+
+///
+/// Input Parameters:
+///   snes - nonlinear solver object
+///   x1 - location at which to evaluate Jacobian
+///   ctx - MatFDColoring context or NULL
+///
+/// Output Parameters:
+///   Jac - Jacobian matrix (not altered in this routine)
+///   Jac_new - newly computed Jacobian matrix to use with preconditioner (generally the same as
+///   Jac)
+static PetscErrorCode ComputeJacobianScaledColor(SNES snes, Vec x1, Mat Jac, Mat Jac_new,
+                                                 void* ctx) {
+  PetscErrorCode err = SNESComputeJacobianDefaultColor(snes, x1, Jac, Jac_new, ctx);
+  CHKERRQ(err);
+
+  if ((err != 0) or (ctx == nullptr)) {
+    return err;
+  }
+
+  // Get the the SNESSolver pointer from the function call context
+  SNESSolver* fctx = nullptr;
+  err = MatFDColoringGetFunction(static_cast<MatFDColoring>(ctx), nullptr,
+                                 reinterpret_cast<void**>(&fctx));
+  CHKERRQ(err);
+
+  // Call the SNESSolver function
+  return fctx->scaleJacobian(Jac_new);
+}
+
+void SNESSolver::updateColoring() {
+  // Re-calculate the coloring
+  MatColoring coloring = NULL;
+  MatColoringCreate(Jfd, &coloring);
+  MatColoringSetType(coloring, MATCOLORINGSL);
+  MatColoringSetFromOptions(coloring);
+
+  // Calculate new index sets
+  ISColoring iscoloring = NULL;
+  MatColoringApply(coloring, &iscoloring);
+  MatColoringDestroy(&coloring);
+
+  // Replace the old coloring with the new one
+  MatFDColoringDestroy(&fdcoloring);
+  MatFDColoringCreate(Jfd, iscoloring, &fdcoloring);
+  MatFDColoringSetFunction(
+      fdcoloring, reinterpret_cast<PetscErrorCode (*)()>(FormFunctionForColoring), this);
+  MatFDColoringSetFromOptions(fdcoloring);
+  MatFDColoringSetUp(Jfd, iscoloring, fdcoloring);
+  ISColoringDestroy(&iscoloring);
+
+  // Replace the CTX pointer in SNES Jacobian
+  if (matrix_free_operator) {
+    // Use matrix-free calculation for operator, finite difference for preconditioner
+    SNESSetJacobian(snes, Jmf, Jfd, ComputeJacobianScaledColor, fdcoloring);
+  } else {
+    SNESSetJacobian(snes, Jfd, Jfd, ComputeJacobianScaledColor, fdcoloring);
+  }
+}
+
 #endif // BOUT_HAS_PETSC

From aed011fd2d65db38992b3e59ac7ccf61f16e51b7 Mon Sep 17 00:00:00 2001
From: Ben Dudson <dudson2@llnl.gov>
Date: Thu, 24 Apr 2025 15:53:26 -0700
Subject: [PATCH 5/9] beuler: Limit stencil to fit in guard cells

If the stencil offset exceeds MXG in X, or MYG in Y, then
the index offsets will go outside the range of the guard cells.

Fixes failure of test-beuler
---
 src/solver/impls/snes/snes.cxx | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index 61d98adafc..024d5fbd08 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -6,6 +6,7 @@
 
 #include <bout/boutcomm.hxx>
 #include <bout/boutexception.hxx>
+#include <bout/globals.hxx>
 #include <bout/msg_stack.hxx>
 #include <bout/utils.hxx>
 
@@ -44,8 +45,11 @@ class ColoringStencil {
     };
     std::vector<std::pair<int, int>> xy_offsets;
     auto loop_bound = std::max({n_square, n_taxi, n_cross});
-    for (int i = -loop_bound; i <= loop_bound; ++i) {
-      for (int j = -loop_bound; j <= loop_bound; ++j) {
+    // Ensure that stencil does not go beyond guard cells
+    auto loop_bound_x = std::min({loop_bound, bout::globals::mesh->xstart});
+    auto loop_bound_y = std::min({loop_bound, bout::globals::mesh->ystart});
+    for (int i = -loop_bound_x; i <= loop_bound_x; ++i) {
+      for (int j = -loop_bound_y; j <= loop_bound_y; ++j) {
         if (inside(i, j)) {
           xy_offsets.emplace_back(i, j);
         }

From c94150e35d689bddf92e6f3f43162e5dcd35e161 Mon Sep 17 00:00:00 2001
From: Ben Dudson <dudson2@llnl.gov>
Date: Thu, 24 Apr 2025 16:19:38 -0700
Subject: [PATCH 6/9] beuler: Fix limits on indices

Revert earlier change to offset array. Instead fix index limits
to be < LocalNx.
---
 src/solver/impls/snes/snes.cxx | 57 ++++++++++++++++------------------
 1 file changed, 26 insertions(+), 31 deletions(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index 024d5fbd08..90e7c5085a 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -45,11 +45,8 @@ class ColoringStencil {
     };
     std::vector<std::pair<int, int>> xy_offsets;
     auto loop_bound = std::max({n_square, n_taxi, n_cross});
-    // Ensure that stencil does not go beyond guard cells
-    auto loop_bound_x = std::min({loop_bound, bout::globals::mesh->xstart});
-    auto loop_bound_y = std::min({loop_bound, bout::globals::mesh->ystart});
-    for (int i = -loop_bound_x; i <= loop_bound_x; ++i) {
-      for (int j = -loop_bound_y; j <= loop_bound_y; ++j) {
+    for (int i = -loop_bound; i <= loop_bound; ++i) {
+      for (int j = -loop_bound; j <= loop_bound; ++j) {
         if (inside(i, j)) {
           xy_offsets.emplace_back(i, j);
         }
@@ -320,8 +317,8 @@ int SNESSolver::init() {
 
       output_progress.write("Setting Jacobian matrix sizes\n");
 
-      int n2d = f2d.size();
-      int n3d = f3d.size();
+      const int n2d = f2d.size();
+      const int n3d = f3d.size();
 
       // Set size of Matrix on each processor to nlocal x nlocal
       MatCreate(BoutComm::get(), &Jfd);
@@ -359,40 +356,39 @@ int SNESSolver::init() {
 
       auto const xy_offsets = ColoringStencil::getOffsets(n_square, n_taxi, n_cross);
       {
-        //This is ugly but can't think of a better and robust way to
-        //count the non-zeros for some arbitery stencil
-        //effectivly the same loop as the one that sets the non-zeros below
+        // This is ugly but can't think of a better and robust way to
+        // count the non-zeros for some arbitrary stencil
+        // effectively the same loop as the one that sets the non-zeros below
         std::vector<std::set<int>> d_nnz_map2d(nlocal);
         std::vector<std::set<int>> o_nnz_map2d(nlocal);
         std::vector<std::set<int>> d_nnz_map3d(nlocal);
         std::vector<std::set<int>> o_nnz_map3d(nlocal);
-        //Loop over every element in 2D to count the *unique* non-zeros
+        // Loop over every element in 2D to count the *unique* non-zeros
         for (int x = mesh->xstart; x <= mesh->xend; x++) {
           for (int y = mesh->ystart; y <= mesh->yend; y++) {
 
-            int ind0 = ROUND(index(x, y, 0)) - Istart;
+            const int ind0 = ROUND(index(x, y, 0)) - Istart;
 
             // 2D fields
             for (int i = 0; i < n2d; i++) {
-              PetscInt row = ind0 + i;
-              //if (row < Istart || row >= Iend) continue;
+              const PetscInt row = ind0 + i;
               // Loop through each point in the stencil
               for (const auto& [x_off, y_off] : xy_offsets) {
-                int xi = x + x_off;
-                int yi = y + y_off;
+                const int xi = x + x_off;
+                const int yi = y + y_off;
                 if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
                     || (yi >= mesh->LocalNy)) {
                   continue;
                 }
 
-                int ind2 = ROUND(index(xi, yi, 0));
+                const int ind2 = ROUND(index(xi, yi, 0));
                 if (ind2 < 0) {
                   continue; // A boundary point
                 }
 
                 // Depends on all variables on this cell
                 for (int j = 0; j < n2d; j++) {
-                  PetscInt col = ind2 + j;
+                  const PetscInt col = ind2 + j;
                   if (col >= Istart && col < Iend) {
                     d_nnz_map2d[row].insert(col);
                   } else {
@@ -403,18 +399,17 @@ int SNESSolver::init() {
             }
             // 3D fields
             for (int z = 0; z < mesh->LocalNz; z++) {
-              int ind = ROUND(index(x, y, z)) - Istart;
+              const int ind = ROUND(index(x, y, z)) - Istart;
 
               for (int i = 0; i < n3d; i++) {
                 PetscInt row = ind + i;
-                //if (row < Istart || row >= Iend) continue;
                 if (z == 0) {
                   row += n2d;
                 }
 
                 // Depends on 2D fields
                 for (int j = 0; j < n2d; j++) {
-                  PetscInt col = ind0 + j;
+                  const PetscInt col = ind0 + j;
                   if (col >= Istart && col < Iend) {
                     d_nnz_map2d[row].insert(col);
                   } else {
@@ -424,11 +419,11 @@ int SNESSolver::init() {
 
                 // Star pattern
                 for (const auto& [x_off, y_off] : xy_offsets) {
-                  int xi = x + x_off;
-                  int yi = y + y_off;
+                  const int xi = x + x_off;
+                  const int yi = y + y_off;
 
-                  if ((xi < 0) || (yi < 0) || (xi > mesh->LocalNx)
-                      || (yi > mesh->LocalNy)) {
+                  if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
+                      || (yi >= mesh->LocalNy)) {
                     continue;
                   }
 
@@ -443,7 +438,7 @@ int SNESSolver::init() {
 
                   // 3D fields on this cell
                   for (int j = 0; j < n3d; j++) {
-                    PetscInt col = ind2 + j;
+                    const PetscInt col = ind2 + j;
                     if (col >= Istart && col < Iend) {
                       d_nnz_map3d[row].insert(col);
                     } else {
@@ -483,17 +478,17 @@ int SNESSolver::init() {
       for (int x = mesh->xstart; x <= mesh->xend; x++) {
         for (int y = mesh->ystart; y <= mesh->yend; y++) {
 
-          int ind0 = ROUND(index(x, y, 0));
+          const int ind0 = ROUND(index(x, y, 0));
 
           // 2D fields
           for (int i = 0; i < n2d; i++) {
-            PetscInt row = ind0 + i;
+            const PetscInt row = ind0 + i;
 
             // Loop through each point in the stencil
             for (const auto& [x_off, y_off] : xy_offsets) {
-              int xi = x + x_off;
-              int yi = y + y_off;
-              if ((xi < 0) || (yi < 0) || (xi > mesh->LocalNx) || (yi > mesh->LocalNy)) {
+              const int xi = x + x_off;
+              const int yi = y + y_off;
+              if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx) || (yi >= mesh->LocalNy)) {
                 continue;
               }
 

From 56298f02cc8709efe8106a98db14e335456dc2cf Mon Sep 17 00:00:00 2001
From: Ben Dudson <dudson2@llnl.gov>
Date: Thu, 24 Apr 2025 16:51:41 -0700
Subject: [PATCH 7/9] beuler: Apply clang-format

---
 src/solver/impls/snes/snes.cxx | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index 90e7c5085a..a7f2d36d8a 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -488,7 +488,8 @@ int SNESSolver::init() {
             for (const auto& [x_off, y_off] : xy_offsets) {
               const int xi = x + x_off;
               const int yi = y + y_off;
-              if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx) || (yi >= mesh->LocalNy)) {
+              if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx)
+                  || (yi >= mesh->LocalNy)) {
                 continue;
               }
 

From 348e65b9e2a09e2090cf209cc0eca07a4b1cfd40 Mon Sep 17 00:00:00 2001
From: Ben Dudson <dudson2@llnl.gov>
Date: Sat, 26 Apr 2025 14:21:43 -0700
Subject: [PATCH 8/9] beuler: Disable matrix_free_operator by default

This uses a matrix-free method to calculate Jacobian-vector products.
It seems to improve performance in 1D, but degrades performance in 2D.
Effect likely problem dependent.
---
 src/solver/impls/snes/snes.cxx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx
index a7f2d36d8a..4dce1e7086 100644
--- a/src/solver/impls/snes/snes.cxx
+++ b/src/solver/impls/snes/snes.cxx
@@ -160,7 +160,7 @@ SNESSolver::SNESSolver(Options* opts)
                       .withDefault<bool>(false)),
       matrix_free_operator((*options)["matrix_free_operator"]
                                .doc("Use matrix free Jacobian-vector operator?")
-                               .withDefault<bool>(true)),
+                               .withDefault<bool>(false)),
       lag_jacobian((*options)["lag_jacobian"]
                        .doc("Re-use the Jacobian this number of SNES iterations")
                        .withDefault(50)),

From ea7695a48de5a1268fa54780c592b01be0271705 Mon Sep 17 00:00:00 2001
From: Ben Dudson <dudson2@llnl.gov>
Date: Sat, 26 Apr 2025 22:37:52 -0700
Subject: [PATCH 9/9] snes: Add documentation to manual

Using @seimtpow's PR comments to add material to the manual.
---
 manual/sphinx/user_docs/time_integration.rst | 58 +++++++++++++++++++-
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/manual/sphinx/user_docs/time_integration.rst b/manual/sphinx/user_docs/time_integration.rst
index 7d0b56abcc..a4db823369 100644
--- a/manual/sphinx/user_docs/time_integration.rst
+++ b/manual/sphinx/user_docs/time_integration.rst
@@ -393,7 +393,9 @@ iterations within a given range.
 +---------------------------+---------------+----------------------------------------------------+
 | predictor                 | true          | Use linear predictor?                              |
 +---------------------------+---------------+----------------------------------------------------+
-| matrix_free               | false         | Use matrix free Jacobian-vector product?           |
+| matrix_free               | false         | Matrix-free preconditioning?                       |
++---------------------------+---------------+----------------------------------------------------+
+| matrix_free_operator      | false         | Use matrix free Jacobian-vector product?           |
 +---------------------------+---------------+----------------------------------------------------+
 | use_coloring              | true          | If ``matrix_free=false``, use coloring to speed up |
 |                           |               | calculation of the Jacobian elements.              |
@@ -402,11 +404,18 @@ iterations within a given range.
 +---------------------------+---------------+----------------------------------------------------+
 | kspsetinitialguessnonzero | false         | If true, Use previous solution as KSP initial      |
 +---------------------------+---------------+----------------------------------------------------+
-| use_precon                | false         | Use user-supplied preconditioner?                  |
+| use_precon                | false         | If ``matrix_free=true``, use user-supplied         |
+|                           |               | preconditioner?                                    |
 |                           |               | If false, the default PETSc preconditioner is used |
 +---------------------------+---------------+----------------------------------------------------+
 | diagnose                  | false         | Print diagnostic information every iteration       |
 +---------------------------+---------------+----------------------------------------------------+
+| stencil:cross             | 0             | If ``matrix_free=false`` and ``use_coloring=true`` |
+| stencil:square            | 0             | Set the size and shape of the Jacobian coloring    |
+| stencil:taxi              | 2             | stencil.                                           |
++---------------------------+---------------+----------------------------------------------------+
+| force_symmetric_coloring  | false         | Ensure that the Jacobian coloring is symmetric     |
++---------------------------+---------------+----------------------------------------------------+
 
 The predictor is linear extrapolation from the last two timesteps. It seems to be
 effective, but can be disabled by setting ``predictor = false``.
@@ -444,6 +453,51 @@ Preconditioner types:
    Enable with command-line args ``-pc_type hypre -pc_hypre_type euclid -pc_hypre_euclid_levels k``
    where ``k`` is the level (1-8 typically).
 
+Jacobian coloring stencil
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The stencil used to create the Jacobian colouring can be varied,
+depending on which numerical operators are in use.
+
+
+``solver:stencil:cross = N``
+e.g. for N == 2
+
+.. code-block:: bash
+
+        *
+        *
+    * * x * *
+        *
+        *
+
+
+``solver:stencil:square = N``
+e.g. for N == 2
+
+.. code-block:: bash
+
+    * * * * *
+    * * * * *
+    * * x * *
+    * * * * *
+    * * * * *
+
+``solver:stencil:taxi = N``
+e.g. for N == 2
+
+.. code-block:: bash
+
+        *
+      * * *
+    * * x * *
+      * * *
+        *
+
+Setting ``solver:force_symmetric_coloring = true``, will make sure
+that the jacobian colouring matrix is symmetric.  This will often
+include a few extra non-zeros that the stencil will miss otherwise
+
 ODE integration
 ---------------