momentum_newton_krylov.c

Go to the documentation of this file.

#include "momentumsolvers.h"
 
typedef enum {
    MOM_NK_JACOBIAN_FINITE_DIFFERENCE
} MomentumNewtonJacobianType;
 
typedef enum {
    MOM_NK_FD_MODE_MATRIX_FREE
} MomentumNewtonFiniteDifferenceMode;
 
typedef struct {
    MomentumNewtonJacobianType type;
    MomentumNewtonFiniteDifferenceMode finite_difference_mode;
    Mat jacobian_operator;
} MomentumNewtonJacobian;
 
typedef enum {
    MOM_NK_PC_MODEL_NONE,
    MOM_NK_PC_MODEL_FROZEN_MOMENTUM_JACOBIAN
} MomentumPreconditionerModel;
 
typedef enum {
    MOM_NK_PC_STRUCTURE_NONE,
    MOM_NK_PC_STRUCTURE_POINT_BLOCK
} MomentumPreconditionerStructure;
 
typedef struct {
    MomentumPreconditionerModel model;
    MomentumPreconditionerStructure structure;
    PetscInt block_size;
    PetscInt stencil_width;
} MomentumPreconditionerDescription;
 
typedef struct {
    PetscErrorCode (*Describe)(UserCtx *, MomentumPreconditionerDescription *);
    PetscErrorCode (*AssembleInterior)(UserCtx *, Vec, Mat);
} MomentumPreconditionerModelOps;
 
typedef struct {
    MomentumPreconditionerDescription description;
    const MomentumPreconditionerModelOps *model_ops;
    Mat preconditioning_matrix;
    PetscBool aliases_jacobian_operator;
    PetscBool owns_preconditioning_matrix;
    const char *petsc_pc_type;
} MomentumPreconditionerEngine;
 
typedef struct MomentumNewtonKrylovContext {
    UserCtx *user;
    SNES snes;
    FILE *history_file;
    FILE *linear_history_file;
    PetscBool have_initial_norm;
    PetscReal initial_norm;
    /* These objects are owned by the solve context, which MomentumSolver_NewtonKrylov() builds
     * on its first call and keeps for the run (UserCtx::momentumNewtonKrylov): nothing they are
     * built from changes between steps. The point-block pattern follows Nvert, which is fixed
     * for the run; an immersed boundary that moves would need it rebuilt. */
    MomentumNewtonJacobian jacobian;
    MomentumPreconditionerEngine preconditioning_engine;
    Vec solution, entry_backup;
} MomentumNewtonKrylovContext;
 
static PetscErrorCode MomentumNewtonKrylov_Validate(UserCtx *user);
static PetscErrorCode MomentumNewtonKrylov_FormResidual(SNES snes, Vec X, Vec F, void *ctx);
static PetscErrorCode MomentumNewtonKrylov_Monitor(SNES snes, PetscInt iteration,
                                                    PetscReal norm, void *ctx);
static PetscErrorCode MomentumNewtonKrylov_LinearMonitor(KSP ksp, PetscInt iteration,
                                                         PetscReal norm, void *ctx);
static void MomentumNewtonKrylov_OpenHistory(MomentumNewtonKrylovContext *ctx);
static void MomentumNewtonKrylov_WriteSummary(const MomentumNewtonKrylovContext *ctx,
                                               SNESConvergedReason reason,
                                               PetscInt nonlinear_its,
                                               PetscInt function_evals,
                                               PetscInt linear_its,
                                               PetscReal final_norm,
                                               PetscBool committed);
static PetscErrorCode MomentumNewtonKrylov_ApplyConstraints(MomentumNewtonKrylovContext *ctx,
                                                            Vec X, Vec F);
static PetscErrorCode MomentumNewtonKrylov_ReadLinearizationConfig(
    MomentumNewtonJacobian *jacobian, MomentumPreconditionerDescription *description);
static PetscErrorCode MomentumNewtonKrylov_FormJacobian(SNES snes, Vec current_solution,
    Mat jacobian_operator, Mat preconditioning_matrix, void *vctx);
static PetscErrorCode FrozenMomentumJacobian_AssemblePointBlocks(
    UserCtx *user, Vec current_solution, Mat preconditioning_matrix);
static PetscErrorCode MomentumPreconditionerEngine_CreateExactPointBlockMatrix(
    UserCtx *user, Mat *preconditioning_matrix);
 
/**
 * @brief Captures SNES iteration norms and optionally writes PICurv history rows.
 * @details SNES supplies the already-computed norm, so this monitor never causes
 * an additional nonlinear residual evaluation. PETSc monitors selected through
 * `-mom_nk_snes_monitor` remain independent and may run alongside this callback.
 */
