SNES solver: Save Jacobian to file

- PetscPreconditioner saves Jacobian using PETSc binary or ASCII
  format.
- Solver saves metadata: A Jacobian global index offset field to
  the dmp files, and a JSON file with the variable information.
- SNES solver outputs Jacobian if `save_jacobian = true`.  An option
  `jacobian_export_kind` select whether the Jacobian is calculated using
  the nonlinear system being solved (that depends on timestep and
  scaling), the scaled RHS function (that depends on variable scaling),
  or the raw `rhs` function is saved.

Author: bendudson

include/bout/petsc_preconditioner.hxx

@@ -10,8 +10,11 @@
 #ifndef BOUT_PETSC_PRECONDITIONER_H
 #define BOUT_PETSC_PRECONDITIONER_H
 
+#include "bout/bout_enum_class.hxx"
 #include "bout/build_defines.hxx"
 
+BOUT_ENUM_CLASS(PetscMatrixExportFormat, binary, ascii);
+
 #if BOUT_HAS_PETSC
 
 #include "bout/petsc_interface.hxx"
@@ -22,6 +25,8 @@
 #include <petscmat.h>
 #include <petscvec.h>
 
+#include <string>
+
 class Options;
 class Field3D;
 
@@ -73,6 +78,13 @@ public:
   Mat jacobian() const { return Jfd; }
   MatFDColoring coloring() const { return fdcoloring; }
 
+  static PetscErrorCode
+  saveMatrix(Mat matrix, const std::string& filename,
+             PetscMatrixExportFormat format = PetscMatrixExportFormat::binary);
+  PetscErrorCode
+  saveMatrix(const std::string& filename,
+             PetscMatrixExportFormat format = PetscMatrixExportFormat::binary) const;
+
   void reset();
 
 private:
@@ -86,6 +98,9 @@ private:
 // unconditionally in PETSc-enabled compilation units.
 class PetscPreconditioner {
 public:
+  void saveMatrix(
+      const std::string& UNUSED(filename),
+      PetscMatrixExportFormat UNUSED(format) = PetscMatrixExportFormat::binary) const {}
   void reset() {}
 };
 

include/bout/solver.hxx

@@ -38,6 +38,7 @@
 
 #include "bout/build_defines.hxx"
 
+#include "bout/bout_enum_class.hxx"
 #include "bout/bout_types.hxx"
 #include "bout/boutexception.hxx"
 #include "bout/globals.hxx"
@@ -101,6 +102,8 @@ constexpr auto SOLVERRKGENERIC = "rkgeneric";
 enum class FieldCategories : std::uint8_t { VARS, DERIVS, MMS };
 enum class SOLVER_VAR_OP : std::uint8_t { LOAD, SET_ID, SAVE };
 
+BOUT_ENUM_CLASS(JacobianExportKind, system, scaled, rhs);
+
 /// A type to set where in the list monitors are added
 enum class MonitorPosition { BACK, FRONT };
 
@@ -368,6 +371,25 @@ public:
 protected:
   friend class SundialsNVectorInterface;
 
+  struct JacobianVariableMetadata {
+    int offset{0};
+    std::string name;
+    std::string location;
+    bool evolve_bndry{false};
+    bool constraint{false};
+    std::string description;
+  };
+
+  struct JacobianMetadata {
+    int format_version{1};
+    std::string solver_name;
+    int n2d{0};
+    int n3d{0};
+    std::vector<JacobianVariableMetadata> variables_2d;
+    std::vector<JacobianVariableMetadata> variables_3d;
+    std::string ordering;
+  };
+
   /// Number of command-line arguments
   static int* pargc;
   /// Command-line arguments
@@ -606,6 +628,12 @@ protected:
 
   /// Returns a Field3D containing the global indices
   Field3D globalIndex(int localStart);
+  Field3D jacobianIndexBase(int localStart = 0);
+  std::vector<JacobianVariableMetadata> getJacobianMetadata2D() const;
+  std::vector<JacobianVariableMetadata> getJacobianMetadata3D() const;
+  JacobianMetadata getJacobianMetadata(const std::string& solver_name) const;
+  void writeJacobianMetadataJson(const std::string& filename,
+                                 const std::string& solver_name) const;
 
   /// Maximum internal timestep
   BoutReal max_dt{-1.0};
@@ -670,6 +698,7 @@ private:
   std::string run_restart_from = "yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy";
   /// Save `run_id` and `run_restart_from` every output
   bool save_repeat_run_id{false};
+  bool save_jacobian_index_base{false};
 
   /// Current iteration (output time-step) number
   int iteration{0};

