From 1a60d3fcba9e048b9a44662405fb010c36feeaef Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Mon, 17 Aug 2026 22:24:01 -0700 Subject: [PATCH 1/4] 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. --- include/bout/petsc_preconditioner.hxx | 15 +++ include/bout/solver.hxx | 29 +++++ src/solver/impls/snes/snes.cxx | 176 ++++++++++++++++++++++---- src/solver/impls/snes/snes.hxx | 31 ++++- src/solver/petsc_preconditioner.cxx | 27 ++++ src/solver/solver.cxx | 133 +++++++++++++++++++ 6 files changed, 381 insertions(+), 30 deletions(-) diff --git a/include/bout/petsc_preconditioner.hxx b/include/bout/petsc_preconditioner.hxx index be283413fc..9fe33f7330 100644 --- a/include/bout/petsc_preconditioner.hxx +++ b/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 #include +#include + 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() {} }; diff --git a/include/bout/solver.hxx b/include/bout/solver.hxx index 0d4b330ec6..10a3d294f1 100644 --- a/include/bout/solver.hxx +++ b/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 variables_2d; + std::vector 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 getJacobianMetadata2D() const; + std::vector 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}; diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index e7d367331a..8c40cf5722 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "petscerror.h" @@ -61,6 +62,15 @@ PetscErrorCode FormFunctionForColoring(void* UNUSED(snes), Vec x, Vec f, void* c return static_cast(ctx)->snes_function(x, f, true); } +PetscErrorCode FormRawFunctionForColoring(void* UNUSED(snes), Vec x, Vec f, void* ctx) { + return static_cast(ctx)->raw_rhs_function(x, f, true); +} + +PetscErrorCode FormScaledFunctionForColoring(void* UNUSED(snes), Vec x, Vec f, + void* ctx) { + return static_cast(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(100.)), asinh_vars((*options)["asinh_vars"] .doc("Apply asinh() to all variables?") - .withDefault(false)) {} + .withDefault(false)), + save_jacobian((*options)["save_jacobian"] + .doc("Save Jacobian matrices for diagnostics?") + .withDefault(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, ¤t_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, ¤t_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(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(ctx)->maybeExportJacobian(Jac_new, x1)); } } // namespace diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index c8042cf2dc..d45ab1cdaf 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -110,6 +110,18 @@ public: /// finite difference approximated Jacobian. PetscErrorCode scaleJacobian(Mat Jac_new); + /// Convert solver coordinates into the physical variables used by the model. + PetscErrorCode toPhysicalState(Vec x, Vec physical_x); + + /// Call the physics model RHS function on a vector of physical variables. + PetscErrorCode raw_rhs_function(Vec x, Vec f, bool linear); + + /// Apply solver-coordinate transforms, call the raw RHS, and transform the + /// resulting derivatives back into solver coordinates. + PetscErrorCode scaled_rhs_function(Vec x, Vec f, bool linear); + + PetscErrorCode maybeExportJacobian(Mat system_jacobian, Vec x_solver); + /// Save diagnostics to output void outputVars(Options& output_options, bool save_repeat = true) override; @@ -121,13 +133,11 @@ private: /// Rescale state (snes_x) so that all quantities are around 1. If /// quantities are near zero then RTOL is used. PetscErrorCode rescale(); - - /// Call the physics model RHS function - /// - /// @param[in] x The state vector. Will be scaled if scale_vars=true - /// @param[out] f The vector for the result f(x) - /// @param[in] linear Specifies that the SNES solver is in a linear (KSP) inner loop - PetscErrorCode rhs_function(Vec x, Vec f, bool linear); + PetscErrorCode saveDiagnosticJacobian(JacobianExportKind kind, Vec x_solver); + std::string getJacobianExportStem(JacobianExportKind kind); + std::string getJacobianMatrixFilename(const std::string& stem) const; + PetscErrorCode exportMatrixAndMetadata(const PetscPreconditioner& preconditioner, + const std::string& stem); BoutSnesOutput output_trigger; ///< Sets when outputs are written @@ -278,6 +288,13 @@ private: bool asinh_vars; ///< Evolve asinh(vars) to compress magnitudes while preserving signs const BoutReal asinh_scale = 1e-5; // Scale below which asinh response becomes ~linear + bool save_jacobian; ///< Save Jacobian matrices for diagnostics + JacobianExportKind jacobian_export_kind; ///< Which Jacobian to save + std::string jacobian_export_prefix; ///< Prefix for Jacobian matrix/metadata outputs + PetscMatrixExportFormat jacobian_export_format; ///< Output format for matrix save + int jacobian_export_counter{0}; ///< Running counter for saved Jacobians + bool jacobian_metadata_written{false}; ///< Has the shared JSON metadata been written? + std::vector resid_2d; ///< Storage for residuals of SNES solve, unpacked from snes_f std::vector diff --git a/src/solver/petsc_preconditioner.cxx b/src/solver/petsc_preconditioner.cxx index 508d001fd3..063d85812e 100644 --- a/src/solver/petsc_preconditioner.cxx +++ b/src/solver/petsc_preconditioner.cxx @@ -6,6 +6,7 @@ #include "bout/assert.hxx" #include "bout/boutcomm.hxx" +#include "bout/boutexception.hxx" #include "bout/field3d.hxx" #include "bout/globals.hxx" #include "bout/mesh.hxx" @@ -16,6 +17,7 @@ #include #include #include +#include #include #include @@ -71,6 +73,31 @@ void PetscPreconditioner::reset() { } } +PetscErrorCode PetscPreconditioner::saveMatrix(Mat matrix, const std::string& filename, + PetscMatrixExportFormat format) { + if (matrix == nullptr) { + throw BoutException("Cannot save Jacobian matrix: matrix has not been created yet"); + } + + PetscViewer viewer{nullptr}; + if (format == PetscMatrixExportFormat::binary) { + PetscCall(PetscViewerBinaryOpen(BoutComm::get(), filename.c_str(), FILE_MODE_WRITE, + &viewer)); + } else { + PetscCall(PetscViewerASCIIOpen(BoutComm::get(), filename.c_str(), &viewer)); + } + + PetscCall(MatView(matrix, viewer)); + PetscCall(PetscViewerDestroy(&viewer)); + + PetscFunctionReturn(PETSC_SUCCESS); +} + +PetscErrorCode PetscPreconditioner::saveMatrix(const std::string& filename, + PetscMatrixExportFormat format) const { + return saveMatrix(Jfd, filename, format); +} + PetscErrorCode PetscPreconditioner::createJacobianPattern(Field3D& index, Options& options, PetscInt nlocal, int n2d, diff --git a/src/solver/solver.cxx b/src/solver/solver.cxx index 1f7b654427..133451c9e7 100644 --- a/src/solver/solver.cxx +++ b/src/solver/solver.cxx @@ -51,6 +51,7 @@ #include #include +#include #include #include #include @@ -79,6 +80,44 @@ int* Solver::pargc = nullptr; char*** Solver::pargv = nullptr; +namespace { +std::string jsonEscape(const std::string& input) { + std::string escaped; + escaped.reserve(input.size()); + + for (const char ch : input) { + switch (ch) { + case '\\': + escaped += "\\\\"; + break; + case '"': + escaped += "\\\""; + break; + case '\b': + escaped += "\\b"; + break; + case '\f': + escaped += "\\f"; + break; + case '\n': + escaped += "\\n"; + break; + case '\r': + escaped += "\\r"; + break; + case '\t': + escaped += "\\t"; + break; + default: + escaped += ch; + break; + } + } + + return escaped; +} +} // namespace + /************************************************************************** * Constructor **************************************************************************/ @@ -94,6 +133,10 @@ Solver::Solver(Options* opts) "timestep, to make it easier to concatenate output " "data sets in time") .withDefault(false)), + save_jacobian_index_base( + (*options)["save_jacobian_index_base"] + .doc("Write the base global index field used for Jacobian diagnostics") + .withDefault(false)), is_nonsplit_model_diffusive( (*options)["is_nonsplit_model_diffusive"] .doc("If not a split operator, treat RHS as diffusive?") @@ -705,6 +748,11 @@ void Solver::outputVars(Options& output_options, bool save_repeat) { "or the previous run did not have a run_id.") .assignRepeat(run_restart_from, "t", save_repeat and save_repeat_run_id, "Solver"); + if (save_jacobian_index_base) { + output_options["jacobian_index_base"].assignRepeat(jacobianIndexBase(), "t", + save_repeat, "Solver"); + } + // Add 2D and 3D evolving fields to output file for (const auto& f : f2d) { // Add to dump file (appending) @@ -1207,6 +1255,91 @@ Field3D Solver::globalIndex(int localStart) { return index; } +Field3D Solver::jacobianIndexBase(int localStart) { return globalIndex(localStart); } + +std::vector Solver::getJacobianMetadata2D() const { + std::vector metadata; + metadata.reserve(f2d.size()); + + for (int i = 0; i < static_cast(f2d.size()); ++i) { + metadata.push_back(JacobianVariableMetadata{i, f2d[i].name, toString(f2d[i].location), + f2d[i].evolve_bndry, f2d[i].constraint, + f2d[i].description}); + } + + return metadata; +} + +std::vector Solver::getJacobianMetadata3D() const { + std::vector metadata; + metadata.reserve(f3d.size()); + + for (int i = 0; i < static_cast(f3d.size()); ++i) { + metadata.push_back(JacobianVariableMetadata{i, f3d[i].name, toString(f3d[i].location), + f3d[i].evolve_bndry, f3d[i].constraint, + f3d[i].description}); + } + + return metadata; +} + +Solver::JacobianMetadata +Solver::getJacobianMetadata(const std::string& solver_name) const { + return JacobianMetadata{1, + solver_name, + n2Dvars(), + n3Dvars(), + getJacobianMetadata2D(), + getJacobianMetadata3D(), + "For each (x,y): 2D variables at z=0, then 3D variables for " + "z=0..Nz-1; evolved boundary points precede RGN_NOBNDRY"}; +} + +void Solver::writeJacobianMetadataJson(const std::string& filename, + const std::string& solver_name) const { + if (MYPE != 0) { + return; + } + + const auto metadata = getJacobianMetadata(solver_name); + std::ofstream json_file(filename); + if (!json_file.is_open()) { + throw BoutException("Failed to open Jacobian metadata file '{}'", filename); + } + + auto write_variables = [&](const std::vector& variables, + const char* name) { + json_file << " \"" << name << "\": [\n"; + for (std::size_t i = 0; i < variables.size(); ++i) { + const auto& variable = variables[i]; + json_file << " {\"offset\": " << variable.offset << ", " + << "\"name\": \"" << jsonEscape(variable.name) << "\", " + << "\"location\": \"" << jsonEscape(variable.location) << "\", " + << "\"evolve_bndry\": " << (variable.evolve_bndry ? "true" : "false") + << ", " + << "\"constraint\": " << (variable.constraint ? "true" : "false") << ", " + << "\"description\": \"" << jsonEscape(variable.description) << "\"}"; + if (i + 1 != variables.size()) { + json_file << ","; + } + json_file << "\n"; + } + json_file << " ]"; + }; + + json_file << "{\n" + << " \"format_version\": " << metadata.format_version << ",\n" + << " \"solver\": \"" << jsonEscape(metadata.solver_name) << "\",\n" + << " \"n2d\": " << metadata.n2d << ",\n" + << " \"n3d\": " << metadata.n3d << ",\n"; + write_variables(metadata.variables_2d, "variables_2d"); + json_file << ",\n"; + write_variables(metadata.variables_3d, "variables_3d"); + json_file << ",\n" + << " \"ordering\": \"" << jsonEscape(metadata.ordering) << "\"\n" + << "}\n"; +} + /************************************************************************** * Running user-supplied functions **************************************************************************/ From bb0ca6266a2d78245105137c825d225efd0167b7 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Tue, 18 Aug 2026 21:28:11 -0700 Subject: [PATCH 2/4] SNES solver: Save Jacobian to data directory Saves to same directory as the dmp file outputs, using `Options::root()["datadir"]` --- src/solver/impls/snes/snes.cxx | 35 +++++++++++++++++----------------- src/solver/impls/snes/snes.hxx | 3 +-- src/solver/solver.cxx | 10 +++++++--- 3 files changed, 25 insertions(+), 23 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index 8c40cf5722..b756463358 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -392,7 +393,9 @@ SNESSolver::SNESSolver(Options* opts) .withDefault(PetscMatrixExportFormat::binary)) {} std::string SNESSolver::getJacobianExportStem(JacobianExportKind kind) { - return fmt::format("{}_{}_{:06d}", jacobian_export_prefix, toString(kind), + // The directory the output data is stored in + const std::string datadir = Options::root()["datadir"]; + return fmt::format("{}/{}_{}_{:06d}", datadir, jacobian_export_prefix, toString(kind), jacobian_export_counter++); } @@ -401,16 +404,19 @@ std::string SNESSolver::getJacobianMatrixFilename(const std::string& stem) const + (jacobian_export_format == PetscMatrixExportFormat::binary ? ".dat" : ".txt"); } -PetscErrorCode -SNESSolver::exportMatrixAndMetadata(const PetscPreconditioner& preconditioner, - const std::string& stem) { +PetscErrorCode SNESSolver::exportMatrixAndMetadata(Mat jacobian, + const std::string& stem) { if (!jacobian_metadata_written) { - writeJacobianMetadataJson(jacobian_export_prefix + "_metadata.json", "snes"); + const std::string datadir = Options::root()["datadir"]; + const std::string metadata_filename = + datadir + "/" + jacobian_export_prefix + "_metadata.json"; + output.write("Jacobian metadata written to {}\n", metadata_filename); + writeJacobianMetadataJson(metadata_filename, "snes"); jacobian_metadata_written = true; } - PetscCall( - preconditioner.saveMatrix(getJacobianMatrixFilename(stem), jacobian_export_format)); + PetscCall(PetscPreconditioner::saveMatrix(jacobian, getJacobianMatrixFilename(stem), + jacobian_export_format)); PetscFunctionReturn(PETSC_SUCCESS); } @@ -440,8 +446,7 @@ PetscErrorCode SNESSolver::saveDiagnosticJacobian(JacobianExportKind kind, Vec x PetscCall(SNESComputeJacobianDefaultColor(snes, x_evaluate, diagnostic_jacobian, diagnostic_jacobian, diagnostic_preconditioner.coloring())); - PetscCall( - exportMatrixAndMetadata(diagnostic_preconditioner, getJacobianExportStem(kind))); + PetscCall(exportMatrixAndMetadata(diagnostic_jacobian, getJacobianExportStem(kind))); if (physical_x != nullptr) { PetscCall(VecDestroy(&physical_x)); @@ -456,14 +461,8 @@ PetscErrorCode SNESSolver::maybeExportJacobian(Mat system_jacobian, Vec x_solver } 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)); + PetscCall(exportMatrixAndMetadata(system_jacobian, + getJacobianExportStem(jacobian_export_kind))); PetscFunctionReturn(PETSC_SUCCESS); } @@ -1787,7 +1786,7 @@ PetscErrorCode ComputeJacobianScaledColor(SNES snes, Vec x1, Mat Jac, Mat Jac_ne PetscErrorCode ComputeJacobianDefaultMaybeExport(SNES snes, Vec x1, Mat Jac, Mat Jac_new, void* ctx) { - PetscErrorCode err = SNESComputeJacobianDefault(snes, x1, Jac, Jac_new, ctx); + const PetscErrorCode err = SNESComputeJacobianDefault(snes, x1, Jac, Jac_new, ctx); CHKERRQ(err); if ((err != 0) or (ctx == nullptr)) { diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index d45ab1cdaf..5c866c2f3a 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -136,8 +136,7 @@ private: PetscErrorCode saveDiagnosticJacobian(JacobianExportKind kind, Vec x_solver); std::string getJacobianExportStem(JacobianExportKind kind); std::string getJacobianMatrixFilename(const std::string& stem) const; - PetscErrorCode exportMatrixAndMetadata(const PetscPreconditioner& preconditioner, - const std::string& stem); + PetscErrorCode exportMatrixAndMetadata(Mat jacobian, const std::string& stem); BoutSnesOutput output_trigger; ///< Sets when outputs are written diff --git a/src/solver/solver.cxx b/src/solver/solver.cxx index 133451c9e7..340680f498 100644 --- a/src/solver/solver.cxx +++ b/src/solver/solver.cxx @@ -57,6 +57,7 @@ #include #include #include +#include // Implementations: #include "impls/adams_bashforth/adams_bashforth.hxx" @@ -748,9 +749,12 @@ void Solver::outputVars(Options& output_options, bool save_repeat) { "or the previous run did not have a run_id.") .assignRepeat(run_restart_from, "t", save_repeat and save_repeat_run_id, "Solver"); - if (save_jacobian_index_base) { - output_options["jacobian_index_base"].assignRepeat(jacobianIndexBase(), "t", - save_repeat, "Solver"); + if (initialised and save_jacobian_index_base) { + // The Jacobian index base offsets are not time-dependent + // Only need to calculate once, but can only be calculated + // once the solver has been initialised. This outputVars + // is called once at the start before Solver is initialised. + output_options["jacobian_index_base"].force(jacobianIndexBase(), "Solver"); } // Add 2D and 3D evolving fields to output file From e932d239a3fbdb95f3ad254bfe5c2c9b6093a6fe Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Tue, 18 Aug 2026 21:30:05 -0700 Subject: [PATCH 3/4] integrated test test-snes-save-jacobian Started a test to check the output Jacobian --- tests/integrated/CMakeLists.txt | 1 + .../test-snes-save-jacobian/CMakeLists.txt | 7 +++++ .../test-snes-save-jacobian/data/BOUT.inp | 12 +++++++++ .../test_snes_save_jacobian.cxx | 27 +++++++++++++++++++ .../test_snes_save_jacobian.py | 2 ++ 5 files changed, 49 insertions(+) create mode 100644 tests/integrated/test-snes-save-jacobian/CMakeLists.txt create mode 100644 tests/integrated/test-snes-save-jacobian/data/BOUT.inp create mode 100644 tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.cxx create mode 100644 tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.py diff --git a/tests/integrated/CMakeLists.txt b/tests/integrated/CMakeLists.txt index d87bd329f0..41a21deac0 100644 --- a/tests/integrated/CMakeLists.txt +++ b/tests/integrated/CMakeLists.txt @@ -40,6 +40,7 @@ add_subdirectory(test-restarting) add_subdirectory(test-slepc-solver) add_subdirectory(test-smooth) add_subdirectory(test-snb) +add_subdirectory(test-snes-save-jacobian) add_subdirectory(test-solver) add_subdirectory(test-squash) add_subdirectory(test-stopCheck) diff --git a/tests/integrated/test-snes-save-jacobian/CMakeLists.txt b/tests/integrated/test-snes-save-jacobian/CMakeLists.txt new file mode 100644 index 0000000000..115ed8b250 --- /dev/null +++ b/tests/integrated/test-snes-save-jacobian/CMakeLists.txt @@ -0,0 +1,7 @@ +bout_add_integrated_test( + test-snes-save-jacobian + SOURCES test_snes_save_jacobian.cxx + CONFLICTS BOUT_USE_METRIC_3D + USE_DATA_BOUT_INP + REQUIRES BOUT_HAS_FFTW +) diff --git a/tests/integrated/test-snes-save-jacobian/data/BOUT.inp b/tests/integrated/test-snes-save-jacobian/data/BOUT.inp new file mode 100644 index 0000000000..03fc71159b --- /dev/null +++ b/tests/integrated/test-snes-save-jacobian/data/BOUT.inp @@ -0,0 +1,12 @@ +MXG = 0 + +[mesh] +nx = 1 +ny = 16 +nz = 1 + +[solver] +type = snes +snes_type = newtonls +save_jacobian_index_base = true +save_jacobian = true diff --git a/tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.cxx b/tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.cxx new file mode 100644 index 0000000000..7d35b949be --- /dev/null +++ b/tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.cxx @@ -0,0 +1,27 @@ +#include +#include +#include +#include + +class TestSolver : public PhysicsModel { +public: + Field3D f, g; + + int init(bool UNUSED(restarting)) override { + solver->add(f, "f"); + solver->add(g, "g"); + + f = 1.0; + g = 0.0; + + return 0; + } + + int rhs(BoutReal UNUSED(time)) override { + ddt(f) = -0.1 * f; + ddt(g) = 0.5 * f - g; + return 0; + } +}; + +BOUTMAIN(TestSolver); diff --git a/tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.py b/tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.py new file mode 100644 index 0000000000..3b6973745c --- /dev/null +++ b/tests/integrated/test-snes-save-jacobian/test_snes_save_jacobian.py @@ -0,0 +1,2 @@ +def test_runtest(assert_success_in_shell): + assert_success_in_shell("./test_snes_save_jacobian") From 70a6e70d07946e0cbce626698611995168d8862f Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Tue, 18 Aug 2026 23:01:10 -0700 Subject: [PATCH 4/4] Add PetscBinaryIO and Jacobian reader --- .../test-snes-save-jacobian/CMakeLists.txt | 3 + .../test-snes-save-jacobian/PetscBinaryIO.py | 588 ++++++++++++++++++ .../test-snes-save-jacobian/read_jacobian.py | 309 +++++++++ 3 files changed, 900 insertions(+) create mode 100755 tests/integrated/test-snes-save-jacobian/PetscBinaryIO.py create mode 100644 tests/integrated/test-snes-save-jacobian/read_jacobian.py diff --git a/tests/integrated/test-snes-save-jacobian/CMakeLists.txt b/tests/integrated/test-snes-save-jacobian/CMakeLists.txt index 115ed8b250..f75c1d7c96 100644 --- a/tests/integrated/test-snes-save-jacobian/CMakeLists.txt +++ b/tests/integrated/test-snes-save-jacobian/CMakeLists.txt @@ -5,3 +5,6 @@ bout_add_integrated_test( USE_DATA_BOUT_INP REQUIRES BOUT_HAS_FFTW ) + +configure_file(read_jacobian.py read_jacobian.py COPYONLY) +configure_file(PetscBinaryIO.py PetscBinaryIO.py COPYONLY) diff --git a/tests/integrated/test-snes-save-jacobian/PetscBinaryIO.py b/tests/integrated/test-snes-save-jacobian/PetscBinaryIO.py new file mode 100755 index 0000000000..e3b329c11a --- /dev/null +++ b/tests/integrated/test-snes-save-jacobian/PetscBinaryIO.py @@ -0,0 +1,588 @@ +"""PetscBinaryIO +=============== + +Provides + 1. PETSc-named objects Vec, Mat, and IS that inherit numpy.ndarray + 2. A class to read and write these objects from PETSc binary files. + +The standard usage of this module should look like: + + >>> import PetscBinaryIO + >>> io = PetscBinaryIO.PetscBinaryIO() + >>> objects = io.readBinaryFile('file.dat') + +or + + >>> import PetscBinaryIO + >>> import numpy + >>> vec = numpy.array([1., 2., 3.]).view(PetscBinaryIO.Vec) + >>> io = PetscBinaryIO.PetscBinaryIO() + >>> io.writeBinaryFile('file.dat', [vec,]) + +to read in objects one at a time use such as + + >>> import PetscBinaryIO + >>> io = PetscBinaryIO.PetscBinaryIO() + >>> fh = open('file.dat') + >>> objecttype = io.readObjectType(fh) + >>> if objecttype == 'Vec': + >>> v = io.readVec(fh) + + Note that one must read in the object type first and then call readVec(), readMat() etc. + + +See also PetscBinaryIO.__doc__ and methods therein. + +This module can also be used as a command line tool for converting matrices. +Run with --help to see options. +""" + +import numpy as np +import functools + +try: + basestring # Python-2 has basestring as a common parent of unicode and str +except NameError: + basestring = str # Python-3 is unicode through and through + + +def update_wrapper_with_doc(wrapper, wrapped): + """Similar to functools.update_wrapper, but also gets the wrapper's __doc__ string""" + wdoc = wrapper.__doc__ + + functools.update_wrapper(wrapper, wrapped) + if wdoc is not None: + if wrapper.__doc__ is None: + wrapper.__doc__ = wdoc + else: + wrapper.__doc__ = wrapper.__doc__ + wdoc + return wrapper + + +def wraps_with_doc(wrapped): + """Similar to functools.wraps, but also gets the wrapper's __doc__ string""" + return functools.partial(update_wrapper_with_doc, wrapped=wrapped) + + +def decorate_with_conf(f): + """Decorates methods to take kwargs for precisions.""" + + @wraps_with_doc(f) + def decorated_f(self, *args, **kwargs): + """ + Additional kwargs: + precision: 'single', 'double', '__float128' for scalars + indices: '32bit', '64-bit' integer size + complexscalars: True/False + + Note these are set in order of preference: + 1. kwargs if given here + 2. PetscBinaryIO class __init__ arguments + 3. PETSC_DIR/PETSC_ARCH defaults + """ + + changed = False + old_precision = self.precision + old_indices = self.indices + old_complexscalars = self.complexscalars + + try: + self.precision = kwargs.pop("precision") + except KeyError: + pass + else: + changed = True + + try: + self.indices = kwargs.pop("indices") + except KeyError: + pass + else: + changed = True + + try: + self.complexscalars = kwargs.pop("complexscalars") + except KeyError: + pass + else: + changed = True + + if changed: + self._update_dtypes() + + result = f(self, *args, **kwargs) + + if changed: + self.precision = old_precision + self.indices = old_indices + self.complexscalars = old_complexscalars + self._update_dtypes() + + return result + + return decorated_f + + +class DoneWithFile(Exception): + pass + + +class Vec(np.ndarray): + """Vec represented as 1D numpy array + + The best way to instantiate this class for use with writeBinaryFile() + is through the numpy view method: + + vec = numpy.array([1,2,3]).view(Vec) + """ + + _classid = 1211214 + + +class MatDense(np.matrix): + """Mat represented as 2D numpy array + + The best way to instantiate this class for use with writeBinaryFile() + is through the numpy view method: + + mat = numpy.array([[1,0],[0,1]]).view(Mat) + """ + + _classid = 1211216 + + +class MatSparse(tuple): + """Mat represented as CSR tuple ((M, N), (rowindices, col, val)) + + This should be instantiated from a tuple: + + mat = MatSparse( ((M,N), (rowindices,col,val)) ) + """ + + _classid = 1211216 + + def __repr__(self): + return "MatSparse: %s" % super(MatSparse, self).__repr__() + + +class IS(np.ndarray): + """IS represented as 1D numpy array + + The best way to instantiate this class for use with writeBinaryFile() + is through the numpy "view" method: + + an_is = numpy.array([3,4,5]).view(IS) + """ + + _classid = 1211218 + + +class PetscBinaryIO(object): + """Reader/Writer class for PETSc binary files. + + Note that by default, precisions for both scalars and indices, as well as + complex scalars, are picked up from the PETSC_DIR/PETSC_ARCH configuration + as set by environmental variables. + + Alternatively, defaults can be overridden at class instantiation, or for + a given method call. + """ + + _classid = { + 1211216: "Mat", + 1211214: "Vec", + 1211218: "IS", + 1211219: "Bag", + 1211213: "Real", + } + + def __init__(self, precision=None, indices=None, complexscalars=None): + if (precision is None) or (indices is None) or (complexscalars is None): + import petsc_conf + + defaultprecision, defaultindices, defaultcomplexscalars = ( + petsc_conf.get_conf() + ) + if precision is None: + if defaultprecision is None: + precision = "double" + else: + precision = defaultprecision + + if indices is None: + if defaultindices is None: + indices = "32bit" + else: + indices = defaultindices + + if complexscalars is None: + if defaultcomplexscalars is None: + complexscalars = False + else: + complexscalars = defaultcomplexscalars + + self.precision = precision + if self.precision == "__float128": + raise RuntimeError( + "__float128 (quadruple) precision is not properly supported. One may use double precision by using -binary_write_double in PETSc and precision='double' here" + ) + self.indices = indices + self.complexscalars = complexscalars + self._update_dtypes() + + def _update_dtypes(self): + if self.indices == "64bit": + self._inttype = np.dtype(">i8") + else: + self._inttype = np.dtype(">i4") + + if self.precision == "__float128": + nbyte = 16 + elif self.precision == "single": + nbyte = 4 + else: + nbyte = 8 + + if self.complexscalars: + name = "c" + nbyte = nbyte * 2 # complex scalar takes twice as many bytes + else: + name = "f" + + self._scalartype = ">{0}{1}".format(name, nbyte) + + @decorate_with_conf + def readReal(self, fh): + """Reads a single real from a binary file handle, must be called after readObjectType().""" + + try: + vals = np.fromfile(fh, dtype=self._scalartype, count=1) + except MemoryError: + raise IOError("Inconsistent or invalid real data in file") + if not len(vals) == 1: + raise IOError("Inconsistent or invalid real data in file") + return vals + + @decorate_with_conf + def readVec(self, fh): + """Reads a PETSc Vec from a binary file handle, must be called after readObjectType().""" + + nz = np.fromfile(fh, dtype=self._inttype, count=1)[0] + try: + vals = np.fromfile(fh, dtype=self._scalartype, count=nz) + except MemoryError: + raise IOError("Inconsistent or invalid Vec data in file") + if not len(vals) == nz: + raise IOError("Inconsistent or invalid Vec data in file") + return vals.view(Vec) + + @decorate_with_conf + def writeVec(self, fh, vec): + """Writes a PETSc Vec to a binary file handle.""" + + metadata = np.array([Vec._classid, len(vec)], dtype=self._inttype) + metadata.tofile(fh) + vec.astype(self._scalartype).tofile(fh) + return + + @decorate_with_conf + def readMatSparse(self, fh): + """Reads a PETSc Mat, returning a sparse representation of the data. Must be called after readObjectType() + + (M,N), (I,J,V) = readMatSparse(fid) + + Input: + fid : file handle to open binary file. + Output: + M,N : matrix size + I,J : arrays of row and column for each nonzero + V: nonzero value + """ + + try: + M, N, nz = np.fromfile(fh, dtype=self._inttype, count=3) + I = np.empty(M + 1, dtype=self._inttype) + I[0] = 0 + rownz = np.fromfile(fh, dtype=self._inttype, count=M) + np.cumsum(rownz, out=I[1:]) + assert I[-1] == nz + + J = np.fromfile(fh, dtype=self._inttype, count=nz) + assert len(J) == nz + V = np.fromfile(fh, dtype=self._scalartype, count=nz) + assert len(V) == nz + except (AssertionError, MemoryError, IndexError): + raise IOError("Inconsistent or invalid Mat data in file") + + return MatSparse(((M, N), (I, J, V))) + + @decorate_with_conf + def writeMatSparse(self, fh, mat): + """Writes a Mat into a PETSc binary file handle""" + + ((M, N), (I, J, V)) = mat + metadata = np.array([MatSparse._classid, M, N, I[-1]], dtype=self._inttype) + rownz = I[1:] - I[:-1] + + assert len(J.shape) == len(V.shape) == len(I.shape) == 1 + assert len(J) == len(V) == I[-1] == rownz.sum() + assert (rownz > -1).all() + + metadata.tofile(fh) + rownz.astype(self._inttype).tofile(fh) + J.astype(self._inttype).tofile(fh) + V.astype(self._scalartype).tofile(fh) + return + + @decorate_with_conf + def readMatDense(self, fh): + """Reads a PETSc Mat, returning a dense representation of the data, must be called after readObjectType()""" + + try: + M, N, nz = np.fromfile(fh, dtype=self._inttype, count=3) + I = np.empty(M + 1, dtype=self._inttype) + I[0] = 0 + rownz = np.fromfile(fh, dtype=self._inttype, count=M) + np.cumsum(rownz, out=I[1:]) + assert I[-1] == nz + + J = np.fromfile(fh, dtype=self._inttype, count=nz) + assert len(J) == nz + V = np.fromfile(fh, dtype=self._scalartype, count=nz) + assert len(V) == nz + + except (AssertionError, MemoryError, IndexError): + raise IOError("Inconsistent or invalid Mat data in file") + + mat = np.zeros((M, N), dtype=self._scalartype) + for row in range(M): + rstart, rend = I[row : row + 2] + mat[row, J[rstart:rend]] = V[rstart:rend] + return mat.view(MatDense) + + @decorate_with_conf + def readMatSciPy(self, fh): + from scipy.sparse import csr_matrix + + (M, N), (I, J, V) = self.readMatSparse(fh) + return csr_matrix((V, J, I), shape=(M, N)) + + @decorate_with_conf + def writeMatSciPy(self, fh, mat): + from scipy.sparse import csr_matrix + + if hasattr(mat, "tocsr"): + mat = mat.tocsr() + assert isinstance(mat, csr_matrix) + V = mat.data + M, N = mat.shape + J = mat.indices + I = mat.indptr + return self.writeMatSparse(fh, (mat.shape, (mat.indptr, mat.indices, mat.data))) + + @decorate_with_conf + def readMat(self, fh, mattype="sparse"): + """Reads a PETSc Mat from binary file handle, must be called after readObjectType() + + optional mattype: 'sparse" or 'dense' + + See also: readMatSparse, readMatDense + """ + + if mattype == "sparse": + return self.readMatSparse(fh) + elif mattype == "dense": + return self.readMatDense(fh) + elif mattype == "scipy.sparse": + return self.readMatSciPy(fh) + else: + raise RuntimeError( + "Invalid matrix type requested: choose sparse/dense/scipy.sparse" + ) + + @decorate_with_conf + def readIS(self, fh): + """Reads a PETSc Index Set from binary file handle, must be called after readObjectType()""" + + try: + nz = np.fromfile(fh, dtype=self._inttype, count=1)[0] + v = np.fromfile(fh, dtype=self._inttype, count=nz) + assert len(v) == nz + except (MemoryError, IndexError): + raise IOError("Inconsistent or invalid IS data in file") + return v.view(IS) + + @decorate_with_conf + def writeIS(self, fh, anis): + """Writes a PETSc IS to binary file handle.""" + + metadata = np.array([IS._classid, len(anis)], dtype=self._inttype) + metadata.tofile(fh) + anis.astype(self._inttype).tofile(fh) + return + + @decorate_with_conf + def readObjectType(self, fid): + """Returns the next object type as a string in the file""" + try: + header = np.fromfile(fid, dtype=self._inttype, count=1)[0] + except (MemoryError, IndexError): + raise DoneWithFile + try: + objecttype = self._classid[header] + except KeyError: + raise IOError( + "Invalid PetscObject CLASSID or object not implemented for python" + ) + return objecttype + + @decorate_with_conf + def readBinaryFile(self, fid, mattype="sparse"): + """Reads a PETSc binary file, returning a tuple of the contained objects. + + objects = self.readBinaryFile(fid, **kwargs) + + Input: + fid : either file name or handle to an open binary file. + + Output: + objects : tuple of objects representing the data in numpy arrays. + + Optional: + mattype : + 'sparse': Return matrices as raw CSR: (M, N), (row, col, val). + 'dense': Return matrices as MxN numpy arrays. + 'scipy.sparse': Return matrices as scipy.sparse objects. + """ + + close = False + + if isinstance(fid, basestring): + fid = open(fid, "rb") + close = True + + objects = [] + try: + while True: + objecttype = self.readObjectType(fid) + + if objecttype == "Vec": + objects.append(self.readVec(fid)) + elif objecttype == "IS": + objects.append(self.readIS(fid)) + elif objecttype == "Mat": + objects.append(self.readMat(fid, mattype)) + elif objecttype == "Real": + objects.append(self.readReal(fid)) + elif objecttype == "Bag": + raise NotImplementedError("Bag Reader not yet implemented") + except DoneWithFile: + pass + finally: + if close: + fid.close() + + return tuple(objects) + + @decorate_with_conf + def writeBinaryFile(self, fid, objects): + """Writes a PETSc binary file containing the objects given. + + readBinaryFile(fid, objects) + + Input: + fid : either file handle to an open binary file, or filename. + objects : list of objects representing the data in numpy arrays, + which must be of type Vec, IS, MatSparse, or MatSciPy. + """ + close = False + if isinstance(fid, basestring): + fid = open(fid, "wb") + close = True + + for petscobj in objects: + if isinstance(petscobj, Vec): + self.writeVec(fid, petscobj) + elif isinstance(petscobj, IS): + self.writeIS(fid, petscobj) + elif isinstance(petscobj, MatSparse): + self.writeMatSparse(fid, petscobj) + elif isinstance(petscobj, MatDense): + if close: + fid.close() + raise NotImplementedError("Writing a dense matrix is not yet supported") + else: + try: + self.writeMatSciPy(fid, petscobj) + except AssertionError: + if close: + fid.close() + raise TypeError( + "Object %s is not a valid PETSc object" % (petscobj.__repr__()) + ) + if close: + fid.close() + return + + +def _convert(infile, outfile, args): + ext = os.path.splitext(infile)[1] + if ext in (".mtx", ".npz"): + if ext == ".mtx": + import scipy.io + + mat = scipy.io.mmread(infile) + else: + import scipy.sparse + + mat = scipy.sparse.load_npz(infile) + if args.symmetrize: + mat = (mat + mat.T) / 2 + with open(outfile, "wb") as fd: + PetscBinaryIO().writeMatSciPy( + fd, + mat, + precision=args.precision, + complexscalars=args.complex, + indices=args.indices, + ) + else: + print("Unknown format: {}".format(infile)) + exit(1) + + +if __name__ == "__main__": + import argparse + import os + + parser = argparse.ArgumentParser("PetscBinaryIO") + subparsers = parser.add_subparsers(title="commands") + convert = subparsers.add_parser( + "convert", help="convert matrices to/from PETSc format" + ) + convert.add_argument("infile", help="file to convert") + convert.add_argument( + "-o", "--outfile", help="name of output file (defaults to {infile}.petsc)" + ) + convert.add_argument( + "--symmetrize", + help="Symmetrize (A+A^T)/2 during conversion", + action="store_true", + ) + convert.add_argument( + "--precision", + help="Precision of scalar values", + choices=["single", "double", "__float128"], + default="double", + ) + convert.add_argument("--complex", help="Use complex scalars", action="store_true") + convert.add_argument( + "--indices", + help="Integer size for indices", + choices=["32bit", "64bit"], + default="32bit", + ) + args = parser.parse_args() + if args.outfile is None: + args.outfile = os.path.splitext(args.infile)[0] + ".petsc" + _convert(args.infile, args.outfile, args) diff --git a/tests/integrated/test-snes-save-jacobian/read_jacobian.py b/tests/integrated/test-snes-save-jacobian/read_jacobian.py new file mode 100644 index 0000000000..be3e6fc514 --- /dev/null +++ b/tests/integrated/test-snes-save-jacobian/read_jacobian.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python3 + +"""Read Jacobian diagnostics written by the SNES save-jacobian test. + +This helper is intentionally lightweight: + +- PETSc binary matrices are read using the bundled ``PetscBinaryIO.py`` +- ``jacobian_index_base`` is read with ``boutdata.collect`` so MPI-written dump + files can be reconstructed in serial +- NumPy is the only required array dependency +- Pandas and SciPy are optional convenience layers +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import numpy as np + + +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parents[2] + +# Allow this script to run directly from the test directory or build directory. +for extra_path in ( + SCRIPT_DIR, + REPO_ROOT / "tools" / "pylib", + REPO_ROOT / "build" / "tools" / "pylib", +): + if extra_path.exists(): + sys.path.insert(0, str(extra_path)) + +import PetscBinaryIO # noqa: E402 + + +RawPetscSparse = tuple[tuple[int, int], tuple[np.ndarray, np.ndarray, np.ndarray]] + + +def load_metadata(path: str | Path) -> dict[str, Any]: + """Load Jacobian variable metadata from JSON.""" + + with Path(path).open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def load_index_base(data_dir: str | Path) -> np.ndarray: + """Load ``jacobian_index_base`` from BOUT++ dump files using ``boutdata``.""" + + try: + from boutdata.collect import collect + except ImportError as exc: + raise ImportError( + "Could not import boutdata.collect. Make sure tools/pylib and " + "build/tools/pylib are on PYTHONPATH, and that the Python " + "NetCDF dependencies required by boutdata (for example " + "'netCDF4') are installed." + ) from exc + + index_base = np.asarray( + collect("jacobian_index_base", path=str(data_dir), info=False, xguards=True) + ) + + if index_base.ndim == 2: + return index_base[:, :, np.newaxis] + if index_base.ndim != 3: + raise ValueError( + f"Expected jacobian_index_base to have 2 or 3 dimensions, got {index_base.ndim}" + ) + return index_base + + +def _find_matrix_file(data_dir: Path) -> Path: + matrix_files = sorted(data_dir.glob("jacobian_*.dat")) + if not matrix_files: + raise FileNotFoundError(f"No PETSc binary matrix files found in {data_dir}") + if len(matrix_files) > 1: + raise ValueError( + f"Found multiple Jacobian matrix files in {data_dir}; specify one explicitly" + ) + return matrix_files[0] + + +def load_petsc_matrix( + path: str | Path, matrix_format: str = "dense" +) -> np.ndarray | RawPetscSparse: + """Load a PETSc binary matrix using the bundled ``PetscBinaryIO`` helper. + + Parameters + ---------- + path: + Path to a PETSc binary matrix file. + matrix_format: + Either ``'dense'`` or ``'sparse'``. + """ + + if matrix_format not in {"dense", "sparse"}: + raise ValueError( + f"matrix_format must be 'dense' or 'sparse', got {matrix_format!r}" + ) + + io = PetscBinaryIO.PetscBinaryIO( + precision="double", indices="32bit", complexscalars=False + ) + objects = io.readBinaryFile(str(path), mattype=matrix_format) + if not objects: + raise ValueError(f"No PETSc objects found in {path}") + + matrix = objects[0] + if matrix_format == "dense": + return np.asarray(matrix) + return matrix + + +def sparse_to_dense(matrix: RawPetscSparse) -> np.ndarray: + """Convert a raw PETSc CSR tuple into a dense NumPy array.""" + + (nrows, ncols), (indptr, indices, values) = matrix + dense = np.zeros((nrows, ncols), dtype=np.asarray(values).dtype) + for row in range(nrows): + start = int(indptr[row]) + end = int(indptr[row + 1]) + dense[row, np.asarray(indices[start:end], dtype=int)] = values[start:end] + return dense + + +def to_numpy_dense(matrix: np.ndarray | RawPetscSparse) -> np.ndarray: + """Return a dense NumPy matrix regardless of the original representation.""" + + if isinstance(matrix, np.ndarray): + return np.asarray(matrix) + return sparse_to_dense(matrix) + + +def to_scipy_csr(matrix: np.ndarray | RawPetscSparse): + """Convert a PETSc sparse matrix representation to ``scipy.sparse.csr_matrix``.""" + + try: + from scipy.sparse import csr_matrix + except ImportError as exc: + raise ImportError("SciPy is required for CSR conversion") from exc + + if isinstance(matrix, np.ndarray): + return csr_matrix(matrix) + + (nrows, ncols), (indptr, indices, values) = matrix + return csr_matrix((values, indices, indptr), shape=(nrows, ncols)) + + +def _make_label(name: str, x: int, y: int, z: int) -> str: + return f"{name}[x={x},y={y},z={z}]" + + +def build_dof_table( + index_base: np.ndarray, metadata: dict[str, Any] +) -> list[dict[str, Any]]: + """Expand ``jacobian_index_base`` plus variable metadata into one record per DOF.""" + + records: list[dict[str, Any]] = [] + variables_2d = metadata.get("variables_2d", []) + variables_3d = metadata.get("variables_3d", []) + n2d = int(metadata.get("n2d", len(variables_2d))) + + nx, ny, nz = index_base.shape + for x in range(nx): + for y in range(ny): + for z in range(nz): + base = int(round(float(index_base[x, y, z]))) + if base < 0: + continue + + if z == 0: + for variable in variables_2d: + global_index = base + int(variable["offset"]) + records.append( + { + "global_index": global_index, + "name": variable["name"], + "field_rank": "2d", + "offset": int(variable["offset"]), + "x": x, + "y": y, + "z": 0, + "label": _make_label(variable["name"], x, y, 0), + } + ) + base_3d = base + n2d + else: + base_3d = base + + for variable in variables_3d: + global_index = base_3d + int(variable["offset"]) + records.append( + { + "global_index": global_index, + "name": variable["name"], + "field_rank": "3d", + "offset": int(variable["offset"]), + "x": x, + "y": y, + "z": z, + "label": _make_label(variable["name"], x, y, z), + } + ) + + records.sort(key=lambda record: record["global_index"]) + return records + + +def load_jacobian( + data_dir: str | Path = "data", + matrix_filename: str | Path | None = None, + matrix_format: str = "dense", +) -> tuple[np.ndarray | RawPetscSparse, list[dict[str, Any]], list[str]]: + """Load the Jacobian matrix and its expanded row/column metadata. + + Returns + ------- + matrix + Dense ``np.ndarray`` if ``matrix_format='dense'``, otherwise the raw PETSc CSR + tuple ``((M, N), (indptr, indices, values))``. + dofs + One dict per matrix row/column, sorted by ``global_index``. + labels + Convenience list of labels, ordered to match rows and columns of ``matrix``. + """ + + data_path = Path(data_dir) + metadata = load_metadata(data_path / "jacobian_metadata.json") + index_base = load_index_base(data_path) + + matrix_path = ( + Path(matrix_filename) + if matrix_filename is not None + else _find_matrix_file(data_path) + ) + matrix = load_petsc_matrix(matrix_path, matrix_format=matrix_format) + dofs = build_dof_table(index_base, metadata) + labels = [record["label"] for record in dofs] + return matrix, dofs, labels + + +def to_pandas(matrix: np.ndarray | RawPetscSparse, dofs: list[dict[str, Any]]): + """Return ``(matrix_df, dof_df)`` using Pandas. + + Pandas is imported lazily so the core reader still works when Pandas is not installed. + """ + + try: + import pandas as pd + except ImportError as exc: + raise ImportError("Pandas is required for the DataFrame adapter") from exc + + dense_matrix = to_numpy_dense(matrix) + labels = [record["label"] for record in dofs] + matrix_df = pd.DataFrame(dense_matrix, index=labels, columns=labels) + dof_df = pd.DataFrame(dofs) + return matrix_df, dof_df + + +def _main() -> None: + parser = argparse.ArgumentParser(description="Read SNES Jacobian test output") + parser.add_argument( + "--data-dir", default="data", help="Directory containing Jacobian outputs" + ) + parser.add_argument( + "--matrix-format", + choices=("dense", "sparse"), + default="dense", + help="Read the PETSc matrix as a dense NumPy array or raw sparse CSR tuple", + ) + parser.add_argument( + "--matrix-file", + default=None, + help="Explicit path to the PETSc binary matrix file. Defaults to the only jacobian_*.dat file", + ) + parser.add_argument( + "--print-matrix", + action="store_true", + help="Print the dense Jacobian matrix after loading", + ) + args = parser.parse_args() + + matrix, dofs, labels = load_jacobian( + data_dir=args.data_dir, + matrix_filename=args.matrix_file, + matrix_format=args.matrix_format, + ) + + if isinstance(matrix, np.ndarray): + shape = matrix.shape + else: + shape = matrix[0] + + print(f"Loaded Jacobian with shape {shape}") + print(f"Loaded {len(dofs)} row/column labels") + print("First few DOFs:") + for record in dofs[: min(10, len(dofs))]: + print(f" {record['global_index']:>4}: {record['label']}") + + if args.print_matrix: + print(to_numpy_dense(matrix)) + + +if __name__ == "__main__": + _main()