static PetscErrorCode MomentumNewtonKrylov_Monitor(SNES snes, PetscInt iteration,
                                                    PetscReal norm, void *vctx)
{
    MomentumNewtonKrylovContext *ctx = (MomentumNewtonKrylovContext *)vctx;
    SimCtx *simCtx = ctx->user->simCtx;
 
    (void)snes;
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    if (iteration == 0 && !ctx->have_initial_norm) {
        ctx->initial_norm = norm;
        ctx->have_initial_norm = PETSC_TRUE;
    }
    if (ctx->history_file) {
        (void)fprintf(ctx->history_file, "%d,%d,%d,%.16e\n",
                      (int)simCtx->step, (int)ctx->user->_this, (int)iteration,
                      (double)norm);
        (void)fflush(ctx->history_file);
    }
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Writes the effective KSP tolerance and PETSc-reported norm for each inner iteration.
 * @details The tolerance is queried after SNES has applied any inexact-Newton forcing
 * update, so the log distinguishes changing Eisenstat--Walker requests from changing
 * linear convergence behavior.
 */
static PetscErrorCode MomentumNewtonKrylov_LinearMonitor(KSP ksp, PetscInt iteration,
                                                         PetscReal norm, void *vctx)
{
    MomentumNewtonKrylovContext *ctx = (MomentumNewtonKrylovContext *)vctx;
    PetscInt newton_iteration = -1;
    PetscReal relative_tolerance = 0.0;
 
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    if (!ctx->linear_history_file) { PROFILE_HELPER_END; PetscFunctionReturn(PETSC_SUCCESS); }
    PetscCall(SNESGetIterationNumber(ctx->snes, &newton_iteration));
    PetscCall(KSPGetTolerances(ksp, &relative_tolerance, NULL, NULL, NULL));
    (void)fprintf(ctx->linear_history_file, "%d,%d,%d,%d,%.16e,%.16e\n",
                  (int)ctx->user->simCtx->step, (int)ctx->user->_this,
                  (int)newton_iteration, (int)iteration,
                  (double)relative_tolerance, (double)norm);
    (void)fflush(ctx->linear_history_file);
    PROFILE_HELPER_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Opens one of the block's Newton CSV records on rank 0, or returns NULL.
 * @details PicurvOpenRunCsv() owns the lifecycle: a fresh run's first open truncates the file,
 *          the header is written to an empty file, and a continued run's first open marks the
 *          seam. A failure only warns: these files are diagnostics, and must not change the
 *          solve's error behaviour.
 */
static FILE *MomentumNewtonKrylov_OpenCsv(const MomentumNewtonKrylovContext *ctx, const char *kind, const char *header)
{
    PROFILE_HELPER_BEGIN;
    SimCtx *simCtx = ctx->user->simCtx;
    char name[128];
    FILE *file = NULL;
 
    if (PetscSNPrintf(name, sizeof(name), "Momentum_Solver_Newton_Krylov_%s_Block_%d.csv", kind, (int)ctx->user->_this) ||
        PicurvOpenRunCsv(simCtx, simCtx->log_dir, name, header, PICURV_CSV_TRUNCATE_ON_FRESH, PETSC_FALSE, &file)) {
        file = NULL;
    }
    PROFILE_HELPER_END;
    return file;
}
 
/** @brief Opens the optional rank-zero Newton iteration-history files (per Newton and per Krylov iteration). */
static void MomentumNewtonKrylov_OpenHistory(MomentumNewtonKrylovContext *ctx)
{
    PROFILE_FUNCTION_BEGIN;
    SimCtx *simCtx = ctx->user->simCtx;
 
    if (!simCtx->mom_nk_monitor_history || simCtx->rank != 0) { PROFILE_FUNCTION_END; return; }
    ctx->history_file = MomentumNewtonKrylov_OpenCsv(ctx, "History", "step,block,newton,nonlinear_norm");
    if (!ctx->history_file) { PROFILE_FUNCTION_END; return; }
    ctx->linear_history_file = MomentumNewtonKrylov_OpenCsv(ctx, "Linear_History",
                                                            "step,block,newton,krylov,requested_rtol,reported_residual_norm");
    PROFILE_FUNCTION_END;
}
 
/**
 * @brief Appends one rank-zero structured Newton result for a physical step.
 * @details File failures are deliberately diagnostic-only: rollback and PETSc
 * cleanup must retain their original error behavior.
 */
static void MomentumNewtonKrylov_WriteSummary(const MomentumNewtonKrylovContext *ctx,
                                               SNESConvergedReason reason,
                                               PetscInt nonlinear_its,
                                               PetscInt function_evals,
                                               PetscInt linear_its,
                                               PetscReal final_norm,
                                               PetscBool committed)
{
    PROFILE_FUNCTION_BEGIN;
    SimCtx *simCtx = ctx->user->simCtx;
    const char *reason_name;
    FILE *file;
 
    if (simCtx->rank != 0) { PROFILE_FUNCTION_END; return; }
    file = MomentumNewtonKrylov_OpenCsv(ctx, "Summary",
                                        "step,block,solver,jacobian,preconditioner,reason,reason_code,newton,evals,krylov,"
                                        "initial_norm,final_norm,state");
    if (!file) { PROFILE_FUNCTION_END; return; }
    reason_name = reason == SNES_CONVERGED_ITERATING
                    ? "SNES_CONVERGED_ITERATING" : SNESConvergedReasons[reason];
    (void)fprintf(file, "%d,%d,Newton Krylov,finite_difference / matrix_free,%s,%s,%d,%d,%d,%d,",
                  (int)simCtx->step, (int)ctx->user->_this,
                  ctx->preconditioning_engine.description.model == MOM_NK_PC_MODEL_NONE ?
                      "none" : "frozen_momentum_jacobian / point_block",
                  reason_name, (int)reason,
                  (int)nonlinear_its, (int)function_evals, (int)linear_its);
    if (ctx->have_initial_norm) (void)fprintf(file, "%.16e", (double)ctx->initial_norm);
    (void)fprintf(file, ",%.16e,%s\n", (double)final_norm,
                  !committed ? "rolled_back" :
                  simCtx->mom_last_converged ? "committed" : "committed_unconverged");
    (void)fclose(file);
    PROFILE_FUNCTION_END;
}
 
/**
 * @brief Rejects configurations outside the audited version-one feature set.
 * @param user Single-block momentum context to validate.
 * @return PetscErrorCode 0 when the configuration is supported.
 */
static PetscErrorCode MomentumNewtonKrylov_Validate(UserCtx *user)
{
    SimCtx   *simCtx;
    PetscReal mask_max = 0.0;
    PetscInt  velocity_dof = 0;
 
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCheck(user != NULL, PETSC_COMM_SELF, PETSC_ERR_ARG_NULL,
               "Newton Krylov requires a non-NULL UserCtx.");
    simCtx = user->simCtx;
    PetscCheck(simCtx != NULL, PETSC_COMM_SELF, PETSC_ERR_ARG_NULL,
               "Newton Krylov requires UserCtx::simCtx.");
    PetscCall(DMDAGetInfo(user->fda, NULL, NULL, NULL, NULL, NULL, NULL, NULL,
                           &velocity_dof, NULL, NULL, NULL, NULL, NULL));
    PetscCheck(velocity_dof == 3, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG,
               "Newton Krylov requires a three-component velocity DMDA (got dof=%" PetscInt_FMT ").",
               velocity_dof);
    PetscCheck(simCtx->block_number == 1, PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Newton Krylov version one supports exactly one block (got %" PetscInt_FMT ").",
               simCtx->block_number);
    PetscCheck(!simCtx->immersed, PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Newton Krylov version one does not support immersed boundaries.");
    PetscCheck(!simCtx->movefsi && !simCtx->rotatefsi, PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Newton Krylov version one does not support moving or rotating bodies/FSI.");
    PetscCheck(!simCtx->moveframe && !simCtx->rotateframe, PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Newton Krylov version one does not support moving or rotating reference frames.");
    for (PetscInt face = 0; face < 6; ++face) {
        const BoundaryFaceConfig *cfg = &user->boundary_faces[face];
        PetscBool supported = PETSC_FALSE;
 
        switch (cfg->handler_type) {
            case BC_HANDLER_WALL_NOSLIP:
            case BC_HANDLER_WALL_FREE_SLIP:
            /* As on an inlet, its velocity is refreshed once per step in PreStep. */
            case BC_HANDLER_WALL_MOVING:
                supported = (PetscBool)(cfg->mathematical_type == WALL);
                break;
            /* The inlet's velocity is refreshed once per step in PreStep, so every
             * residual evaluation within the step sees the same boundary value. */
            case BC_HANDLER_INLET_PRESCRIBED_VELOCITY:
                supported = (PetscBool)(cfg->mathematical_type == INLET);
                break;
            case BC_HANDLER_OUTLET_CONSERVATION:
                supported = (PetscBool)(cfg->mathematical_type == OUTLET);
                break;
            case BC_HANDLER_PERIODIC_GEOMETRIC:
            /* Both driven handlers freeze their correction once per timestep in
             * PreStep, which the Newton solve never re-enters, so the momentum
             * source stays constant across every residual evaluation. The paired
             * face checks below still apply. */
            case BC_HANDLER_PERIODIC_DRIVEN_CONSTANT_FLUX:
            case BC_HANDLER_PERIODIC_DRIVEN_INITIAL_FLUX:
                supported = (PetscBool)(cfg->mathematical_type == PERIODIC);
                break;
            default:
                supported = PETSC_FALSE;
                break;
        }
        PetscCheck(supported, PETSC_COMM_WORLD, PETSC_ERR_SUP,
                   "Newton Krylov version one does not support boundary face %" PetscInt_FMT " with mathematical type %" PetscInt_FMT " and handler %" PetscInt_FMT ".",
                   face, (PetscInt)cfg->mathematical_type, (PetscInt)cfg->handler_type);
    }
 
    PetscCheck((user->boundary_faces[BC_FACE_NEG_X].mathematical_type == PERIODIC) ==
               (user->boundary_faces[BC_FACE_POS_X].mathematical_type == PERIODIC),
               PETSC_COMM_WORLD, PETSC_ERR_SUP, "Newton Krylov requires paired x-periodic faces.");
    PetscCheck((user->boundary_faces[BC_FACE_NEG_Y].mathematical_type == PERIODIC) ==
               (user->boundary_faces[BC_FACE_POS_Y].mathematical_type == PERIODIC),
               PETSC_COMM_WORLD, PETSC_ERR_SUP, "Newton Krylov requires paired y-periodic faces.");
    PetscCheck((user->boundary_faces[BC_FACE_NEG_Z].mathematical_type == PERIODIC) ==
               (user->boundary_faces[BC_FACE_POS_Z].mathematical_type == PERIODIC),
               PETSC_COMM_WORLD, PETSC_ERR_SUP, "Newton Krylov requires paired z-periodic faces.");
 
    PetscCheck(user->Nvert != NULL, PETSC_COMM_SELF, PETSC_ERR_ARG_NULL,
               "Newton Krylov requires the cell-mask vector Nvert.");
    /* Solid cells no longer disqualify the solver: MomentumRowIsSolidMasked() gives
       their rows the same constrained treatment as any other row carrying no unknown.
       The immersed-boundary method itself remains unsupported above, because its
       velocity reconstruction does not run inside this solver's residual. */
    PetscCall(VecMax(user->Nvert, NULL, &mask_max));
    LOG_ALLOW(GLOBAL, LOG_DEBUG, "Newton Krylov solid-cell mask maximum Nvert=%g.\n",
              (double)mask_max);
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Reads application-owned Jacobian and preconditioner mathematics.
 */
static PetscErrorCode MomentumNewtonKrylov_ReadLinearizationConfig(
    MomentumNewtonJacobian *jacobian, MomentumPreconditionerDescription *description)
{
    char type[48] = "finite_difference";
    char finite_difference_mode[32] = "matrix_free";
    char model[48] = "none";
    char structure[32] = "none";
    PetscBool set = PETSC_FALSE, match = PETSC_FALSE;
 
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCheck(jacobian != NULL && description != NULL, PETSC_COMM_SELF,
               PETSC_ERR_ARG_NULL, "Newton Krylov linearization configuration is NULL.");
    PetscCall(PetscOptionsGetString(NULL, NULL, "-mom_nk_jacobian_type",
                                    type, sizeof(type), &set));
    PetscCall(PetscStrcasecmp(type, "finite_difference", &match));
    PetscCheck(match, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG,
               "-mom_nk_jacobian_type must be 'finite_difference' (got '%s').", type);
    jacobian->type = MOM_NK_JACOBIAN_FINITE_DIFFERENCE;
    PetscCall(PetscOptionsGetString(NULL, NULL, "-mom_nk_jacobian_fd_mode",
                                    finite_difference_mode,
                                    sizeof(finite_difference_mode), &set));
    PetscCall(PetscStrcasecmp(finite_difference_mode, "matrix_free", &match));
    PetscCheck(match, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG,
               "-mom_nk_jacobian_fd_mode must be 'matrix_free' (got '%s').",
               finite_difference_mode);
    jacobian->finite_difference_mode = MOM_NK_FD_MODE_MATRIX_FREE;
 
    PetscCall(PetscOptionsGetString(NULL, NULL, "-mom_nk_preconditioner_model",
                                    model, sizeof(model), &set));
    PetscCall(PetscStrcasecmp(model, "none", &match));
    if (match) description->model = MOM_NK_PC_MODEL_NONE;
    if (!match) {
        PetscCall(PetscStrcasecmp(model, "frozen_momentum_jacobian", &match));
        if (match) description->model = MOM_NK_PC_MODEL_FROZEN_MOMENTUM_JACOBIAN;
    }
    PetscCheck(match, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG,
               "-mom_nk_preconditioner_model must be 'none' or "
               "'frozen_momentum_jacobian' (got '%s').", model);
    PetscCall(PetscOptionsGetString(NULL, NULL, "-mom_nk_preconditioner_structure",
                                    structure, sizeof(structure), &set));
    PetscCall(PetscStrcasecmp(structure, "none", &match));
    if (match) description->structure = MOM_NK_PC_STRUCTURE_NONE;
    if (!match) {
        PetscCall(PetscStrcasecmp(structure, "point_block", &match));
        if (match) description->structure = MOM_NK_PC_STRUCTURE_POINT_BLOCK;
    }
    PetscCheck(match, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG,
               "-mom_nk_preconditioner_structure must be 'none' or 'point_block' "
               "(got '%s').", structure);
    PetscCheck((description->model == MOM_NK_PC_MODEL_NONE &&
                description->structure == MOM_NK_PC_STRUCTURE_NONE) ||
               (description->model == MOM_NK_PC_MODEL_FROZEN_MOMENTUM_JACOBIAN &&
                description->structure == MOM_NK_PC_STRUCTURE_POINT_BLOCK),
               PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Unsupported Newton Krylov preconditioner model/structure combination: "
               "model='%s', structure='%s'.", model, structure);
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Row classification with solid-cell masking folded in.
 *
 * A masked row's residual is identically zero for every X, so it has a zero Jacobian
 * row and a zero column and its unknown is undetermined. It therefore needs the same
 * identity treatment as any other row carrying no unknown. Checked after the boundary
 * classification so a periodic duplicate keeps its representative column.
 */
static MomentumRowType MomentumNewtonKrylov_ClassifyRow(
    UserCtx *user, const PetscReal ***nvert, PetscInt i, PetscInt j, PetscInt k,
    PetscInt component, PetscInt *ri, PetscInt *rj, PetscInt *rk)
{
    PROFILE_HELPER_BEGIN;
    const MomentumRowType type = ClassifyMomentumRow(user, i, j, k, component, ri, rj, rk);
 
    if (type == MOM_ROW_PHYSICAL && MomentumRowIsSolidMasked(nvert, i, j, k, component))
        { PROFILE_HELPER_END; return MOM_ROW_FIXED_HOMOGENEOUS; }
    PROFILE_HELPER_END;
    return type;
}
 
/** @brief Returns the squared Euclidean norm of one metric vector. */
static inline PetscReal FrozenMomentumJacobian_MetricNormSquared(Cmpnts metric)
{
    return metric.x * metric.x + metric.y * metric.y + metric.z * metric.z;
}
 
/**
 * @brief Returns the face-averaged eddy viscosity the residual uses on one face.
 *
 * The preconditioner has to reproduce the operator it preconditions, so the interior
 * face average here tracks `ComputeViscous()`.
 *
 * The WALL branch below does not, and is retained only because it cannot currently
 * change an assembled entry. It is unreachable by construction, not by configuration:
 * `FrozenMomentumJacobian_PointBlock()` puts the eddy viscosity only on the three
 * diagonal entries, and the row-at-a-time insertion in
 * `FrozenMomentumJacobian_AssemblePointBlocks()` means row c carries only axis == c.
 * The two coordinates that trigger the branch are exactly those
 * `ClassifyMomentumRow()` reports as non-MOM_ROW_PHYSICAL, and only physical rows
 * assemble a block; the WALL-on-a-periodic-axis loophole is closed by the pairing
 * check in `BoundarySystem_Validate()`. Both rules state the same boundary
 * fact: the wall-normal flux at a no-slip wall is not an unknown.
 *
 * On a wall-resolved face `ComputeViscous()` uses the molecular gradient flux with no eddy
 * viscosity, which this branch matches. On a wall-modelled face it applies the wall
 * stress `ApplyWallFunction()` stored, with no gradient flux at all: that flux depends on
 * the reference cell, the second from the wall, through the law, and not on the cells
 * either side of the face. Widening the stencil would make the branch reachable, and the
 * block would then need the law's derivative with respect to the reference velocity in
 * place of any viscosity here. Tracked as issue #8.
 */
static PetscReal FrozenMomentumJacobian_FaceEddyViscosity(
    const UserCtx *user, const PetscReal ***nu_t, PetscInt axis,
    PetscInt i, PetscInt j, PetscInt k)
{
    PROFILE_HELPER_BEGIN;
    const BCFace neg_face[3] = {BC_FACE_NEG_X, BC_FACE_NEG_Y, BC_FACE_NEG_Z};
    const BCFace pos_face[3] = {BC_FACE_POS_X, BC_FACE_POS_Y, BC_FACE_POS_Z};
    const PetscInt coord[3] = {i, j, k};
    const PetscInt size[3] = {user->info.mx, user->info.my, user->info.mz};
    PetscReal neighbour;
 
    if (nu_t == NULL) { PROFILE_HELPER_END; return 0.0; }
    if ((BoundaryFaceHasWallShear(&user->boundary_faces[neg_face[axis]]) && coord[axis] == 0) ||
        (BoundaryFaceHasWallShear(&user->boundary_faces[pos_face[axis]]) && coord[axis] == size[axis] - 2))
        { PROFILE_HELPER_END; return 0.0; }
 
    neighbour = (axis == 0) ? nu_t[k][j][i + 1]
              : (axis == 1) ? nu_t[k][j + 1][i]
                            : nu_t[k + 1][j][i];
    PROFILE_HELPER_END;
    return 0.5 * (nu_t[k][j][i] + neighbour);
}
 
/** @brief Returns the frozen-momentum point block for the current residual convention. */
static void FrozenMomentumJacobian_PointBlock(const UserCtx *user,
    const Cmpnts ***ucont, const Cmpnts ***csi, const Cmpnts ***eta,
    const Cmpnts ***zet, const PetscReal ***aj, const PetscReal ***nu_t,
    PetscInt i, PetscInt j, PetscInt k, PetscScalar block[9])
{
    PROFILE_HELPER_BEGIN;
    const SimCtx *simCtx = user->simCtx;
    const PetscReal dtc = MomentumBDFCoefficient((SimCtx *)simCtx) / simCtx->dt;
    const PetscReal molecular = simCtx->ren > 0.0 ? 1.0 / simCtx->ren : 0.0;
    /* The residual diffuses with nu + nu_t, so the block must too; nu_t alone is the
       eddy contribution. Omitting it left the preconditioner modelling a viscous
       diagonal smaller than the operator's by the eddy-to-molecular ratio, which on a
       developed LES is order one or more. */
    const PetscReal nu_eff_i = molecular + FrozenMomentumJacobian_FaceEddyViscosity(user, nu_t, 0, i, j, k);
    const PetscReal nu_eff_j = molecular + FrozenMomentumJacobian_FaceEddyViscosity(user, nu_t, 1, i, j, k);
    const PetscReal nu_eff_k = molecular + FrozenMomentumJacobian_FaceEddyViscosity(user, nu_t, 2, i, j, k);
    const PetscReal AJip = 0.5 * (aj[k][j][i] + aj[k][j][i + 1]);
    const PetscReal AJjp = 0.5 * (aj[k][j][i] + aj[k][j + 1][i]);
    const PetscReal AJkp = 0.5 * (aj[k][j][i] + aj[k + 1][j][i]);
    const PetscReal g11ip = csi[k][j][i].x * csi[k][j][i].x +
                             csi[k][j][i].y * csi[k][j][i].y +
                             csi[k][j][i].z * csi[k][j][i].z;
    const PetscReal g22ip = 0.25 * (
        FrozenMomentumJacobian_MetricNormSquared(eta[k][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(eta[k][j][i + 1]) +
        FrozenMomentumJacobian_MetricNormSquared(eta[k][j - 1][i]) +
        FrozenMomentumJacobian_MetricNormSquared(eta[k][j - 1][i + 1]));
    const PetscReal g33ip = 0.25 * (
        FrozenMomentumJacobian_MetricNormSquared(zet[k][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(zet[k][j][i + 1]) +
        FrozenMomentumJacobian_MetricNormSquared(zet[k - 1][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(zet[k - 1][j][i + 1]));
    const PetscReal g11jp = 0.25 * (
        FrozenMomentumJacobian_MetricNormSquared(csi[k][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(csi[k][j + 1][i]) +
        FrozenMomentumJacobian_MetricNormSquared(csi[k][j][i - 1]) +
        FrozenMomentumJacobian_MetricNormSquared(csi[k][j + 1][i - 1]));
    const PetscReal g22jp = eta[k][j][i].x * eta[k][j][i].x +
                             eta[k][j][i].y * eta[k][j][i].y +
                             eta[k][j][i].z * eta[k][j][i].z;
    const PetscReal g33jp = 0.25 * (
        FrozenMomentumJacobian_MetricNormSquared(zet[k][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(zet[k][j + 1][i]) +
        FrozenMomentumJacobian_MetricNormSquared(zet[k - 1][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(zet[k - 1][j + 1][i]));
    const PetscReal g11kp = 0.25 * (
        FrozenMomentumJacobian_MetricNormSquared(csi[k][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(csi[k + 1][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(csi[k][j][i - 1]) +
        FrozenMomentumJacobian_MetricNormSquared(csi[k + 1][j][i - 1]));
    const PetscReal g22kp = 0.25 * (
        FrozenMomentumJacobian_MetricNormSquared(eta[k][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(eta[k + 1][j][i]) +
        FrozenMomentumJacobian_MetricNormSquared(eta[k][j - 1][i]) +
        FrozenMomentumJacobian_MetricNormSquared(eta[k + 1][j - 1][i]));
    const PetscReal g33kp = zet[k][j][i].x * zet[k][j][i].x +
                             zet[k][j][i].y * zet[k][j][i].y +
                             zet[k][j][i].z * zet[k][j][i].z;
    const PetscReal U0jp = 0.25 * (ucont[k][j][i].x + ucont[k][j][i - 1].x +
                                   ucont[k][j + 1][i].x + ucont[k][j + 1][i - 1].x);
    const PetscReal U0kp = 0.25 * (ucont[k][j][i].x + ucont[k][j][i - 1].x +
                                   ucont[k + 1][j][i].x + ucont[k + 1][j][i - 1].x);
    const PetscReal U1ip = 0.25 * (ucont[k][j][i].y + ucont[k][j - 1][i].y +
                                   ucont[k][j][i + 1].y + ucont[k][j - 1][i + 1].y);
    const PetscReal U1kp = 0.25 * (ucont[k][j][i].y + ucont[k][j - 1][i].y +
                                   ucont[k + 1][j][i].y + ucont[k + 1][j - 1][i].y);
    const PetscReal U2ip = 0.25 * (ucont[k][j][i].z + ucont[k - 1][j][i].z +
                                   ucont[k][j][i + 1].z + ucont[k - 1][j][i + 1].z);
    const PetscReal U2jp = 0.25 * (ucont[k][j][i].z + ucont[k - 1][j][i].z +
                                   ucont[k][j + 1][i].z + ucont[k - 1][j + 1][i].z);
    PetscReal A[6][4] = {{0.0}};
    PetscReal Su, Sv, Sw, nui, nuj, nuk;
 
    A[0][0] =  0.125 * aj[k][j][i]         * ucont[k][j][i].y;
    A[0][1] = -0.125 * aj[k][j - 1][i]     * ucont[k][j - 1][i].y;
    A[0][2] =  0.125 * aj[k][j][i + 1]     * ucont[k][j][i + 1].y;
    A[0][3] = -0.125 * aj[k][j - 1][i + 1] * ucont[k][j - 1][i + 1].y;
    A[1][0] =  0.125 * aj[k][j][i]         * ucont[k][j][i].z;
    A[1][1] = -0.125 * aj[k - 1][j][i]     * ucont[k - 1][j][i].z;
    A[1][2] =  0.125 * aj[k][j][i + 1]     * ucont[k][j][i + 1].z;
    A[1][3] = -0.125 * aj[k - 1][j][i + 1] * ucont[k - 1][j][i + 1].z;
    A[2][0] = -0.125 * aj[k][j + 1][i - 1] * ucont[k][j + 1][i - 1].x;
    A[2][1] = -0.125 * aj[k][j][i - 1]     * ucont[k][j][i - 1].x;
    A[2][2] =  0.125 * aj[k][j + 1][i]     * ucont[k][j + 1][i].x;
    A[2][3] =  0.125 * aj[k][j][i]         * ucont[k][j][i].x;
    A[3][0] =  0.125 * aj[k][j][i]         * ucont[k][j][i].z;
    A[3][1] = -0.125 * aj[k - 1][j][i]     * ucont[k - 1][j][i].z;
    A[3][2] =  0.125 * aj[k][j + 1][i]     * ucont[k][j + 1][i].z;
    A[3][3] = -0.125 * aj[k - 1][j + 1][i] * ucont[k - 1][j + 1][i].z;
    A[4][0] = -0.125 * aj[k + 1][j][i - 1] * ucont[k + 1][j][i - 1].x;
    A[4][1] = -0.125 * aj[k][j][i - 1]     * ucont[k][j][i - 1].x;
    A[4][2] =  0.125 * aj[k + 1][j][i]     * ucont[k + 1][j][i].x;
    A[4][3] =  0.125 * aj[k][j][i]         * ucont[k][j][i].x;
    A[5][0] = -0.125 * aj[k + 1][j - 1][i] * ucont[k + 1][j - 1][i].y;
    A[5][1] = -0.125 * aj[k][j - 1][i]     * ucont[k][j - 1][i].y;
    A[5][2] =  0.125 * aj[k + 1][j][i]     * ucont[k + 1][j][i].y;
    A[5][3] =  0.125 * aj[k][j][i]         * ucont[k][j][i].y;
    Su = A[0][0] + A[0][1] + A[0][2] + A[0][3] +
         A[1][0] + A[1][1] + A[1][2] + A[1][3];
    Sv = A[2][0] + A[2][1] + A[2][2] + A[2][3] +
         A[3][0] + A[3][1] + A[3][2] + A[3][3];
    Sw = A[4][0] + A[4][1] + A[4][2] + A[4][3] +
         A[5][0] + A[5][1] + A[5][2] + A[5][3];
    nui = AJip * AJip * (g11ip + g22ip + g33ip) * nu_eff_i;
    nuj = AJjp * AJjp * (g11jp + g22jp + g33jp) * nu_eff_j;
    nuk = AJkp * AJkp * (g11kp + g22kp + g33kp) * nu_eff_k;
 
    /* MomentumNewtonKrylov_FormResidual forms F = -R, fixing this block's sign. */
    /* `block` is row-major: each row is one residual component and each column
       is one same-cell Ucont component. */
    block[0] = dtc + nui + Su; block[1] = 0.5 * AJjp * U0jp; block[2] = 0.5 * AJkp * U0kp;
    block[3] = 0.5 * AJip * U1ip; block[4] = dtc + nuj + Sv; block[5] = 0.5 * AJkp * U1kp;
    block[6] = 0.5 * AJip * U2ip; block[7] = 0.5 * AJjp * U2jp; block[8] = dtc + nuk + Sw;
    PROFILE_HELPER_END;
}
 
/** @brief Describes the audited frozen-coefficient point-block model. */
static PetscErrorCode FrozenMomentumJacobian_DescribePointBlock(
    UserCtx *user, MomentumPreconditionerDescription *description)
{
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    (void)user;
    PetscCheck(description != NULL, PETSC_COMM_SELF, PETSC_ERR_ARG_NULL,
               "Preconditioner description is NULL.");
    description->model = MOM_NK_PC_MODEL_FROZEN_MOMENTUM_JACOBIAN;
    description->structure = MOM_NK_PC_STRUCTURE_POINT_BLOCK;
    description->block_size = 3;
    description->stencil_width = 0;
    PROFILE_HELPER_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Inserts only the audited interior frozen-momentum point blocks. */
static PetscErrorCode FrozenMomentumJacobian_AssemblePointBlocks(
    UserCtx *user, Vec current_solution, Mat preconditioning_matrix)
{
    DMDALocalInfo info = user->info;
    Cmpnts ***ucont = NULL, ***csi = NULL, ***eta = NULL, ***zet = NULL;
    PetscReal ***aj = NULL, ***nvert = NULL, ***nu_t = NULL;
    SimCtx *simCtx = user->simCtx;
    /* The eddy viscosity enters the preconditioner only when a turbulence model is
       actually producing one; without it the field may not even be allocated. */
    const PetscBool has_eddy_viscosity =
        (PetscBool)(simCtx->les && user->lNu_t != NULL);
    PetscErrorCode ierr = PETSC_SUCCESS, cleanup_ierr;
 
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    /*
     * current_solution is the current SNES trial Ucont Vec: it is layout-compatible
     * with user->fda/user->Ucont, but is not necessarily the canonical user->Ucont
     * selected by UpdateLocalGhosts(FIELD_ID_UCONT). Scatter that trial Vec directly with
     * user->fda so PETSc applies its MPI ownership, periodic topology, component
     * ordering, and ghost mapping without canonical-field synchronization or mutation
     * of current_solution. UpdateLocalGhosts additionally repairs the component-normal
     * staggered buffers Uxi(i=-1,mx), Ueta(j=-1,my), and Uzeta(k=-1,mz); the audited
     * point-block velocity stencil reads none of those planes, so the repair cannot
     * change a coefficient. Revisit this choice and extend the periodic-localization
     * tests if the model stencil is expanded to read any repaired normal buffer.
     */
    ierr = DMGlobalToLocalBegin(user->fda, current_solution, INSERT_VALUES, user->lUcont);
    if (ierr) goto cleanup;
    ierr = DMGlobalToLocalEnd(user->fda, current_solution, INSERT_VALUES, user->lUcont);
    if (ierr) goto cleanup;
    ierr = DMDAVecGetArrayRead(user->fda, user->lUcont, &ucont); if (ierr) goto cleanup;
    ierr = DMDAVecGetArrayRead(user->fda, user->lCsi, &csi); if (ierr) goto cleanup;
    ierr = DMDAVecGetArrayRead(user->fda, user->lEta, &eta); if (ierr) goto cleanup;
    ierr = DMDAVecGetArrayRead(user->fda, user->lZet, &zet); if (ierr) goto cleanup;
    ierr = DMDAVecGetArrayRead(user->da, user->lAj, &aj); if (ierr) goto cleanup;
    ierr = DMDAVecGetArrayRead(user->da, user->lNvert, &nvert); if (ierr) goto cleanup;
    if (has_eddy_viscosity) {
        ierr = DMDAVecGetArrayRead(user->da, user->lNu_t, &nu_t); if (ierr) goto cleanup;
    }
    for (PetscInt k = info.zs; k < info.zs + info.zm; ++k) {
        for (PetscInt j = info.ys; j < info.ys + info.ym; ++j) {
            for (PetscInt i = info.xs; i < info.xs + info.xm; ++i) {
                for (PetscInt component = 0; component < 3; ++component) {
                    MatStencil row = {.i = i, .j = j, .k = k, .c = component};
                    PetscInt ri, rj, rk;
                    MomentumRowType type = MomentumNewtonKrylov_ClassifyRow(
                        user, (const PetscReal ***)nvert, i, j, k, component, &ri, &rj, &rk);
                    if (type == MOM_ROW_PHYSICAL) {
                        PetscScalar block[9];
                        MatStencil cols[3] = {
                            {.i = i, .j = j, .k = k, .c = 0},
                            {.i = i, .j = j, .k = k, .c = 1},
                            {.i = i, .j = j, .k = k, .c = 2}
                        };
                        FrozenMomentumJacobian_PointBlock(user, (const Cmpnts ***)ucont,
                            (const Cmpnts ***)csi, (const Cmpnts ***)eta, (const Cmpnts ***)zet,
                            (const PetscReal ***)aj, (const PetscReal ***)nu_t, i, j, k, block);
                        /* A TwoD-masked row is the time derivative alone in the residual,
                           so its block row is that derivative's diagonal. */
                        if (simCtx->TwoD == component + 1) {
                            for (PetscInt c = 0; c < 3; ++c)
                                block[3 * component + c] = (c == component)
                                    ? MomentumBDFCoefficient(simCtx) / simCtx->dt : 0.0;
                        }
                        ierr = MatSetValuesStencil(preconditioning_matrix, 1, &row, 3, cols,
                                                  &block[3 * component], INSERT_VALUES);
                        if (ierr) goto cleanup;
                    }
                }
            }
        }
    }
cleanup:
    if (nu_t) { cleanup_ierr = DMDAVecRestoreArrayRead(user->da, user->lNu_t, &nu_t); if (!ierr) ierr = cleanup_ierr; }
    if (nvert) { cleanup_ierr = DMDAVecRestoreArrayRead(user->da, user->lNvert, &nvert); if (!ierr) ierr = cleanup_ierr; }
    if (aj) { cleanup_ierr = DMDAVecRestoreArrayRead(user->da, user->lAj, &aj); if (!ierr) ierr = cleanup_ierr; }
    if (zet) { cleanup_ierr = DMDAVecRestoreArrayRead(user->fda, user->lZet, &zet); if (!ierr) ierr = cleanup_ierr; }
    if (eta) { cleanup_ierr = DMDAVecRestoreArrayRead(user->fda, user->lEta, &eta); if (!ierr) ierr = cleanup_ierr; }
    if (csi) { cleanup_ierr = DMDAVecRestoreArrayRead(user->fda, user->lCsi, &csi); if (!ierr) ierr = cleanup_ierr; }
    if (ucont) { cleanup_ierr = DMDAVecRestoreArrayRead(user->fda, user->lUcont, &ucont); if (!ierr) ierr = cleanup_ierr; }
    PROFILE_HELPER_END;
    PetscFunctionReturn(ierr);
}
 
/** @brief Inserts all common fixed, homogeneous, and periodic-duplicate rows. */
static PetscErrorCode MomentumPreconditionerEngine_ApplyConstraintRows(
    UserCtx *user, Mat preconditioning_matrix)
{
    DMDALocalInfo info = user->info;
    const PetscReal ***nvert = NULL;
 
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    PetscCall(DMDAVecGetArrayRead(user->da, user->lNvert, &nvert));
    for (PetscInt k = info.zs; k < info.zs + info.zm; ++k) {
        for (PetscInt j = info.ys; j < info.ys + info.ym; ++j) {
            for (PetscInt i = info.xs; i < info.xs + info.xm; ++i) {
                for (PetscInt component = 0; component < 3; ++component) {
                    PetscInt ri, rj, rk;
                    MomentumRowType type = MomentumNewtonKrylov_ClassifyRow(
                        user, nvert, i, j, k, component, &ri, &rj, &rk);
                    if (type != MOM_ROW_PHYSICAL) {
                        MatStencil row = {.i = i, .j = j, .k = k, .c = component};
                        MatStencil columns[2] = {row, row};
                        PetscScalar values[2] = {1.0, -1.0};
                        PetscInt column_count = 1;
                        if (type == MOM_ROW_PERIODIC_DUPLICATE) {
                            columns[column_count++] = (MatStencil){
                                .i = ri, .j = rj, .k = rk, .c = component
                            };
                        }
                        PetscCall(MatSetValuesStencil(preconditioning_matrix, 1, &row,
                                                      column_count, columns, values,
                                                      INSERT_VALUES));
                    }
                }
            }
        }
    }
    PetscCall(DMDAVecRestoreArrayRead(user->da, user->lNvert, &nvert));
    PROFILE_HELPER_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
static const MomentumPreconditionerModelOps frozen_momentum_point_block_ops = {
    FrozenMomentumJacobian_DescribePointBlock,
    FrozenMomentumJacobian_AssemblePointBlocks
};
 
/** @brief Creates the selected Jacobian operator; currently PETSc MFFD only. */
static PetscErrorCode MomentumNewtonJacobian_Create(SNES snes,
    MomentumNewtonJacobian *jacobian)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCheck(jacobian->type == MOM_NK_JACOBIAN_FINITE_DIFFERENCE &&
               jacobian->finite_difference_mode == MOM_NK_FD_MODE_MATRIX_FREE,
               PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Unsupported Newton Krylov Jacobian type/finite-difference-mode combination.");
    PetscCall(MatCreateSNESMF(snes, &jacobian->jacobian_operator));
    PetscCall(MatSetOptionsPrefix(jacobian->jacobian_operator, "mom_nk_"));
    PetscCall(PetscObjectSetName((PetscObject)jacobian->jacobian_operator,
                                 "momentum_jacobian_finite_difference_matrix_free"));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Updates the matrix-free finite-difference operator base. */
static PetscErrorCode MomentumNewtonJacobian_Update(SNES snes, Vec current_solution,
    MomentumNewtonJacobian *jacobian)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCall(MatMFFDComputeJacobian(snes, current_solution, jacobian->jacobian_operator,
                                     jacobian->jacobian_operator, NULL));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Registers the application orchestration callback and both SNES matrices. */
static PetscErrorCode MomentumNewtonJacobian_Register(SNES snes,
    MomentumNewtonJacobian *jacobian, MomentumPreconditionerEngine *engine,
    MomentumNewtonKrylovContext *ctx)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCall(SNESSetJacobian(snes, jacobian->jacobian_operator,
                              engine->preconditioning_matrix,
                              MomentumNewtonKrylov_FormJacobian, ctx));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Destroys a partially or fully created Jacobian operator. */
static PetscErrorCode MomentumNewtonJacobian_Destroy(MomentumNewtonJacobian *jacobian)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCall(MatDestroy(&jacobian->jacobian_operator));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Validates a model/structure and creates or aliases its matrix. */
static PetscErrorCode MomentumPreconditionerEngine_Create(UserCtx *user,
    Mat jacobian_operator, const MomentumPreconditionerDescription *requested,
    MomentumPreconditionerEngine *engine)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    engine->description = *requested;
    if (requested->model == MOM_NK_PC_MODEL_NONE &&
        requested->structure == MOM_NK_PC_STRUCTURE_NONE) {
        engine->description.block_size = 0;
        engine->description.stencil_width = 0;
        engine->preconditioning_matrix = jacobian_operator;
        engine->aliases_jacobian_operator = PETSC_TRUE;
        engine->owns_preconditioning_matrix = PETSC_FALSE;
        engine->petsc_pc_type = PCNONE;
    } else if (requested->model == MOM_NK_PC_MODEL_FROZEN_MOMENTUM_JACOBIAN &&
               requested->structure == MOM_NK_PC_STRUCTURE_POINT_BLOCK) {
        engine->model_ops = &frozen_momentum_point_block_ops;
        PetscCall(engine->model_ops->Describe(user, &engine->description));
        PetscCall(MomentumPreconditionerEngine_CreateExactPointBlockMatrix(
            user, &engine->preconditioning_matrix));
        engine->owns_preconditioning_matrix = PETSC_TRUE;
        PetscCall(PetscObjectSetName((PetscObject)engine->preconditioning_matrix,
                                     "momentum_preconditioner_frozen_point_block"));
        engine->aliases_jacobian_operator = PETSC_FALSE;
        engine->petsc_pc_type = PCPBJACOBI;
    } else {
        SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP,
                "Unsupported Newton Krylov preconditioner model/structure combination.");
    }
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Runs model insertion, common row handling, and final assembly. */
static PetscErrorCode MomentumPreconditionerEngine_Assemble(
    MomentumPreconditionerEngine *engine, UserCtx *user, Vec current_solution)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    if (engine->aliases_jacobian_operator) { PROFILE_FUNCTION_END; PetscFunctionReturn(PETSC_SUCCESS); }
    PetscCall(MatZeroEntries(engine->preconditioning_matrix));
    PetscCall(engine->model_ops->AssembleInterior(user, current_solution,
                                                   engine->preconditioning_matrix));
    PetscCall(MomentumPreconditionerEngine_ApplyConstraintRows(
        user, engine->preconditioning_matrix));
    PetscCall(MatAssemblyBegin(engine->preconditioning_matrix, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(engine->preconditioning_matrix, MAT_FINAL_ASSEMBLY));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Applies the validated model/structure-to-PETSc-PC mapping. */
static PetscErrorCode MomentumPreconditionerEngine_ConfigurePetscPC(
    MomentumPreconditionerEngine *engine, PC pc)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCall(PCSetType(pc, engine->petsc_pc_type));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Rejects raw options that select an unvalidated PETSc PC backend. */
static PetscErrorCode MomentumPreconditionerEngine_ValidatePetscPC(
    MomentumPreconditionerEngine *engine, PC pc)
{
    const char *actual_type = NULL;
    PetscBool matches = PETSC_FALSE;
 
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCall(PCGetType(pc, &actual_type));
    PetscCall(PetscStrcmp(actual_type, engine->petsc_pc_type, &matches));
    PetscCheck(matches, PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Newton Krylov preconditioner model/structure requires internal PETSc PC "
               "type '%s', but raw option processing selected '%s'.",
               engine->petsc_pc_type, actual_type ? actual_type : "(unset)");
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Destroys only a separately owned preconditioning matrix. */
static PetscErrorCode MomentumPreconditionerEngine_Destroy(
    MomentumPreconditionerEngine *engine)
{
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    if (!engine->owns_preconditioning_matrix) engine->preconditioning_matrix = NULL;
    PetscCall(MatDestroy(&engine->preconditioning_matrix));
    engine->aliases_jacobian_operator = PETSC_FALSE;
    engine->owns_preconditioning_matrix = PETSC_FALSE;
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Updates the Jacobian and then assembles any separate preconditioning matrix. */
static PetscErrorCode MomentumNewtonKrylov_FormJacobian(SNES snes, Vec current_solution,
    Mat jacobian_operator, Mat preconditioning_matrix, void *vctx)
{
    MomentumNewtonKrylovContext *ctx = (MomentumNewtonKrylovContext *)vctx;
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    (void)jacobian_operator;
    (void)preconditioning_matrix;
    PetscCall(MomentumNewtonJacobian_Update(snes, current_solution, &ctx->jacobian));
    PetscCall(MomentumPreconditionerEngine_Assemble(
        &ctx->preconditioning_engine, ctx->user, current_solution));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Creates the frozen point-block P matrix with its exact scalar pattern.
 * @details Row and column ownership follows the velocity DMDA global vector.
 * Physical rows reserve the three same-point components, fixed rows reserve
 * only their diagonal, and periodic duplicate rows additionally reserve their
 * wrapped representative. DMDA AO, local mapping, and stencil metadata are
 * retained so the existing insertion paths keep their exact ordering.
 */
static PetscErrorCode MomentumPreconditionerEngine_CreateExactPointBlockMatrix(
    UserCtx *user, Mat *preconditioning_matrix)
{
    DMDALocalInfo info;
    ISLocalToGlobalMapping local_to_global = NULL;
    Mat matrix = NULL;
    MPI_Comm comm;
    PetscInt local_size, global_size, ownership_start, ownership_end;
    PetscInt *diagonal_nnz = NULL, *offdiagonal_nnz = NULL;
    PetscInt ghost_starts[4] = {0, 0, 0, 0}, ghost_sizes[3] = {0, 0, 0};
    const PetscReal ***nvert = NULL;
    PetscErrorCode ierr = PETSC_SUCCESS, cleanup_ierr;
 
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    PetscCheck(preconditioning_matrix != NULL, PETSC_COMM_SELF, PETSC_ERR_ARG_NULL,
               "Point-block matrix output is NULL.");
    *preconditioning_matrix = NULL;
    comm = PetscObjectComm((PetscObject)user->fda);
    info = user->info;
    PetscCall(VecGetLocalSize(user->Ucont, &local_size));
    PetscCall(VecGetSize(user->Ucont, &global_size));
    PetscCall(VecGetOwnershipRange(user->Ucont, &ownership_start, &ownership_end));
    PetscCheck(local_size == ownership_end - ownership_start, comm, PETSC_ERR_PLIB,
               "Velocity ownership range does not match its local size.");
    ierr = PetscCalloc2(local_size, &diagonal_nnz,
                        local_size, &offdiagonal_nnz); if (ierr) goto cleanup;
    ierr = DMGetLocalToGlobalMapping(user->fda, &local_to_global); if (ierr) goto cleanup;
    ghost_starts[0] = info.gxs; ghost_starts[1] = info.gys; ghost_starts[2] = info.gzs;
    ghost_sizes[0] = info.gxm; ghost_sizes[1] = info.gym; ghost_sizes[2] = info.gzm;
    ierr = DMDAVecGetArrayRead(user->da, user->lNvert, &nvert); if (ierr) goto cleanup;
 
    for (PetscInt k = info.zs; k < info.zs + info.zm; ++k) {
        for (PetscInt j = info.ys; j < info.ys + info.ym; ++j) {
            for (PetscInt i = info.xs; i < info.xs + info.xm; ++i) {
                for (PetscInt component = 0; component < 3; ++component) {
                    PetscInt ri, rj, rk, column_count;
                    MatStencil row_stencil = {.i = i, .j = j, .k = k, .c = component};
                    MatStencil column_stencils[3];
                    PetscInt row_local, row, column_locals[3], columns[3];
                    MomentumRowType type = MomentumNewtonKrylov_ClassifyRow(
                        user, nvert, i, j, k, component, &ri, &rj, &rk);
 
                    if (type == MOM_ROW_PHYSICAL) {
                        column_count = 3;
                        for (PetscInt column_component = 0; column_component < 3;
                             ++column_component) {
                            column_stencils[column_component] = (MatStencil){
                                .i = i, .j = j, .k = k, .c = column_component
                            };
                        }
                    } else {
                        column_count = 1;
                        column_stencils[0] = row_stencil;
                        if (type == MOM_ROW_PERIODIC_DUPLICATE) {
                            column_stencils[column_count++] = (MatStencil){
                                .i = ri, .j = rj, .k = rk, .c = component
                            };
                        }
                    }
                    row_local = component + 3 * (
                        (i - ghost_starts[0]) + ghost_sizes[0] * (
                        (j - ghost_starts[1]) + ghost_sizes[1] *
                        (k - ghost_starts[2])));
                    for (PetscInt column_index = 0; column_index < column_count;
                         ++column_index) {
                        const MatStencil column = column_stencils[column_index];
                        if (!(column.i >= ghost_starts[0] &&
                              column.i < ghost_starts[0] + ghost_sizes[0] &&
                              column.j >= ghost_starts[1] &&
                              column.j < ghost_starts[1] + ghost_sizes[1] &&
                              column.k >= ghost_starts[2] &&
                              column.k < ghost_starts[2] + ghost_sizes[2])) {
                            ierr = PetscError(comm, __LINE__, PETSC_FUNCTION_NAME, __FILE__,
                                              PETSC_ERR_ARG_OUTOFRANGE,
                                              PETSC_ERROR_INITIAL,
                                              "Point-block column lies outside the DMDA ghost stencil.");
                            goto cleanup;
                        }
                        column_locals[column_index] = column.c + 3 * (
                            (column.i - ghost_starts[0]) + ghost_sizes[0] * (
                            (column.j - ghost_starts[1]) + ghost_sizes[1] *
                            (column.k - ghost_starts[2])));
                    }
                    ierr = ISLocalToGlobalMappingApply(local_to_global, 1,
                                                        &row_local, &row);
                    if (ierr) goto cleanup;
                    ierr = ISLocalToGlobalMappingApply(local_to_global, column_count,
                                                        column_locals, columns);
                    if (ierr) goto cleanup;
                    if (!(row >= ownership_start && row < ownership_end)) {
                        ierr = PetscError(comm, __LINE__, PETSC_FUNCTION_NAME, __FILE__,
                                          PETSC_ERR_PLIB, PETSC_ERROR_INITIAL,
                                          "DMDA-mapped point-block row is not locally owned.");
                        goto cleanup;
                    }
                    for (PetscInt column_index = 0; column_index < column_count;
                         ++column_index) {
                        PetscBool duplicate = PETSC_FALSE;
                        for (PetscInt previous = 0; previous < column_index; ++previous)
                            if (columns[previous] == columns[column_index]) duplicate = PETSC_TRUE;
                        if (columns[column_index] < 0) {
                            ierr = PetscError(comm, __LINE__, PETSC_FUNCTION_NAME, __FILE__,
                                              PETSC_ERR_PLIB, PETSC_ERROR_INITIAL,
                                              "DMDA-mapped point-block column is invalid.");
                            goto cleanup;
                        }
                        if (duplicate) continue;
                        if (columns[column_index] >= ownership_start &&
                            columns[column_index] < ownership_end)
                            ++diagonal_nnz[row - ownership_start];
                        else
                            ++offdiagonal_nnz[row - ownership_start];
                    }
                }
            }
        }
    }
 
    ierr = MatCreateAIJ(comm, local_size, local_size, global_size, global_size,
                        0, diagonal_nnz, 0, offdiagonal_nnz, &matrix);
    if (ierr) goto cleanup;
    ierr = MatSetBlockSize(matrix, 3); if (ierr) goto cleanup;
    ierr = MatSetLocalToGlobalMapping(matrix, local_to_global, local_to_global);
    if (ierr) goto cleanup;
    ierr = MatSetStencil(matrix, 3, ghost_sizes, ghost_starts, 3);
    if (ierr) goto cleanup;
    ierr = MatSetDM(matrix, user->fda); if (ierr) goto cleanup;
    ierr = MatSetOption(matrix, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
    if (ierr) goto cleanup;
    ierr = PetscFree2(diagonal_nnz, offdiagonal_nnz); if (ierr) goto cleanup;
 
    *preconditioning_matrix = matrix;
    matrix = NULL;
 
cleanup:
    if (nvert) {
        cleanup_ierr = DMDAVecRestoreArrayRead(user->da, user->lNvert, &nvert);
        if (!ierr) ierr = cleanup_ierr;
    }
    cleanup_ierr = MatDestroy(&matrix); if (!ierr) ierr = cleanup_ierr;
    cleanup_ierr = PetscFree2(diagonal_nnz, offdiagonal_nnz);
    if (!ierr) ierr = cleanup_ierr;
    PROFILE_HELPER_END;
    PetscFunctionReturn(ierr);
}
 
/**
 * @brief Replaces every non-independent residual row with an explicit equation.
 *
 * Which rows are non-independent is decided by ClassifyMomentumRow(), shared with
 * the residual path; only the action taken differs. Conditioned face-normal rows
 * use F=X-Uconditioned, unconditioned legacy dummy/tangential rows use F=X, and
 * periodic duplicates use Fdup=Xdup-Xrep. These equations prevent the zero Jacobian
 * rows that EnforceRHSBoundaryConditions()'s zeroing would produce in a matrix-free
 * Newton operator. Immersed and interface rows are rejected before this callback is
 * installed. A TwoD-masked component needs no constraint: ComputeRHS() zeroes its
 * spatial terms, so its row is the time derivative alone, which holds a velocity that
 * is constant in time exactly, as the Picard solver's skipped row does.
 * @param ctx Active solve context.
 * @param X Unconditioned PETSc trial state.
 * @param F Residual vector to update in place.
 * @return PetscErrorCode 0 on success.
 */
static PetscErrorCode MomentumNewtonKrylov_ApplyConstraints(MomentumNewtonKrylovContext *ctx,
                                                            Vec X, Vec F)
{
    UserCtx       *user = ctx->user;
    DMDALocalInfo  info = user->info;
    Vec            local_x = NULL;
    Cmpnts       ***x = NULL, ***conditioned = NULL, ***f = NULL, ***lx = NULL;
    const PetscInt xs = info.xs, xe = info.xs + info.xm;
    const PetscInt ys = info.ys, ye = info.ys + info.ym;
    const PetscInt zs = info.zs, ze = info.zs + info.zm;
 
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    PetscCall(DMGetLocalVector(user->fda, &local_x));
    PetscCall(DMGlobalToLocalBegin(user->fda, X, INSERT_VALUES, local_x));
    PetscCall(DMGlobalToLocalEnd(user->fda, X, INSERT_VALUES, local_x));
    PetscCall(DMDAVecGetArrayRead(user->fda, X, &x));
    PetscCall(DMDAVecGetArrayRead(user->fda, user->Ucont, &conditioned));
    PetscCall(DMDAVecGetArray(user->fda, F, &f));
    PetscCall(DMDAVecGetArrayRead(user->fda, local_x, &lx));
 
    for (PetscInt k = zs; k < ze; ++k) {
        for (PetscInt j = ys; j < ye; ++j) {
            for (PetscInt i = xs; i < xe; ++i) {
                PetscScalar *fv = &f[k][j][i].x;
                const PetscScalar *xv = &x[k][j][i].x;
                const PetscScalar *cv = &conditioned[k][j][i].x;
 
                for (PetscInt component = 0; component < 3; ++component) {
                    PetscInt ri, rj, rk;
                    MomentumRowType row = ClassifyMomentumRow(
                        user, i, j, k, component, &ri, &rj, &rk);
                    const PetscScalar *rv = &lx[rk][rj][ri].x;
 
                    if (row == MOM_ROW_FIXED_CONDITIONED) fv[component] = xv[component] - cv[component];
                    else if (row == MOM_ROW_FIXED_HOMOGENEOUS) fv[component] = xv[component];
                    else if (row == MOM_ROW_PERIODIC_DUPLICATE) fv[component] = xv[component] - rv[component];
                }
            }
        }
    }
 
    PetscCall(DMDAVecRestoreArrayRead(user->fda, local_x, &lx));
    PetscCall(DMDAVecRestoreArray(user->fda, F, &f));
    PetscCall(DMDAVecRestoreArrayRead(user->fda, user->Ucont, &conditioned));
    PetscCall(DMDAVecRestoreArrayRead(user->fda, X, &x));
    PetscCall(DMRestoreLocalVector(user->fda, &local_x));
    PROFILE_HELPER_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Adapts a PETSc trial vector to the existing momentum residual path.
 *
 * A matrix-free SNES residual must be a deterministic function of the trial
 * vector X alone: F(X) may not depend on any state left by a previous residual
 * or MFFD evaluation, or finite-difference Jacobian actions become inconsistent.
 * To honor that contract this callback fully derives the Cartesian velocity
 * state (Ucat/lUcat) from X before the first boundary sweep -- see the inline
 * comment below for why ApplyBoundaryConditions()'s own internal reconstruction
 * is not sufficient for the first outlet pass.
 *
 * State invariants:
 *   - On entry, X is the only input that determines the result; user->Ucont,
 *     user->Ucat and their local ghosts are treated as scratch and are fully
 *     overwritten from X.
 *   - Supported handlers may overwrite flux totals and other diagnostics on
 *     every call, but those values must not affect a later call at the same X;
 *     the deterministic seed guarantees this.
 *   - No histories, pressure, viscosity, or controller state advance here.
 *
 * Side effects: overwrites user->Ucont/lUcont, user->Ucat/lUcat, user->Rhs, the
 * boundary Ubcs targets, and boundary flux/area diagnostics; writes F.
 *
 * @param snes Calling nonlinear solver.
 * @param X Trial solution (read-only).
 * @param F Residual output.
 * @param vctx Pointer to MomentumNewtonKrylovContext.
 * @return PetscErrorCode 0 on success.
 */
static PetscErrorCode MomentumNewtonKrylov_FormResidual(SNES snes, Vec X, Vec F, void *vctx)
{
    MomentumNewtonKrylovContext *ctx = (MomentumNewtonKrylovContext *)vctx;
    UserCtx *user = ctx->user;
    const FieldId staggered_fields[] = {FIELD_ID_UCONT};
    const FieldId cell_fields[] = {FIELD_ID_UCAT};
 
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    (void)snes;
    PetscCall(VecCopy(X, user->Ucont));
    PetscCall(SynchronizePeriodicStaggeredFields(user, 1, staggered_fields));
 
    /* Deterministic pre-boundary Cartesian seed. Establish the full Ucat/lUcat
     * state from the current X before any boundary handler runs:
     *
     *     X -> Ucont/lUcont -> Ucat -> periodic Ucat -> lUcat -> boundaries
     *
     * Why this is required, and why it is NOT redundant with the reconstruction
     * already performed inside ApplyBoundaryConditions():
     *
     *  1. A matrix-free SNES residual must be a deterministic function of X. If
     *     the Cartesian state is left over from a previous residual/MFFD call,
     *     F(X) becomes history dependent and the finite-difference Jacobian
     *     action Jv = (F(X+hv)-F(X))/h is invalidated.
     *  2. The conservation-outlet handler reads lUcat during the FIRST boundary
     *     sweep (it measures the uncorrected outflow and builds the outlet
     *     profile from the Cartesian field). Without this seed it would read the
     *     stale lUcat from the preceding evaluation.
     *  3. ApplyBoundaryConditions() does reconstruct Ucat/lUcat, but only AFTER
     *     each handler sweep (Contra2Cart runs after BoundarySystem_ExecuteStep
     *     within every pass). Those internal updates therefore prepare passes 2
     *     and 3 -- they cannot prepare the very first outlet read of pass 1.
     *  4. Contra2Cart() rebuilds the global Ucat interior from the current
     *     lUcont, but it does not by itself refresh lUcat (nor lUcont; that was
     *     done by the SynchronizePeriodicStaggeredFields call above).
     *  5. SynchronizePeriodicCellFields(FIELD_ID_UCAT) must run before the ghost
     *     scatter so periodic duplicate planes are finalized consistently (it is
     *     a no-op when no direction is periodic, as on the straight duct).
     *  6. UpdateLocalGhosts(FIELD_ID_UCAT) is required because the outlet handler reads
     *     lUcat -- the local ghosted vector -- not merely the global Ucat.
     *
     * Do NOT "simplify" this to a bare Contra2Cart(user), and do NOT delete it
     * as apparently redundant with ApplyBoundaryConditions(): the three internal
     * boundary passes remain necessary (they refresh the Cartesian state after
     * each boundary correction), but only this sequence makes pass 1's input a
     * deterministic function of X. */
    PetscCall(Contra2Cart(user));
    PetscCall(SynchronizePeriodicCellFields(user, 1, cell_fields));
    PetscCall(UpdateLocalGhosts(user, FIELD_ID_UCAT));
 
    PetscCall(ApplyBoundaryConditions(user));
    PetscCall(ComputeTotalResidual(user));
    PetscCall(VecCopy(user->Rhs, F));
    PetscCall(VecScale(F, -1.0));
    PetscCall(MomentumNewtonKrylov_ApplyConstraints(ctx, X, F));
    PROFILE_HELPER_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/** @brief Builds the solve context MomentumSolver_NewtonKrylov() keeps for the run. */
static PetscErrorCode MomentumNewtonKrylov_Build(UserCtx *user)
{
    MomentumNewtonKrylovContext      *ctx;
    MomentumPreconditionerDescription preconditioner_description = {0};
    KSP                               ksp = NULL;
    PC                                pc = NULL;
 
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCall(PetscNew(&user->momentumNewtonKrylov));
    ctx = user->momentumNewtonKrylov;
    ctx->user = user;
    PetscCall(VecDuplicate(user->Ucont, &ctx->solution));
    PetscCall(VecDuplicate(user->Ucont, &ctx->entry_backup));
    PetscCall(SNESCreate(PetscObjectComm((PetscObject)user->Ucont), &ctx->snes));
    PetscCall(MomentumNewtonKrylov_ReadLinearizationConfig(&ctx->jacobian, &preconditioner_description));
    PetscCall(SNESSetOptionsPrefix(ctx->snes, "mom_nk_"));
    PetscCall(SNESSetType(ctx->snes, SNESNEWTONLS));
    PetscCall(SNESSetDM(ctx->snes, user->fda));
    /* The residual and Jacobian callbacks are registered by MomentumSolver_NewtonKrylov() before
       every solve, so none is registered here. */
    PetscCall(MomentumNewtonJacobian_Create(ctx->snes, &ctx->jacobian));
    PetscCall(MomentumPreconditionerEngine_Create(user, ctx->jacobian.jacobian_operator,
                                                  &preconditioner_description,
                                                  &ctx->preconditioning_engine));
    PetscCall(SNESGetKSP(ctx->snes, &ksp));
    PetscCall(KSPSetType(ksp, KSPGMRES));
    PetscCall(KSPGetPC(ksp, &pc));
    PetscCall(MomentumPreconditionerEngine_ConfigurePetscPC(&ctx->preconditioning_engine, pc));
    PetscCall(SNESSetFromOptions(ctx->snes));
    PetscCall(SNESMonitorSet(ctx->snes, MomentumNewtonKrylov_Monitor, ctx, NULL));
    PetscCall(KSPMonitorSet(ksp, MomentumNewtonKrylov_LinearMonitor, ctx, NULL));
    PetscCall(MomentumPreconditionerEngine_ValidatePetscPC(&ctx->preconditioning_engine, pc));
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Implementation of \ref MomentumSolver_NewtonKrylov_Destroy().
 * @details Full API contract (arguments, ownership, side effects) is documented with
 *          the header declaration in `include/momentumsolvers.h`.
 * @see MomentumSolver_NewtonKrylov_Destroy()
 */
PetscErrorCode MomentumSolver_NewtonKrylov_Destroy(UserCtx *user)
{
    MomentumNewtonKrylovContext *ctx = user->momentumNewtonKrylov;
 
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    if (ctx) {
        PetscCall(VecDestroy(&ctx->entry_backup));
        PetscCall(VecDestroy(&ctx->solution));
        PetscCall(MomentumPreconditionerEngine_Destroy(&ctx->preconditioning_engine));
        PetscCall(MomentumNewtonJacobian_Destroy(&ctx->jacobian));
        PetscCall(SNESDestroy(&ctx->snes));
        PetscCall(PetscFree(user->momentumNewtonKrylov));
    }
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/**
 * @brief Measure the residual at the iterate a failed SNES solve holds, for the
 *        failure policy in MomentumSolver_NewtonKrylov().
 *
 * SNES judges the residual by absolute and relative tolerances that do not scale with the
 * grid or the flow. Near a steady state the residual can reach the round-off floor of its
 * own evaluation; the line search then cannot decrease it and SNES reports divergence for
 * a state that is as converged as the discretization allows. The caller therefore judges
 * the held iterate itself: by the dual-time solver's dimensionless test,
 * `|R|_inf <= mom_resid_atol * a0 |U|_inf / dt`, and by whether its residual is finite and
 * no larger than the residual the step started from.
 *
 * @param[in]  ctx        Solve context; its SNES supplies the residual function and vector.
 * @param[in]  solution   Iterate to test; it is left unchanged.
 * @param[out] residual   `|R|_inf` at @p solution.
 * @param[out] residual_2 `|R|_2` at @p solution, the norm SNES reports as its initial norm.
 * @param[out] reference  `a0 |U|_inf / dt` at @p solution.
 * @param[out] satisfied  True when `mom_resid_atol > 0` and the dimensionless test holds.
 */
static PetscErrorCode MomentumNewtonKrylov_AssessIterate(MomentumNewtonKrylovContext *ctx, Vec solution,
                                                         PetscReal *residual, PetscReal *residual_2,
                                                         PetscReal *reference, PetscBool *satisfied)
{
    SimCtx   *simCtx = ctx->user->simCtx;
    Vec       F = NULL;
    PetscReal ucont_inf;
 
    PetscFunctionBeginUser;
    PROFILE_HELPER_BEGIN;
    PetscCall(SNESGetFunction(ctx->snes, &F, NULL, NULL));
    PetscCall(SNESComputeFunction(ctx->snes, solution, F));
    PetscCall(VecNorm(F, NORM_INFINITY, residual));
    PetscCall(VecNorm(F, NORM_2, residual_2));
    PetscCall(VecNorm(solution, NORM_INFINITY, &ucont_inf));
    *reference = MomentumBDFCoefficient(simCtx) * ucont_inf / simCtx->dt;
    *satisfied = (PetscBool)(simCtx->mom_resid_atol > 0.0 && *reference > 0.0 &&
                             *residual <= simCtx->mom_resid_atol * *reference);
    PROFILE_HELPER_END;
    PetscFunctionReturn(PETSC_SUCCESS);
}
 
/*
 * Runs one matrix-free Newton--Krylov momentum solve, with the solve context built on the
 * first call and kept for the run. The public header owns the rendered API contract; this
 * definition retains the detailed lifecycle and rollback behavior below.
 */
PetscErrorCode MomentumSolver_NewtonKrylov(UserCtx *user, IBMNodes *ibm, FSInfo *fsi)
{
    PetscErrorCode ierr = PETSC_SUCCESS, cleanup_ierr;
    SimCtx        *simCtx;
    SNES           snes = NULL;
    Vec            solution = NULL, entry_backup = NULL;
    MomentumNewtonKrylovContext *ctx = NULL;
    PetscBool      restore_entry = PETSC_FALSE;
    PetscBool      solve_started = PETSC_FALSE;
    PetscBool      committed = PETSC_FALSE;
    PetscBool      converged = PETSC_FALSE;
    SNESConvergedReason reason = SNES_CONVERGED_ITERATING;
    PetscInt       nonlinear_its = 0, function_evals = 0, linear_its = 0;
    PetscReal      final_norm = PETSC_MAX_REAL;
    const FieldId staggered_fields[] = {FIELD_ID_UCONT};
 
    PetscFunctionBeginUser;
    PROFILE_FUNCTION_BEGIN;
    PetscCall(MomentumNewtonKrylov_Validate(user));
    PetscCheck(ibm == NULL && fsi == NULL, PETSC_COMM_WORLD, PETSC_ERR_SUP,
               "Newton Krylov version one does not accept IBM or FSI objects.");
    simCtx = user->simCtx;
 
    /* Rhs is kept for the run and zeroed for each step. */
    if (!user->Rhs) { ierr = VecDuplicate(user->Ucont, &user->Rhs); if (ierr) goto cleanup; }
    ierr = VecZeroEntries(user->Rhs); if (ierr) goto cleanup;
 
    ierr = SynchronizePeriodicStaggeredFields(user, 1, staggered_fields); if (ierr) goto cleanup;
    ierr = ApplyBoundaryConditions(user); if (ierr) goto cleanup;
    if (!user->momentumNewtonKrylov) { ierr = MomentumNewtonKrylov_Build(user); if (ierr) goto cleanup; }
    ctx = user->momentumNewtonKrylov;
    ctx->have_initial_norm = PETSC_FALSE;
    ctx->initial_norm = 0.0;
    snes = ctx->snes;
    /* The residual and Jacobian callbacks live on fda's DMSNES, which every SNES on fda shares;
       registering them before each solve keeps another solver on fda from leaving its own in place. */
    ierr = SNESSetFunction(snes, NULL, MomentumNewtonKrylov_FormResidual, ctx); if (ierr) goto cleanup;
    ierr = MomentumNewtonJacobian_Register(snes, &ctx->jacobian,
                                            &ctx->preconditioning_engine, ctx); if (ierr) goto cleanup;
    solution = ctx->solution;
    entry_backup = ctx->entry_backup;
    ierr = VecCopy(user->Ucont, entry_backup); if (ierr) goto cleanup;
    restore_entry = PETSC_TRUE;
    ierr = VecCopy(user->Ucont, solution); if (ierr) goto cleanup;
 
    LOG_ALLOW(GLOBAL, LOG_INFO,
              "Newton Krylov Jacobian: finite_difference / matrix_free; "
              "Preconditioner: %s; PETSc Jacobian matrix type: MATMFFD; "
              "PETSc PC type: %s.\n",
              ctx->preconditioning_engine.description.model == MOM_NK_PC_MODEL_NONE ?
                  "none" : "frozen_momentum_jacobian / point_block",
              ctx->preconditioning_engine.petsc_pc_type);
 
    MomentumNewtonKrylov_OpenHistory(ctx);
    solve_started = PETSC_TRUE;
    ierr = SNESSolve(snes, NULL, solution);
    if (ierr) goto cleanup;
    ierr = SNESGetConvergedReason(snes, &reason); if (ierr) goto cleanup;
    ierr = SNESGetIterationNumber(snes, &nonlinear_its); if (ierr) goto cleanup;
    ierr = SNESGetNumberFunctionEvals(snes, &function_evals); if (ierr) goto cleanup;
    ierr = SNESGetLinearSolveIterations(snes, &linear_its); if (ierr) goto cleanup;
    ierr = SNESGetFunctionNorm(snes, &final_norm); if (ierr) goto cleanup;
 
    /* Failure policy, shared with the dual-time solver: a solve SNES did not declare
       converged still commits its iterate when the residual there is finite and no larger
       than at entry, and the step is reported as not converged. Only a non-finite or
       grown residual rolls the step back and stops the run. */
    if (reason > 0) {
        committed = PETSC_TRUE;
        converged = PETSC_TRUE;
    } else {
        PetscReal held_residual = 0.0, held_residual_2 = 0.0, held_reference = 0.0;
        ierr = MomentumNewtonKrylov_AssessIterate(ctx, solution, &held_residual, &held_residual_2,
                                                  &held_reference, &converged); if (ierr) goto cleanup;
        committed = (PetscBool)(!PetscIsInfOrNanReal(held_residual_2) && ctx->have_initial_norm &&
                                held_residual_2 <= ctx->initial_norm);
        if (!committed) converged = PETSC_FALSE;
        if (converged) {
            LOG(GLOBAL, LOG_WARNING,
                "Newton Krylov step %" PetscInt_FMT ": SNES reported %s, but |R|_inf = %.3e is within "
                "residual_absolute_tol (%.1e) x a0|U|/dt = %.3e; the step is committed as converged.\n",
                simCtx->step, SNESConvergedReasons[reason], (double)held_residual,
                (double)simCtx->mom_resid_atol, (double)held_reference);
        } else if (committed) {
            LOG(GLOBAL, LOG_WARNING,
                "Newton Krylov step %" PetscInt_FMT ": SNES reported %s with |R|_inf = %.3e, above "
                "residual_absolute_tol (%.1e) x a0|U|/dt = %.3e; |R|_2 = %.3e is no larger than %.3e at "
                "entry, so the step continues from the SNES iterate, not converged.\n",
                simCtx->step, SNESConvergedReasons[reason], (double)held_residual,
                (double)simCtx->mom_resid_atol, (double)held_reference,
                (double)held_residual_2, (double)ctx->initial_norm);
        } else {
            if (ctx->have_initial_norm) {
                LOG(GLOBAL, LOG_ERROR,
                    "Newton Krylov step %" PetscInt_FMT ": SNES reported %s and |R|_2 at its iterate is "
                    "%.3e against %.3e at entry; the step is rolled back.\n",
                    simCtx->step, SNESConvergedReasons[reason], (double)held_residual_2,
                    (double)ctx->initial_norm);
            } else {
                LOG(GLOBAL, LOG_ERROR,
                    "Newton Krylov step %" PetscInt_FMT ": SNES reported %s before an entry residual norm "
                    "was available (|R|_2 at its iterate is %.3e); the step is rolled back.\n",
                    simCtx->step, SNESConvergedReasons[reason], (double)held_residual_2);
            }
        }
    }
 
    ierr = VecCopy(committed ? solution : entry_backup, user->Ucont); if (ierr) goto cleanup;
    simCtx->mom_last_converged = converged;
    ierr = SynchronizePeriodicStaggeredFields(user, 1, staggered_fields); if (ierr) goto cleanup;
    ierr = ApplyBoundaryConditions(user); if (ierr) goto cleanup;
    restore_entry = PETSC_FALSE;
 
    LOG_ALLOW(GLOBAL, LOG_INFO,
              "Newton Krylov momentum solve: reason=%s (%" PetscInt_FMT "), Newton iterations=%" PetscInt_FMT ", residual evaluations=%" PetscInt_FMT ", Krylov iterations=%" PetscInt_FMT ", final norm=%.6e, state=%s.\n",
              SNESConvergedReasons[reason], (PetscInt)reason, nonlinear_its, function_evals,
              linear_its, (double)final_norm,
              !committed ? "rolled back" : converged ? "committed" : "committed, not converged");
    if (!committed) ierr = PETSC_ERR_CONV_FAILED;
 
cleanup:
    /* A PETSc solve error can bypass the normal statistics path. Query whatever
       SNES retained without replacing the primary error so a failed attempt is
       still represented in the structured log. */
    if (solve_started && snes) {
        (void)SNESGetConvergedReason(snes, &reason);
        (void)SNESGetIterationNumber(snes, &nonlinear_its);
        (void)SNESGetNumberFunctionEvals(snes, &function_evals);
        (void)SNESGetLinearSolveIterations(snes, &linear_its);
        (void)SNESGetFunctionNorm(snes, &final_norm);
    }
    if (restore_entry && entry_backup) {
        cleanup_ierr = VecCopy(entry_backup, user->Ucont);
        if (!ierr) ierr = cleanup_ierr;
        cleanup_ierr = SynchronizePeriodicStaggeredFields(user, 1, staggered_fields);
        if (!ierr) ierr = cleanup_ierr;
        cleanup_ierr = ApplyBoundaryConditions(user);
        if (!ierr) ierr = cleanup_ierr;
        simCtx->mom_last_converged = PETSC_FALSE;
        committed = PETSC_FALSE;
    }
    if (ctx && ctx->history_file) {
        (void)fclose(ctx->history_file);
        ctx->history_file = NULL;
    }
    if (ctx && ctx->linear_history_file) {
        (void)fclose(ctx->linear_history_file);
        ctx->linear_history_file = NULL;
    }
    if (solve_started) {
        MomentumNewtonKrylov_WriteSummary(ctx, reason, nonlinear_its, function_evals,
                                          linear_its, final_norm, committed);
    }
    /* A failed step, or a failed first build, releases what the solver holds, so the next call
       starts afresh; a committed step keeps it for the next. */
    if (ierr) {
        cleanup_ierr = VecDestroy(&user->Rhs); if (!ierr) ierr = cleanup_ierr;
        cleanup_ierr = MomentumSolver_NewtonKrylov_Destroy(user); if (!ierr) ierr = cleanup_ierr;
    }
    PROFILE_FUNCTION_END;
    PetscFunctionReturn(ierr);
}