From 1126b28121275240d43ab754d221eed7cc207456 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Tue, 22 Oct 2024 11:57:46 -0700 Subject: [PATCH 1/9] beuler solver: scale_rhs option Scales the RHS (time derivatives) in the Jacobian calculation. Speeds up the 1D-recycling example by about x2. --- src/solver/impls/snes/snes.cxx | 147 +++++++++++++++++++++++++++++---- src/solver/impls/snes/snes.hxx | 10 +++ 2 files changed, 143 insertions(+), 14 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index 2143aad96a..3fff616e34 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -57,6 +57,35 @@ static PetscErrorCode snesPCapply(PC pc, Vec x, Vec y) { PetscFunctionReturn(s->precon(x, y)); } +/// +/// Input Parameters: +/// snes - nonlinear solver object +/// x1 - location at which to evaluate Jacobian +/// ctx - MatFDColoring context or NULL +/// +/// Output Parameters: +/// J - Jacobian matrix (not altered in this routine) +/// B - newly computed Jacobian matrix to use with preconditioner (generally the same as +/// J) +PetscErrorCode SNESComputeJacobianScaledColor(SNES snes, Vec x1, Mat J, Mat B, + void* ctx) { + PetscErrorCode err = SNESComputeJacobianDefaultColor(snes, x1, J, B, ctx); + CHKERRQ(err); + + if ((err != 0) or (ctx == nullptr)) { + return err; + } + + // Get the the SNESSolver pointer from the function call context + SNESSolver* fctx; + err = MatFDColoringGetFunction(static_cast(ctx), nullptr, + reinterpret_cast(&fctx)); + CHKERRQ(err); + + // Call the SNESSolver function + return fctx->scaleJacobian(B); +} + SNESSolver::SNESSolver(Options* opts) : Solver(opts), timestep( @@ -123,7 +152,10 @@ SNESSolver::SNESSolver(Options* opts) .withDefault(50)), use_coloring((*options)["use_coloring"] .doc("Use matrix coloring to calculate Jacobian?") - .withDefault(true)) {} + .withDefault(true)), + scale_rhs((*options)["scale_rhs"] + .doc("Scale time derivatives?") + .withDefault(false)) {} int SNESSolver::init() { @@ -163,12 +195,26 @@ 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); } // Nonlinear solver interface (SNES) @@ -227,17 +273,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 + // Set size of Matrix on each processor to nlocal x nlocal MatCreate(BoutComm::get(), &Jmf); - MatSetSizes(Jmf, localN, localN, PETSC_DETERMINE, PETSC_DETERMINE); + MatSetSizes(Jmf, nlocal, nlocal, PETSC_DETERMINE, PETSC_DETERMINE); MatSetFromOptions(Jmf); - std::vector d_nnz(localN); - std::vector o_nnz(localN); + std::vector d_nnz(nlocal); + std::vector o_nnz(nlocal); // Set values for most points const int ncells_x = (mesh->LocalNx > 1) ? 2 : 0; @@ -260,7 +305,7 @@ int SNESSolver::init() { } output_info.write("Star pattern: {} non-zero entries\n", star_pattern); - for (int i = 0; i < localN; i++) { + for (int i = 0; i < nlocal; i++) { // Non-zero elements on this processor d_nnz[i] = star_pattern; // Non-zero elements on neighboring processor @@ -274,7 +319,7 @@ int SNESSolver::init() { for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { int localIndex = ROUND(index(mesh->xstart, y, z)); - ASSERT2((localIndex >= 0) && (localIndex < localN)); + 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); @@ -286,7 +331,7 @@ int SNESSolver::init() { for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { int localIndex = ROUND(index(mesh->xstart, y, z)); - ASSERT2((localIndex >= 0) && (localIndex < localN)); + 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); @@ -300,7 +345,7 @@ int SNESSolver::init() { for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { int localIndex = ROUND(index(mesh->xend, y, z)); - ASSERT2((localIndex >= 0) && (localIndex < localN)); + 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); @@ -312,7 +357,7 @@ int SNESSolver::init() { for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { int localIndex = ROUND(index(mesh->xend, y, z)); - ASSERT2((localIndex >= 0) && (localIndex < localN)); + 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); @@ -591,7 +636,7 @@ int SNESSolver::init() { MatFDColoringSetUp(Jmf, iscoloring, fdcoloring); ISColoringDestroy(&iscoloring); - SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianDefaultColor, fdcoloring); + SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianScaledColor, fdcoloring); } else { // Brute force calculation @@ -971,6 +1016,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; } @@ -1017,4 +1068,72 @@ PetscErrorCode SNESSolver::precon(Vec x, Vec f) { return 0; } +PetscErrorCode SNESSolver::scaleJacobian(Mat B) { + if (!scale_rhs) { + return 0; // Not scaling the RHS values + } + + int ierr; + + // Get index of rows owned by this processor + int rstart, rend; + MatGetOwnershipRange(B, &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(B, 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(B, 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(B, jac_row_inv_norms, nullptr); + CHKERRQ(ierr); + + return 0; +} + #endif // BOUT_HAS_PETSC diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index 601eaaaa25..98a9dd2db0 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -82,6 +82,12 @@ public: /// @param[out] f The result of the operation PetscErrorCode precon(Vec x, Vec f); + /// Scale an approximate Jacobian, + /// and update the internal RHS scaling factors + /// This is called by SNESComputeJacobianScaledColor with the + /// finite difference approximated Jacobian. + PetscErrorCode scaleJacobian(Mat B); + private: BoutReal timestep; ///< Internal timestep BoutReal dt; ///< Current timestep used in snes_function @@ -130,6 +136,10 @@ private: bool matrix_free; ///< Use matrix free Jacobian int lag_jacobian; ///< Re-use Jacobian bool use_coloring; ///< Use matrix coloring + + bool scale_rhs; ///< Scale time derivatives? + Vec rhs_scaling_factors; ///< Factors to multiply RHS function + Vec jac_row_inv_norms; ///< 1 / Norm of the rows of the Jacobian }; #else From 422017533ecf431985ae8428f702c966b288114d Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 23 Oct 2024 11:55:49 -0700 Subject: [PATCH 2/9] beuler: Add variable scaling (Jacobian column) Uses the magnitude of each variable to update its normalisation. Doesn't improve convergence for 1D-recycling test case (makes worse) --- src/solver/impls/snes/snes.cxx | 126 +++++++++++++++++++++++++++++---- src/solver/impls/snes/snes.hxx | 5 ++ 2 files changed, 119 insertions(+), 12 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index 3fff616e34..cce20180e2 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -154,8 +154,12 @@ SNESSolver::SNESSolver(Options* opts) .doc("Use matrix coloring to calculate Jacobian?") .withDefault(true)), scale_rhs((*options)["scale_rhs"] - .doc("Scale time derivatives?") - .withDefault(false)) {} + .doc("Scale time derivatives (Jacobian row scaling)?") + .withDefault(false)), + scale_vars((*options)["scale_vars"] + .doc("Scale variables (Jacobian column scaling)?") + .withDefault(false)) {} + int SNESSolver::init() { @@ -217,6 +221,18 @@ int SNESSolver::init() { 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) output_info.write("Create SNES\n"); SNESCreate(BoutComm::get(), &snes); @@ -750,6 +766,7 @@ int SNESSolver::init() { int SNESSolver::run() { TRACE("SNESSolver::run()"); + int ierr; // Set initial guess at the solution from variables { BoutReal* xdata = nullptr; @@ -766,7 +783,56 @@ 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; + 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); @@ -877,9 +943,20 @@ int SNESSolver::run() { // 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; - int ierr = VecGetArrayRead(snes_x, &xdata); + ierr = VecGetArrayRead(scaled_x, &xdata); + CHKERRQ(ierr); + load_vars(const_cast(xdata)); + ierr = VecRestoreArrayRead(scaled_x, &xdata); + CHKERRQ(ierr); + } else { + const BoutReal* xdata = nullptr; + ierr = VecGetArrayRead(snes_x, &xdata); CHKERRQ(ierr); load_vars(const_cast(xdata)); ierr = VecRestoreArrayRead(snes_x, &xdata); @@ -935,7 +1012,19 @@ 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(xdata)); + ierr = VecRestoreArrayRead(scaled_x, &xdata); + CHKERRQ(ierr); + } else { const BoutReal* xdata = nullptr; int ierr = VecGetArrayRead(snes_x, &xdata); CHKERRQ(ierr); @@ -956,12 +1045,25 @@ 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(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(xdata)); + ierr = VecRestoreArrayRead(scaled_x, &xdata); + CHKERRQ(ierr); + } else { + const BoutReal* xdata = nullptr; + int ierr = VecGetArrayRead(x, &xdata); + CHKERRQ(ierr); + load_vars(const_cast(xdata)); + ierr = VecRestoreArrayRead(x, &xdata); + CHKERRQ(ierr); + } try { // Call RHS function @@ -979,7 +1081,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); diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index 98a9dd2db0..07807c52b6 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -140,6 +140,11 @@ private: bool scale_rhs; ///< Scale time derivatives? Vec rhs_scaling_factors; ///< Factors to multiply RHS function Vec jac_row_inv_norms; ///< 1 / Norm of the rows of the Jacobian + + bool scale_vars; ///< Scale individual variables? + Vec var_scaling_factors; ///< Factors to multiply variables when passing to user + Vec scaled_x; ///< The values passed to the user RHS + }; #else From 743ca86a8cd021e7693a2038b5bbebef2bf054da Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 23 Oct 2024 11:57:54 -0700 Subject: [PATCH 3/9] beuler: Add prune_jacobian option When using coloring, identify near-zero values and remove them from the Jacobian. This could improve efficiency by reducing the effort needed to recompute the Jacobian and preconditioner. Currently over-prunes based on initial conditions, resulting in a worse preconditioner. Need to occasionally reset and add back in the removed elements, maybe following a convergence failure. --- src/solver/impls/snes/snes.cxx | 81 +++++++++++++++++++++++++++++++++- src/solver/impls/snes/snes.hxx | 7 ++- 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index cce20180e2..f19f00902c 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -16,6 +16,8 @@ #include #include "petscsnes.h" +#include "petscmat.h" + /* * PETSc callback function, which evaluates the nonlinear * function to be solved by SNES. @@ -153,6 +155,16 @@ SNESSolver::SNESSolver(Options* opts) use_coloring((*options)["use_coloring"] .doc("Use matrix coloring to calculate Jacobian?") .withDefault(true)), + jacobian_recalculated(false), + prune_jacobian((*options)["prune_jacobian"] + .doc("Remove small elements in the Jacobian?") + .withDefault(false)), + prune_abstol((*options)["prune_abstol"] + .doc("Prune values with absolute values smaller than this") + .withDefault(1e-16)), + prune_fraction((*options)["prune_fraction"] + .doc("Prune if fraction of small elements is larger than this") + .withDefault(0.2)), scale_rhs((*options)["scale_rhs"] .doc("Scale time derivatives (Jacobian row scaling)?") .withDefault(false)), @@ -673,6 +685,8 @@ int SNESSolver::init() { } // 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); @@ -996,6 +1010,69 @@ 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(Jmf, &rstart, &rend); + + PetscInt ncols; + const PetscScalar* vals; + for (int row = rstart; row < rend; row++) { + MatGetRow(Jmf, row, &ncols, nullptr, &vals); + for (int col = 0; col < ncols; col++) { + if (std::abs(vals[col]) < prune_abstol) { + ++small_elements; + } + ++total_elements; + } + MatRestoreRow(Jmf, row, &ncols, nullptr, &vals); + } + + if (small_elements > prune_fraction * total_elements) { + output.write("Pruning Jacobian elements: {} / {}\n", + small_elements, total_elements); + + // Prune Jacobian, keeping diagonal elements + //ierr = MatEliminateZeros(Jmf, PETSC_TRUE); CHKERRQ(ierr); + ierr = MatFilter(Jmf, prune_abstol, PETSC_TRUE, PETSC_TRUE); + + // Re-calculate the coloring + MatColoring coloring = NULL; + MatColoringCreate(Jmf, &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(Jmf, iscoloring, &fdcoloring); + MatFDColoringSetFunction(fdcoloring, + reinterpret_cast(FormFunctionForColoring), + this); + MatFDColoringSetFromOptions(fdcoloring); + MatFDColoringSetUp(Jmf, iscoloring, fdcoloring); + ISColoringDestroy(&iscoloring); + + // Replace the CTX pointer in SNES Jacobian + SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianScaledColor, fdcoloring); + } + } +#endif // PETSC_VERSION_GE(3,20,0) + if (looping) { if (nl_its <= lower_its) { // Increase timestep slightly @@ -1171,6 +1248,8 @@ PetscErrorCode SNESSolver::precon(Vec x, Vec f) { } PetscErrorCode SNESSolver::scaleJacobian(Mat B) { + jacobian_recalculated = true; + if (!scale_rhs) { return 0; // Not scaling the RHS values } @@ -1196,7 +1275,7 @@ PetscErrorCode SNESSolver::scaleJacobian(Mat B) { PetscInt ncols; const PetscScalar* vals; - for (int row = rstart; row < rend; row++) { + for (int row = rstart; row < rend; ++row) { MatGetRow(B, row, &ncols, nullptr, &vals); // Calculate a norm of this row of the Jacobian diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index 07807c52b6..d262d39db9 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -122,7 +122,7 @@ private: BoutReal time1{-1.0}; ///< Time of previous solution SNES snes; ///< SNES context - Mat Jmf; ///< Matrix-free Jacobian + Mat Jmf; ///< Jacobian MatFDColoring fdcoloring; ///< Matrix coloring context, used for finite difference ///< Jacobian evaluation @@ -137,6 +137,11 @@ private: int lag_jacobian; ///< Re-use Jacobian bool use_coloring; ///< Use matrix coloring + bool jacobian_recalculated; ///< Flag set when Jacobian is recalculated + bool prune_jacobian; ///< Remove small elements in the Jacobian? + BoutReal prune_abstol; ///< Prune values with absolute values smaller than this + BoutReal prune_fraction; ///< Prune if fraction of small elements is larger than this + bool scale_rhs; ///< Scale time derivatives? Vec rhs_scaling_factors; ///< Factors to multiply RHS function Vec jac_row_inv_norms; ///< 1 / Norm of the rows of the Jacobian From 1df2f3467a4c2f14fc214194fff805e0cd7b6848 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 23 Oct 2024 19:34:11 -0700 Subject: [PATCH 4/9] beuler: Add matrix_free_operator flag Works well, true by default. Configures SNES so that a matrix-free approximation is used for the Jacobian-vector operator. The finite difference Jacobian is used to construct the preconditioner. --- src/solver/impls/snes/snes.cxx | 191 ++++++++++++++++++--------------- src/solver/impls/snes/snes.hxx | 15 ++- 2 files changed, 117 insertions(+), 89 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index f19f00902c..bc82556466 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -149,6 +149,9 @@ SNESSolver::SNESSolver(Options* opts) matrix_free((*options)["matrix_free"] .doc("Use matrix free Jacobian?") .withDefault(false)), + matrix_free_operator((*options)["matrix_free_operator"] + .doc("Use matrix free Jacobian-vector operator?") + .withDefault(true)), lag_jacobian((*options)["lag_jacobian"] .doc("Re-use the Jacobian this number of SNES iterations") .withDefault(50)), @@ -262,7 +265,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. @@ -278,12 +281,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 @@ -305,9 +313,9 @@ int SNESSolver::init() { int n3d = f3d.size(); // Set size of Matrix on each processor to nlocal x nlocal - MatCreate(BoutComm::get(), &Jmf); - MatSetSizes(Jmf, nlocal, nlocal, PETSC_DETERMINE, PETSC_DETERMINE); - MatSetFromOptions(Jmf); + MatCreate(BoutComm::get(), &Jfd); + MatSetSizes(Jfd, nlocal, nlocal, PETSC_DETERMINE, PETSC_DETERMINE); + MatSetFromOptions(Jfd); std::vector d_nnz(nlocal); std::vector o_nnz(nlocal); @@ -495,14 +503,14 @@ 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); + 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); // 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 @@ -548,7 +556,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); } } @@ -568,7 +576,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); } @@ -594,7 +602,7 @@ int SNESSolver::init() { // 3D fields on this cell for (int j = 0; j < n3d; 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); if (ierr != 0) { output.write("ERROR: {} : ({}, {}) -> ({}, {}) : {} -> {}\n", row, x, y, xi, yi, ind2, ind2 + n3d - 1); @@ -615,7 +623,7 @@ int SNESSolver::init() { } for (int j = 0; j < n3d; 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); } @@ -626,7 +634,7 @@ int SNESSolver::init() { } for (int j = 0; j < n3d; 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); } } @@ -639,34 +647,21 @@ 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); 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(FormFunctionForColoring), - this); - MatFDColoringSetFromOptions(fdcoloring); - MatFDColoringSetUp(Jmf, iscoloring, fdcoloring); - ISColoringDestroy(&iscoloring); - - SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianScaledColor, 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 @@ -674,14 +669,15 @@ 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 @@ -888,6 +884,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 @@ -918,6 +921,18 @@ 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); @@ -950,9 +965,6 @@ 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 @@ -997,9 +1009,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(reason)); @@ -1012,7 +1021,6 @@ int SNESSolver::run() { #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 @@ -1023,52 +1031,36 @@ int SNESSolver::run() { // Get index of rows owned by this processor int rstart, rend; - MatGetOwnershipRange(Jmf, &rstart, &rend); + MatGetOwnershipRange(Jfd, &rstart, &rend); PetscInt ncols; const PetscScalar* vals; for (int row = rstart; row < rend; row++) { - MatGetRow(Jmf, row, &ncols, nullptr, &vals); + 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(Jmf, row, &ncols, nullptr, &vals); + MatRestoreRow(Jfd, row, &ncols, nullptr, &vals); } if (small_elements > prune_fraction * total_elements) { - output.write("Pruning Jacobian elements: {} / {}\n", - small_elements, total_elements); + if (diagnose) { + output.write("\nPruning Jacobian elements: {} / {}\n", + small_elements, total_elements); + } // Prune Jacobian, keeping diagonal elements - //ierr = MatEliminateZeros(Jmf, PETSC_TRUE); CHKERRQ(ierr); - ierr = MatFilter(Jmf, prune_abstol, PETSC_TRUE, PETSC_TRUE); - - // Re-calculate the coloring - MatColoring coloring = NULL; - MatColoringCreate(Jmf, &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(Jmf, iscoloring, &fdcoloring); - MatFDColoringSetFunction(fdcoloring, - reinterpret_cast(FormFunctionForColoring), - this); - MatFDColoringSetFromOptions(fdcoloring); - MatFDColoringSetUp(Jmf, iscoloring, fdcoloring); - ISColoringDestroy(&iscoloring); - - // Replace the CTX pointer in SNES Jacobian - SNESSetJacobian(snes, Jmf, Jmf, SNESComputeJacobianScaledColor, fdcoloring); + //ierr = MatEliminateZeros(Jfd, PETSC_TRUE); CHKERRQ(ierr); + 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) @@ -1317,4 +1309,35 @@ PetscErrorCode SNESSolver::scaleJacobian(Mat B) { return 0; } +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(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, SNESComputeJacobianScaledColor, fdcoloring); + } else { + SNESSetJacobian(snes, Jfd, Jfd, SNESComputeJacobianScaledColor, fdcoloring); + } +} + #endif // BOUT_HAS_PETSC diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index d262d39db9..fd00100f81 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -4,9 +4,9 @@ * using PETSc for the SNES interface * ************************************************************************** - * Copyright 2015, 2021 B.D.Dudson + * Copyright 2015-2024 BOUT++ contributors * - * Contact: Ben Dudson, bd512@york.ac.uk + * Contact: Ben Dudson, dudson2@llnl.gov * * This file is part of BOUT++. * @@ -122,8 +122,9 @@ private: BoutReal time1{-1.0}; ///< Time of previous solution SNES snes; ///< SNES context - Mat Jmf; ///< Jacobian - MatFDColoring fdcoloring; ///< Matrix coloring context, used for finite difference + Mat Jmf; ///< Matrix Free Jacobian + Mat Jfd; ///< Finite Difference Jacobian + MatFDColoring fdcoloring {NULL}; ///< Matrix coloring context ///< Jacobian evaluation bool use_precon; ///< Use preconditioner @@ -133,7 +134,9 @@ private: std::string pc_type; ///< Preconditioner type std::string pc_hypre_type; ///< Hypre preconditioner type std::string line_search_type; ///< Line search type + bool matrix_free; ///< Use matrix free Jacobian + bool matrix_free_operator; ///< Use matrix free Jacobian in the operator? int lag_jacobian; ///< Re-use Jacobian bool use_coloring; ///< Use matrix coloring @@ -141,6 +144,9 @@ private: bool prune_jacobian; ///< Remove small elements in the Jacobian? BoutReal prune_abstol; ///< Prune values with absolute values smaller than this BoutReal prune_fraction; ///< Prune if fraction of small elements is larger than this + bool jacobian_pruned {false}; ///< Has the Jacobian been pruned? + Mat Jfd_original; ///< Used to reset the Jacobian if over-pruned + void updateColoring(); ///< Updates the coloring using Jfd bool scale_rhs; ///< Scale time derivatives? Vec rhs_scaling_factors; ///< Factors to multiply RHS function @@ -149,7 +155,6 @@ private: bool scale_vars; ///< Scale individual variables? Vec var_scaling_factors; ///< Factors to multiply variables when passing to user Vec scaled_x; ///< The values passed to the user RHS - }; #else From 393bbe5513d5c1a8f7a4118302317b2fbbacece5 Mon Sep 17 00:00:00 2001 From: bendudson Date: Thu, 24 Oct 2024 06:18:36 +0000 Subject: [PATCH 5/9] Apply clang-format changes --- src/solver/impls/snes/snes.cxx | 60 +++++++++++++++++++--------------- src/solver/impls/snes/snes.hxx | 20 ++++++------ 2 files changed, 43 insertions(+), 37 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index bc82556466..4fbd34aba9 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -15,8 +15,8 @@ #include -#include "petscsnes.h" #include "petscmat.h" +#include "petscsnes.h" /* * PETSc callback function, which evaluates the nonlinear @@ -150,8 +150,8 @@ SNESSolver::SNESSolver(Options* opts) .doc("Use matrix free Jacobian?") .withDefault(false)), matrix_free_operator((*options)["matrix_free_operator"] - .doc("Use matrix free Jacobian-vector operator?") - .withDefault(true)), + .doc("Use matrix free Jacobian-vector operator?") + .withDefault(true)), lag_jacobian((*options)["lag_jacobian"] .doc("Re-use the Jacobian this number of SNES iterations") .withDefault(50)), @@ -160,21 +160,20 @@ SNESSolver::SNESSolver(Options* opts) .withDefault(true)), jacobian_recalculated(false), prune_jacobian((*options)["prune_jacobian"] - .doc("Remove small elements in the Jacobian?") - .withDefault(false)), + .doc("Remove small elements in the Jacobian?") + .withDefault(false)), prune_abstol((*options)["prune_abstol"] .doc("Prune values with absolute values smaller than this") .withDefault(1e-16)), prune_fraction((*options)["prune_fraction"] - .doc("Prune if fraction of small elements is larger than this") - .withDefault(0.2)), + .doc("Prune if fraction of small elements is larger than this") + .withDefault(0.2)), scale_rhs((*options)["scale_rhs"] .doc("Scale time derivatives (Jacobian row scaling)?") .withDefault(false)), scale_vars((*options)["scale_vars"] - .doc("Scale variables (Jacobian column scaling)?") - .withDefault(false)) {} - + .doc("Scale variables (Jacobian column scaling)?") + .withDefault(false)) {} int SNESSolver::init() { @@ -656,7 +655,8 @@ int SNESSolver::init() { 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); + ierr = MatDuplicate(Jfd, MAT_SHARE_NONZERO_PATTERN, &Jfd_original); + CHKERRQ(ierr); } } else { // Brute force calculation @@ -808,25 +808,32 @@ int SNESSolver::run() { // Take ownership of snes_x and var_scaling_factors data PetscScalar* snes_x_data; - ierr = VecGetArray(snes_x, &snes_x_data); CHKERRQ(ierr); + ierr = VecGetArray(snes_x, &snes_x_data); + CHKERRQ(ierr); PetscScalar* x1_data; - ierr = VecGetArray(x1, &x1_data); CHKERRQ(ierr); + ierr = VecGetArray(x1, &x1_data); + CHKERRQ(ierr); PetscScalar* var_scaling_factors_data; - ierr = VecGetArray(var_scaling_factors, &var_scaling_factors_data); CHKERRQ(ierr); + 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]); + 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); + 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 @@ -842,7 +849,7 @@ int SNESSolver::run() { } } ++loop_count; - + // Copy the state (snes_x) into initial values (x0) VecCopy(snes_x, x0); @@ -928,7 +935,8 @@ int SNESSolver::run() { if (diagnose) { output.write("\nRestoring Jacobian\n"); } - ierr = MatCopy(Jfd_original, Jfd, DIFFERENT_NONZERO_PATTERN); CHKERRQ(ierr); + 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. @@ -1048,8 +1056,8 @@ int SNESSolver::run() { if (small_elements > prune_fraction * total_elements) { if (diagnose) { - output.write("\nPruning Jacobian elements: {} / {}\n", - small_elements, total_elements); + output.write("\nPruning Jacobian elements: {} / {}\n", small_elements, + total_elements); } // Prune Jacobian, keeping diagonal elements @@ -1083,8 +1091,7 @@ int SNESSolver::run() { // 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); + int ierr = VecPointwiseMult(scaled_x, snes_x, var_scaling_factors); CHKERRQ(ierr); const BoutReal* xdata = nullptr; @@ -1324,9 +1331,8 @@ void SNESSolver::updateColoring() { // Replace the old coloring with the new one MatFDColoringDestroy(&fdcoloring); MatFDColoringCreate(Jfd, iscoloring, &fdcoloring); - MatFDColoringSetFunction(fdcoloring, - reinterpret_cast(FormFunctionForColoring), - this); + MatFDColoringSetFunction( + fdcoloring, reinterpret_cast(FormFunctionForColoring), this); MatFDColoringSetFromOptions(fdcoloring); MatFDColoringSetUp(Jfd, iscoloring, fdcoloring); ISColoringDestroy(&iscoloring); diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index fd00100f81..7482fd4be4 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -124,8 +124,8 @@ private: SNES snes; ///< SNES context Mat Jmf; ///< Matrix Free Jacobian Mat Jfd; ///< Finite Difference Jacobian - MatFDColoring fdcoloring {NULL}; ///< Matrix coloring context - ///< Jacobian evaluation + MatFDColoring fdcoloring{NULL}; ///< Matrix coloring context + ///< Jacobian evaluation bool use_precon; ///< Use preconditioner std::string ksp_type; ///< Linear solver type @@ -134,19 +134,19 @@ private: std::string pc_type; ///< Preconditioner type std::string pc_hypre_type; ///< Hypre preconditioner type std::string line_search_type; ///< Line search type - + bool matrix_free; ///< Use matrix free Jacobian bool matrix_free_operator; ///< Use matrix free Jacobian in the operator? int lag_jacobian; ///< Re-use Jacobian bool use_coloring; ///< Use matrix coloring bool jacobian_recalculated; ///< Flag set when Jacobian is recalculated - bool prune_jacobian; ///< Remove small elements in the Jacobian? - BoutReal prune_abstol; ///< Prune values with absolute values smaller than this - BoutReal prune_fraction; ///< Prune if fraction of small elements is larger than this - bool jacobian_pruned {false}; ///< Has the Jacobian been pruned? - Mat Jfd_original; ///< Used to reset the Jacobian if over-pruned - void updateColoring(); ///< Updates the coloring using Jfd + bool prune_jacobian; ///< Remove small elements in the Jacobian? + BoutReal prune_abstol; ///< Prune values with absolute values smaller than this + BoutReal prune_fraction; ///< Prune if fraction of small elements is larger than this + bool jacobian_pruned{false}; ///< Has the Jacobian been pruned? + Mat Jfd_original; ///< Used to reset the Jacobian if over-pruned + void updateColoring(); ///< Updates the coloring using Jfd bool scale_rhs; ///< Scale time derivatives? Vec rhs_scaling_factors; ///< Factors to multiply RHS function @@ -154,7 +154,7 @@ private: bool scale_vars; ///< Scale individual variables? Vec var_scaling_factors; ///< Factors to multiply variables when passing to user - Vec scaled_x; ///< The values passed to the user RHS + Vec scaled_x; ///< The values passed to the user RHS }; #else From 07a1b60b9bb9309ded60da62efebf980df79bf69 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 24 Oct 2024 08:44:50 -0700 Subject: [PATCH 6/9] Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Peter Hill --- src/solver/impls/snes/snes.cxx | 5 ++--- src/solver/impls/snes/snes.hxx | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index 4fbd34aba9..226566d2b5 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -79,7 +79,7 @@ PetscErrorCode SNESComputeJacobianScaledColor(SNES snes, Vec x1, Mat J, Mat B, } // Get the the SNESSolver pointer from the function call context - SNESSolver* fctx; + SNESSolver* fctx = nullptr; err = MatFDColoringGetFunction(static_cast(ctx), nullptr, reinterpret_cast(&fctx)); CHKERRQ(err); @@ -807,7 +807,7 @@ int SNESSolver::run() { VecGetOwnershipRange(snes_x, &istart, &iend); // Take ownership of snes_x and var_scaling_factors data - PetscScalar* snes_x_data; + PetscScalar* snes_x_data = nullptr; ierr = VecGetArray(snes_x, &snes_x_data); CHKERRQ(ierr); PetscScalar* x1_data; @@ -1061,7 +1061,6 @@ int SNESSolver::run() { } // Prune Jacobian, keeping diagonal elements - //ierr = MatEliminateZeros(Jfd, PETSC_TRUE); CHKERRQ(ierr); ierr = MatFilter(Jfd, prune_abstol, PETSC_TRUE, PETSC_TRUE); // Update the coloring from Jfd matrix diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index 7482fd4be4..e7b3d1e8cc 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -124,7 +124,7 @@ private: SNES snes; ///< SNES context Mat Jmf; ///< Matrix Free Jacobian Mat Jfd; ///< Finite Difference Jacobian - MatFDColoring fdcoloring{NULL}; ///< Matrix coloring context + MatFDColoring fdcoloring{nullptr}; ///< Matrix coloring context ///< Jacobian evaluation bool use_precon; ///< Use preconditioner From adf2311e0dd8b66514cca825c81913ac3e5e01c7 Mon Sep 17 00:00:00 2001 From: bendudson Date: Thu, 24 Oct 2024 15:45:21 +0000 Subject: [PATCH 7/9] Apply clang-format changes --- src/solver/impls/snes/snes.hxx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/solver/impls/snes/snes.hxx b/src/solver/impls/snes/snes.hxx index e7b3d1e8cc..9e68e2bc79 100644 --- a/src/solver/impls/snes/snes.hxx +++ b/src/solver/impls/snes/snes.hxx @@ -125,7 +125,7 @@ private: Mat Jmf; ///< Matrix Free Jacobian Mat Jfd; ///< Finite Difference Jacobian MatFDColoring fdcoloring{nullptr}; ///< Matrix coloring context - ///< Jacobian evaluation + ///< Jacobian evaluation bool use_precon; ///< Use preconditioner std::string ksp_type; ///< Linear solver type From dec8ae7614ff0614e352229e41abe0f5c9bbe756 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Fri, 25 Oct 2024 14:04:10 -0700 Subject: [PATCH 8/9] beuler solver: Implement tidy suggestions - Making more indices const - Renaming Jacobian scaling function, removing SNES prefix to avoid potential collision with PETSc, and making function static. --- src/solver/impls/snes/snes.cxx | 116 ++++++++++++++++----------------- 1 file changed, 58 insertions(+), 58 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index 226566d2b5..9448e38076 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -59,35 +59,6 @@ static PetscErrorCode snesPCapply(PC pc, Vec x, Vec y) { PetscFunctionReturn(s->precon(x, y)); } -/// -/// Input Parameters: -/// snes - nonlinear solver object -/// x1 - location at which to evaluate Jacobian -/// ctx - MatFDColoring context or NULL -/// -/// Output Parameters: -/// J - Jacobian matrix (not altered in this routine) -/// B - newly computed Jacobian matrix to use with preconditioner (generally the same as -/// J) -PetscErrorCode SNESComputeJacobianScaledColor(SNES snes, Vec x1, Mat J, Mat B, - void* ctx) { - PetscErrorCode err = SNESComputeJacobianDefaultColor(snes, x1, J, B, 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(ctx), nullptr, - reinterpret_cast(&fctx)); - CHKERRQ(err); - - // Call the SNESSolver function - return fctx->scaleJacobian(B); -} - SNESSolver::SNESSolver(Options* opts) : Solver(opts), timestep( @@ -353,7 +324,7 @@ int SNESSolver::init() { // Lower X boundary for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { - int localIndex = ROUND(index(mesh->xstart, y, 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++) { @@ -365,7 +336,7 @@ int SNESSolver::init() { // On another processor for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { - int localIndex = ROUND(index(mesh->xstart, y, 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++) { @@ -379,7 +350,7 @@ int SNESSolver::init() { // Upper X boundary for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { - int localIndex = ROUND(index(mesh->xend, y, 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++) { @@ -391,7 +362,7 @@ int SNESSolver::init() { // On another processor for (int y = mesh->ystart; y <= mesh->yend; y++) { for (int z = 0; z < mesh->LocalNz; z++) { - int localIndex = ROUND(index(mesh->xend, y, 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++) { @@ -478,7 +449,7 @@ int SNESSolver::init() { // A boundary, so no communication // z = 0 case - int localIndex = ROUND(index(it.ind, mesh->yend, 0)); + const int localIndex = ROUND(index(it.ind, mesh->yend, 0)); if (localIndex < 0) { continue; // Out of domain } @@ -489,7 +460,7 @@ int SNESSolver::init() { } for (int z = 1; z < mesh->LocalNz; z++) { - int localIndex = ROUND(index(it.ind, mesh->yend, z)); + const int localIndex = ROUND(index(it.ind, mesh->yend, z)); // Only 3D fields for (int i = 0; i < n3d; i++) { @@ -525,28 +496,28 @@ int SNESSolver::init() { output_progress.write("Marking non-zero Jacobian entries\n"); - PetscScalar val = 1.0; + const PetscScalar val = 1.0; 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 5-point stencil for (const auto& xyoffset : xyoffsets) { - int xi = x + xyoffset.first; - int yi = y + xyoffset.second; + const int xi = x + xyoffset.first; + const int yi = y + xyoffset.second; 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 @@ -554,7 +525,7 @@ int SNESSolver::init() { // Depends on all variables on this cell for (int j = 0; j < n2d; j++) { - PetscInt col = ind2 + j; + const PetscInt col = ind2 + j; ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES); CHKERRQ(ierr); } @@ -564,7 +535,7 @@ int SNESSolver::init() { // 3D fields for (int z = 0; z < mesh->LocalNz; z++) { - int ind = ROUND(index(x, y, z)); + const int ind = ROUND(index(x, y, z)); for (int i = 0; i < n3d; i++) { PetscInt row = ind + i; @@ -574,15 +545,15 @@ int SNESSolver::init() { // Depends on 2D fields for (int j = 0; j < n2d; j++) { - PetscInt col = ind0 + j; + const PetscInt col = ind0 + j; ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES); CHKERRQ(ierr); } // Star pattern for (const auto& xyoffset : xyoffsets) { - int xi = x + xyoffset.first; - int yi = y + xyoffset.second; + const int xi = x + xyoffset.first; + const int yi = y + xyoffset.second; if ((xi < 0) || (yi < 0) || (xi >= mesh->LocalNx) || (yi >= mesh->LocalNy)) { @@ -600,7 +571,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; ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES); if (ierr != 0) { output.write("ERROR: {} : ({}, {}) -> ({}, {}) : {} -> {}\n", row, x, @@ -621,7 +592,7 @@ int SNESSolver::init() { ind2 += n2d; } for (int j = 0; j < n3d; j++) { - PetscInt col = ind2 + j; + const PetscInt col = ind2 + j; ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES); CHKERRQ(ierr); } @@ -632,7 +603,7 @@ int SNESSolver::init() { ind2 += n2d; } for (int j = 0; j < n3d; j++) { - PetscInt col = ind2 + j; + const PetscInt col = ind2 + j; ierr = MatSetValues(Jfd, 1, &row, 1, &col, &val, INSERT_VALUES); CHKERRQ(ierr); } @@ -1128,14 +1099,14 @@ PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) { const BoutReal* xdata = nullptr; ierr = VecGetArrayRead(scaled_x, &xdata); CHKERRQ(ierr); - load_vars(const_cast(xdata)); + load_vars(const_cast(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(xdata)); + load_vars(const_cast(xdata)); // const_cast needed due to load_vars API. Not writing to xdata. ierr = VecRestoreArrayRead(x, &xdata); CHKERRQ(ierr); } @@ -1245,7 +1216,7 @@ PetscErrorCode SNESSolver::precon(Vec x, Vec f) { return 0; } -PetscErrorCode SNESSolver::scaleJacobian(Mat B) { +PetscErrorCode SNESSolver::scaleJacobian(Mat Jac_new) { jacobian_recalculated = true; if (!scale_rhs) { @@ -1256,7 +1227,7 @@ PetscErrorCode SNESSolver::scaleJacobian(Mat B) { // Get index of rows owned by this processor int rstart, rend; - MatGetOwnershipRange(B, &rstart, &rend); + MatGetOwnershipRange(Jac_new, &rstart, &rend); // Check that the vector has the same ownership range int istart, iend; @@ -1274,7 +1245,7 @@ PetscErrorCode SNESSolver::scaleJacobian(Mat B) { PetscInt ncols; const PetscScalar* vals; for (int row = rstart; row < rend; ++row) { - MatGetRow(B, row, &ncols, nullptr, &vals); + MatGetRow(Jac_new, row, &ncols, nullptr, &vals); // Calculate a norm of this row of the Jacobian PetscScalar norm = 0.0; @@ -1290,7 +1261,7 @@ PetscErrorCode SNESSolver::scaleJacobian(Mat B) { // Store in the vector as 1 / norm row_inv_norm_data[row - rstart] = 1. / norm; - MatRestoreRow(B, row, &ncols, nullptr, &vals); + MatRestoreRow(Jac_new, row, &ncols, nullptr, &vals); } ierr = VecRestoreArray(jac_row_inv_norms, &row_inv_norm_data); @@ -1309,12 +1280,41 @@ PetscErrorCode SNESSolver::scaleJacobian(Mat B) { } // Scale the Jacobian rows by multiplying on the left by 1/norm - ierr = MatDiagonalScale(B, jac_row_inv_norms, nullptr); + 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(ctx), nullptr, + reinterpret_cast(&fctx)); + CHKERRQ(err); + + // Call the SNESSolver function + return fctx->scaleJacobian(Jac_new); +} + void SNESSolver::updateColoring() { // Re-calculate the coloring MatColoring coloring = NULL; @@ -1339,9 +1339,9 @@ void SNESSolver::updateColoring() { // 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, SNESComputeJacobianScaledColor, fdcoloring); + SNESSetJacobian(snes, Jmf, Jfd, ComputeJacobianScaledColor, fdcoloring); } else { - SNESSetJacobian(snes, Jfd, Jfd, SNESComputeJacobianScaledColor, fdcoloring); + SNESSetJacobian(snes, Jfd, Jfd, ComputeJacobianScaledColor, fdcoloring); } } From 89c32b7f2164e85cf62591af3487fa1e9c3d153f Mon Sep 17 00:00:00 2001 From: bendudson Date: Fri, 25 Oct 2024 21:07:01 +0000 Subject: [PATCH 9/9] Apply clang-format changes --- src/solver/impls/snes/snes.cxx | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/solver/impls/snes/snes.cxx b/src/solver/impls/snes/snes.cxx index 9448e38076..80d007898f 100644 --- a/src/solver/impls/snes/snes.cxx +++ b/src/solver/impls/snes/snes.cxx @@ -1099,14 +1099,16 @@ PetscErrorCode SNESSolver::snes_function(Vec x, Vec f, bool linear) { const BoutReal* xdata = nullptr; ierr = VecGetArrayRead(scaled_x, &xdata); CHKERRQ(ierr); - load_vars(const_cast(xdata)); // const_cast needed due to load_vars API. Not writing to xdata. + load_vars(const_cast( + 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(xdata)); // const_cast needed due to load_vars API. Not writing to xdata. + load_vars(const_cast( + xdata)); // const_cast needed due to load_vars API. Not writing to xdata. ierr = VecRestoreArrayRead(x, &xdata); CHKERRQ(ierr); }