src/solver/impls/snes/snes.cxx

@@ -21,6 +21,7 @@
 #include <algorithm>
 #include <cmath>
 #include <cstddef>
+#include <fmt/format.h>
 #include <vector>
 
 #include "petscerror.h"
@@ -61,6 +62,15 @@ PetscErrorCode FormFunctionForColoring(void* UNUSED(snes), Vec x, Vec f, void* c
   return static_cast<SNESSolver*>(ctx)->snes_function(x, f, true);
 }
 
+PetscErrorCode FormRawFunctionForColoring(void* UNUSED(snes), Vec x, Vec f, void* ctx) {
+  return static_cast<SNESSolver*>(ctx)->raw_rhs_function(x, f, true);
+}
+
+PetscErrorCode FormScaledFunctionForColoring(void* UNUSED(snes), Vec x, Vec f,
+                                             void* ctx) {
+  return static_cast<SNESSolver*>(ctx)->scaled_rhs_function(x, f, true);
+}
+
 PetscErrorCode snesPCapply(PC pc, Vec x, Vec y) {
   // Get the context
   SNESSolver* s;
@@ -71,6 +81,8 @@ PetscErrorCode snesPCapply(PC pc, Vec x, Vec y) {
 
 PetscErrorCode ComputeJacobianScaledColor(SNES snes, Vec x1, Mat Jac, Mat Jac_new,
                                           void* ctx);
+PetscErrorCode ComputeJacobianDefaultMaybeExport(SNES snes, Vec x1, Mat Jac, Mat Jac_new,
+                                                 void* ctx);
 } // namespace
 
 PetscErrorCode SNESSolver::FDJinitialise() {
@@ -111,9 +123,9 @@ PetscErrorCode SNESSolver::FDJinitialise() {
         nullptr, &Jfd);
 
     if (matrix_free_operator) {
-      SNESSetJacobian(snes, Jmf, Jfd, SNESComputeJacobianDefault, this);
+      SNESSetJacobian(snes, Jmf, Jfd, ComputeJacobianDefaultMaybeExport, this);
     } else {
-      SNESSetJacobian(snes, Jfd, Jfd, SNESComputeJacobianDefault, this);
+      SNESSetJacobian(snes, Jfd, Jfd, ComputeJacobianDefaultMaybeExport, this);
     }
 
     MatSetOption(Jfd, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);
@@ -364,7 +376,99 @@ SNESSolver::SNESSolver(Options* opts)
                             .withDefault<BoutReal>(100.)),
       asinh_vars((*options)["asinh_vars"]
                      .doc("Apply asinh() to all variables?")
-                     .withDefault<bool>(false)) {}
+                     .withDefault<bool>(false)),
+      save_jacobian((*options)["save_jacobian"]
+                        .doc("Save Jacobian matrices for diagnostics?")
+                        .withDefault<bool>(false)),
+      jacobian_export_kind((*options)["jacobian_export_kind"]
+                               .doc("Which Jacobian to save: system, scaled, rhs")
+                               .withDefault(JacobianExportKind::system)),
+      jacobian_export_prefix(
+          (*options)["jacobian_export_prefix"]
+              .doc("Prefix for saved Jacobian matrix and metadata files")
+              .withDefault("jacobian")),
+      jacobian_export_format((*options)["jacobian_export_format"]
+                                 .doc("Format for saved Jacobian matrices: binary, ascii")
+                                 .withDefault(PetscMatrixExportFormat::binary)) {}
+
+std::string SNESSolver::getJacobianExportStem(JacobianExportKind kind) {
+  return fmt::format("{}_{}_{:06d}", jacobian_export_prefix, toString(kind),
+                     jacobian_export_counter++);
+}
+
+std::string SNESSolver::getJacobianMatrixFilename(const std::string& stem) const {
+  return stem
+         + (jacobian_export_format == PetscMatrixExportFormat::binary ? ".dat" : ".txt");
+}
+
+PetscErrorCode
+SNESSolver::exportMatrixAndMetadata(const PetscPreconditioner& preconditioner,
+                                    const std::string& stem) {
+  if (!jacobian_metadata_written) {
+    writeJacobianMetadataJson(jacobian_export_prefix + "_metadata.json", "snes");
+    jacobian_metadata_written = true;
+  }
+
+  PetscCall(
+      preconditioner.saveMatrix(getJacobianMatrixFilename(stem), jacobian_export_format));
+  PetscFunctionReturn(PETSC_SUCCESS);
+}
+
+PetscErrorCode SNESSolver::saveDiagnosticJacobian(JacobianExportKind kind, Vec x_solver) {
+  PetscPreconditioner diagnostic_preconditioner;
+  Field3D index = globalIndex(0);
+  PetscCall(diagnostic_preconditioner.createJacobianPattern(
+      index, *options, nlocal, n2Dvars(), n3Dvars(), BoutComm::get()));
+
+  if (kind == JacobianExportKind::rhs) {
+    PetscCall(diagnostic_preconditioner.updateColoring(FormRawFunctionForColoring, this));
+  } else {
+    PetscCall(
+        diagnostic_preconditioner.updateColoring(FormScaledFunctionForColoring, this));
+  }
+
+  Vec x_evaluate = x_solver;
+  Vec physical_x{nullptr};
+  if (kind == JacobianExportKind::rhs) {
+    PetscCall(VecDuplicate(x_solver, &physical_x));
+    PetscCall(toPhysicalState(x_solver, physical_x));
+    x_evaluate = physical_x;
+  }
+
+  Mat diagnostic_jacobian = diagnostic_preconditioner.jacobian();
+  PetscCall(MatZeroEntries(diagnostic_jacobian));
+  PetscCall(SNESComputeJacobianDefaultColor(snes, x_evaluate, diagnostic_jacobian,
+                                            diagnostic_jacobian,
+                                            diagnostic_preconditioner.coloring()));
+  PetscCall(
+      exportMatrixAndMetadata(diagnostic_preconditioner, getJacobianExportStem(kind)));
+
+  if (physical_x != nullptr) {
+    PetscCall(VecDestroy(&physical_x));
+  }
+
+  PetscFunctionReturn(PETSC_SUCCESS);
+}
+
+PetscErrorCode SNESSolver::maybeExportJacobian(Mat system_jacobian, Vec x_solver) {
+  if (!save_jacobian) {
+    PetscFunctionReturn(PETSC_SUCCESS);
+  }
+
+  if (jacobian_export_kind == JacobianExportKind::system) {
+    if (!jacobian_metadata_written) {
+      writeJacobianMetadataJson(jacobian_export_prefix + "_metadata.json", "snes");
+      jacobian_metadata_written = true;
+    }
+    PetscCall(PetscPreconditioner::saveMatrix(
+        system_jacobian,
+        getJacobianMatrixFilename(getJacobianExportStem(jacobian_export_kind)),
+        jacobian_export_format));
+    PetscFunctionReturn(PETSC_SUCCESS);
+  }
+
+  PetscFunctionReturn(saveDiagnosticJacobian(jacobian_export_kind, x_solver));
+}
 
 int SNESSolver::init() {
   Solver::init();
@@ -1128,13 +1232,13 @@ PetscErrorCode SNESSolver::updateResiduals(Vec x) {
   const BoutReal* current_residual = nullptr;
   if (diagnose) {
     // Call RHS function to get time derivatives
-    PetscCall(rhs_function(x, deriv, false));
+    PetscCall(scaled_rhs_function(x, deriv, false));
 
     // Reading the residual vectors
     PetscCall(VecGetArrayRead(deriv, &current_residual));
   } else {
     // Call RHS function to get time derivatives
-    PetscCall(rhs_function(x, snes_f, false));
+    PetscCall(scaled_rhs_function(x, snes_f, false));
 
     // Reading the residual vectors
     PetscCall(VecGetArrayRead(snes_f, &current_residual));
@@ -1415,34 +1519,34 @@ BoutReal SNESSolver::updatePseudoTimestep(BoutReal previous_timestep,
   throw BoutException("SNESSolver::updatePseudoTimestep invalid BoutPTCStrategy");
 }
 
-PetscErrorCode SNESSolver::rhs_function(Vec x, Vec f, bool linear) {
-  // Get data from PETSc into BOUT++ fields
+PetscErrorCode SNESSolver::toPhysicalState(Vec x, Vec physical_x) {
   if (scale_vars) {
-    // scaled_x <- x * var_scaling_factors
-    PetscCall(VecPointwiseMult(scaled_x, x, var_scaling_factors));
-  } else if (asinh_vars) {
-    PetscCall(VecCopy(x, scaled_x));
+    PetscCall(VecPointwiseMult(physical_x, x, var_scaling_factors));
   } else {
-    scaled_x = x;
+    PetscCall(VecCopy(x, physical_x));
   }
 
   if (asinh_vars) {
     PetscInt size;
-    PetscCall(VecGetLocalSize(scaled_x, &size));
+    PetscCall(VecGetLocalSize(physical_x, &size));
 
-    BoutReal* scaled_data = nullptr;
-    PetscCall(VecGetArray(scaled_x, &scaled_data));
+    BoutReal* physical_data = nullptr;
+    PetscCall(VecGetArray(physical_x, &physical_data));
     for (PetscInt i = 0; i != size; ++i) {
-      scaled_data[i] = asinh_scale * std::sinh(scaled_data[i]);
+      physical_data[i] = asinh_scale * std::sinh(physical_data[i]);
     }
-    PetscCall(VecRestoreArray(scaled_x, &scaled_data));
+    PetscCall(VecRestoreArray(physical_x, &physical_data));
   }
 
+  return PETSC_SUCCESS;
+}
+
+PetscErrorCode SNESSolver::raw_rhs_function(Vec x, Vec f, bool linear) {
   const BoutReal* xdata = nullptr;
-  PetscCall(VecGetArrayRead(scaled_x, &xdata));
+  PetscCall(VecGetArrayRead(x, &xdata));
   // const_cast needed due to load_vars API. Not writing to xdata.
   load_vars(const_cast<BoutReal*>(xdata));
-  PetscCall(VecRestoreArrayRead(scaled_x, &xdata));
+  PetscCall(VecRestoreArrayRead(x, &xdata));
 
   try {
     // Call RHS function
@@ -1460,6 +1564,18 @@ PetscErrorCode SNESSolver::rhs_function(Vec x, Vec f, bool linear) {
   BoutReal* fdata = nullptr;
   PetscCall(VecGetArray(f, &fdata));
   save_derivs(fdata);
+  PetscCall(VecRestoreArray(f, &fdata));
+
+  return PETSC_SUCCESS;
+}
+
+PetscErrorCode SNESSolver::scaled_rhs_function(Vec x, Vec f, bool linear) {
+  if (!scale_vars && !asinh_vars) {
+    return raw_rhs_function(x, f, linear);
+  }
+
+  PetscCall(toPhysicalState(x, scaled_x));
+  PetscCall(raw_rhs_function(scaled_x, f, linear));
 
   if (asinh_vars) {
     // Modify time-derivatives for asinh(var) using chain rule
@@ -1472,14 +1588,15 @@ PetscErrorCode SNESSolver::rhs_function(Vec x, Vec f, bool linear) {
     PetscCall(VecGetLocalSize(f, &size));
     const BoutReal* scaled_data = nullptr;
     PetscCall(VecGetArrayRead(scaled_x, &scaled_data));
+    BoutReal* fdata = nullptr;
+    PetscCall(VecGetArray(f, &fdata));
     for (PetscInt i = 0; i != size; ++i) {
       fdata[i] /= std::sqrt(SQ(scaled_data[i]) + SQ(asinh_scale));
     }
+    PetscCall(VecRestoreArray(f, &fdata));
     PetscCall(VecRestoreArrayRead(scaled_x, &scaled_data));
   }
 
-  PetscCall(VecRestoreArray(f, &fdata));
-
   if (scale_vars) {
     PetscCall(VecPointwiseDivide(f, f, var_scaling_factors));
   }
@@ -1490,7 +1607,7 @@ PetscErrorCode SNESSolver::rhs_function(Vec x, Vec f, bool linear) {
 PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) {
 
   // Call the RHS function
-  if (rhs_function(x, f, linear) != PETSC_SUCCESS) {
+  if (scaled_rhs_function(x, f, linear) != PETSC_SUCCESS) {
     // Tell SNES that the input was out of domain
     SNESSetFunctionDomainError(snes);
     // Note: Returning non-zero error here leaves vectors in locked state
@@ -1664,7 +1781,20 @@ PetscErrorCode ComputeJacobianScaledColor(SNES snes, Vec x1, Mat Jac, Mat Jac_ne
   CHKERRQ(err);
 
   // Call the SNESSolver function
-  return fctx->scaleJacobian(Jac_new);
+  PetscCall(fctx->scaleJacobian(Jac_new));
+  PetscFunctionReturn(fctx->maybeExportJacobian(Jac_new, x1));
+}
+
+PetscErrorCode ComputeJacobianDefaultMaybeExport(SNES snes, Vec x1, Mat Jac, Mat Jac_new,
+                                                 void* ctx) {
+  PetscErrorCode err = SNESComputeJacobianDefault(snes, x1, Jac, Jac_new, ctx);
+  CHKERRQ(err);
+
+  if ((err != 0) or (ctx == nullptr)) {
+    return err;
+  }
+
+  PetscFunctionReturn(static_cast<SNESSolver*>(ctx)->maybeExportJacobian(Jac_new, x1));
 }
 } // namespace