Implicit Core

Numerical-version policy: the active double-gyre implementation is implicit-dg-v2-pre-safe-sqrt. See Implicit double-gyre numerical version policy.

The implicit core uses a semi-implicit free-surface formulation that lifts the restrictive CFL condition imposed by surface gravity waves. By solving an elliptic equation for the free-surface elevation $\eta$, the implicit core achieves stable integration with substantially larger time steps -- making it the preferred core for neural network training and long climate-scale simulations.

Architecture

The implicit core follows a two-level class hierarchy. A reference implementation provides the baseline numerical methods, while an accelerated subclass adds configurable PCG preconditioning, cached geometric terms, and reusable scratch buffers. The active V2 profile uses the spectral preconditioner; Jacobi remains available for V1 replay controls:

BaseCore
  └── ReferencePOM2KCore (implicit/reference_core.py)
        ├── Inherits: GridMixin, CasesMixin, InitializationMixin,
        │             BaroclinicMixin, TracersMixin, MomentumMixin,
        │             DiagnosticsMixin, CheckpointingMixin,
        │             TimeLoopMixin, ImplicitBarotropicMixin
        │
        └── POM2K_Core (implicit/core.py)   ← Accelerated
              ├── Cached geometry terms
              ├── Spectral/Jacobi-preconditioned PCG
              └── Reusable scratch buffers

# From src/neuralpom/cores/implicit/core.py
class POM2K_Core(_ReferencePOM2KCore):
    """
    Accelerated implicit core used as the default implicit
    implementation in NeuralPOM.

    The underlying elliptic problem is unchanged. This class
    upgrades the default implicit core with cached geometry terms,
    configurable PCG preconditioning, and reusable scratch buffers for
    simulation efficiency.
    """
    def __init__(self, im, jm, *args, **kwargs):
        super().__init__(im, jm, *args, **kwargs)
        self._accel_cache_ready = False
        self._accel_cache = None
        self._cg_work = None
        self.eta_solver_rel_tol = 1e-5

Mode Splitting

The implicit core uses a fractional step approach to separate the barotropic and baroclinic dynamics, but unlike the explicit core, it takes only a single barotropic step per baroclinic timestep. The mode splitting proceeds in five stages:

Stage 1: Baroclinic Step (Pre-Barotropic)

Before the barotropic solver, the following 3D calculations are performed:

1. Horizontal advection and lateral viscosity (Smagorinsky scheme)
2. Baroclinic pressure gradient ($\mathrm{\nabla }{p}_{\text{bc}}$ from density field)
3. Depth-averaged forcing terms (adx2d, ady2d, drx2d, dry2d, aam2d)
4. Depth-averaged advection (advave): computes advua, advva, wubot, wvbot

These 2D-integrated quantities serve as forcing for the barotropic system.

Stage 2: Predictor (Provisional Barotropic Velocities ${\mathbf{u}}^{\ast }$, ${\mathbf{v}}^{\ast }$)

The explicit part of the barotropic momentum equation is computed. All terms that do not involve ${\eta }^{n+1}$ are evaluated to obtain provisional velocities ${\mathbf{u}}^{\ast }$ and ${\mathbf{v}}^{\ast }$:

# From src/neuralpom/cores/implicit/core.py,
#    _compute_barotropic_predictor()
uaf_prov_expl = (
    self.adx2d[s_i, s_j]              # depth-averaged advection (x)
    + self.advua[s_i, s_j]            # depth-averaged 2D advection (x)
    - self.aru[s_i, s_j] * cor_term   # Coriolis
    + grad_term_u_expl                # explicit part of pressure gradient
    + self.drx2d[s_i, s_j]            # baroclinic pressure gradient (x)
    - self.aru[s_i, s_j] * (-self.wusurf + self.wubot)  # surface + bottom stress
)

u_star = (old_flux - 4*dt*uaf_prov_expl) / (denom + 1e-20)

The explicit pressure gradient term uses only the old and current elevation fields:

$$\text{grad\_u\_expl}=\frac{g}{4}(\overline{dy})(\bar{D})[(1-2\alpha )({\eta }^n-{\eta }_{i-1}^n)+\alpha ({\eta }^{n-1}-{\eta }_{i-1}^{n-1})+(p_a/{\rho }_{0}g-p_{a,i-1}/{\rho }_{0}g)]$$

Stage 3: Elliptic Solve for ${\eta }^{n+1}$

The continuity equation is combined with the semi-implicit momentum formulation to yield a 2D elliptic equation for ${\eta }^{n+1}$:

$$\boxed{{\displaystyle {\eta }^{n+1}-4(\mathrm{\Delta }{t}_i{)}^{2}g\alpha \mathrm{\nabla }\cdot (D\mathrm{\nabla }{\eta }^{n+1})=\text{RHS}}}$$

The right-hand side captures all explicit contributions:

$$\text{RHS}={\eta }^{n-1}-2\mathrm{\Delta }{t}_i[\mathrm{\nabla }\cdot (D{\overline{\mathbf{u}}}^{\ast })+Q_f]$$

where ${\overline{\mathbf{u}}}^{\ast }$ is the provisional barotropic velocity and $Q_f$ is the surface volume flux. This operator is discretized on the Arakawa-C grid with land masking via dum and dvm:

# From src/neuralpom/cores/implicit/implicit_barotropic.py
def apply_elliptic_operator(self, eta, d_2d):
    # H at U/V faces (Arakawa C)
    H_u_east = 0.5 * (d_2d[s_i_p1, s_j] + d_2d[s_i, s_j])
    H_u_west = 0.5 * (d_2d[s_i, s_j] + d_2d[s_i_m1, s_j])
    # ...

    # Fluxes with land mask
    flux_x_east = (H_u_east
        * (eta[s_i_p1, s_j] - eta[s_i, s_j])
        / (dx_east + 1e-20) * dy_east
        * self.dum[s_i_p1, s_j])

    div_flux = (flux_x_east - flux_x_west
                + flux_y_north - flux_y_south) / art

    coeff = 4.0 * (self.dti ** 2) * self.grav * self.alpha
    out[s_i, s_j] = eta[s_i, s_j] - coeff * div_flux
    return out

The operator $\mathcal{L}$ acting on $\eta$ can be expressed as:

$$\mathcal{L}(\eta {)}_{i,j}={\eta }_{i,j}-\frac{4(\mathrm{\Delta }{t}_i{)}^{2}g\alpha }{{\text{art}}_{i,j}}\sum _{\text{faces}}(H_{\text{face}}\cdot \mathrm{\Delta }{\eta }_{\text{face}}\cdot {\text{mask}}_{\text{face}})$$

The coefficient $4(\mathrm{\Delta }{t}_i{)}^{2}g\alpha$ arises from $g\mathrm{\Delta }{t}_i$ (acceleration) times $2\mathrm{\Delta }{t}_i$ (leapfrog) times $2\alpha$ (centering parameter).

Stage 4: CG/PCG Solver

The elliptic system $\mathcal{L}(\eta )=\text{RHS}$ is symmetric positive-definite for $\alpha >0$, making it well-suited for the Conjugate Gradient method.

Warm Start

The solver uses a linear extrapolation in time as the initial guess, which dramatically reduces iteration counts:

# Linear extrapolation: eta^{n+1} ~ 2*eta^n - eta^{n-1}
x = (2.0 * self.el - self.elb).clone()

Spectral/Jacobi-Preconditioned PCG (Accelerated Core)

The accelerated implicit core stores pre-computed grid geometry in a cache. The active V2 profile uses a spectral preconditioner; a Jacobi (diagonal) preconditioner remains selectable for reference controls:

# From src/neuralpom/cores/implicit/core.py
def _build_acceleration_cache(self):
    self._accel_cache = {
        "dy_over_dx_east": dy_east / (dx_east + 1e-20),
        "dy_over_dx_west": dy_west / (dx_west + 1e-20),
        "dx_over_dy_north": dx_north / (dy_north + 1e-20),
        "dx_over_dy_south": dx_south / (dy_south + 1e-20),
        "mask_u_east": self.dum[s_i_p1, s_j],
        "mask_u_west": self.dum[s_i, s_j],
        "mask_v_north": self.dvm[s_i, s_j_p1],
        "mask_v_south": self.dvm[s_i, s_j],
        "inv_art": 1.0 / (self.art[s_i, s_j] + 1e-20),
        "implicit_coeff": 4.0 * (self.dti ** 2) * self.grav * self.alpha,
        # ...
    }

The Jacobi preconditioner is $\mathbf{M}=\text{diag}(\mathcal{L})$, giving:

$${\mathbf{M}}_{i,j}^{-1}=\frac{1}{1+\frac{4(\mathrm{\Delta }{t}_i{)}^{2}g\alpha }{{\text{art}}_{i,j}}\sum _{\text{faces}}{H}_{\text{face}}\cdot \frac{\text{length}}{\text{width}}\cdot {\text{mask}}_{\text{face}}}$$

The PCG algorithm:

r = RHS - L(x)
z = M^{-1} r      # Jacobi preconditioner
p = z
rho_old = r^T z

for it = 1, 2, ..., max_iter:
    Ap = L(p)
    alpha = rho_old / (p^T Ap)
    x = x + alpha * p
    r = r - alpha * Ap
    if |r|^2 <= tol^2:  break
    z = M^{-1} r
    rho_new = r^T z
    beta = rho_new / rho_old
    p = z + beta * p
    rho_old = rho_new

Autograd and CG

When gradients are enabled (torch.is_grad_enabled()), the accelerated core falls back to the reference CG implementation from _ReferencePOM2KCore.solve_implicit_eta_cg. This ensures that autograd can trace through the entire CG iteration. When gradients are disabled, the accelerated PCG path uses the selected preconditioner for maximum throughput.

Stage 5: Corrector (Final Barotropic Velocities)

With ${\eta }^{n+1}$ known, the implicit part of the pressure gradient is applied to correct the provisional velocities:

# From src/neuralpom/cores/implicit/core.py,
#    _compute_barotropic_corrector()
grad_term_u_impl = 0.25 * self.grav * (dy_sum) * (d_sum) * \
    self.alpha * (self.elf[s_i, s_j] - self.elf[s_i_m1, s_j])

uaf_inner = (u_star * denom_u_star - 4.0 * dti * grad_term_u_impl) \
    / (denom_u_final + 1e-20) * self.dum[s_i, s_j]

The final velocities satisfy:

$$\begin{array}{rl}{\mathbf{u}}^{n+1}& ={\mathbf{u}}^{\ast }-g\mathrm{\Delta }{t}_i\alpha \mathrm{\nabla }{\eta }^{n+1}\\ {\mathbf{v}}^{n+1}& ={\mathbf{v}}^{\ast }-g\mathrm{\Delta }{t}_i\alpha \mathrm{\nabla }{\eta }^{n+1}\end{array}$$

Boundary Conditions for the Elliptic System

The elliptic solver enforces boundary conditions specific to each problem type:

# From src/neuralpom/cores/implicit/implicit_barotropic.py
def _enforce_eta_neumann_bc(self, eta):
    if self.iproblem == 4 or self.iproblem == 5:
        # Double Gyre (closed basin):
        # Dirichlet eta=0 on outermost T-ring
        eta[0, :] = 0.0
        eta[im_end, :] = 0.0
        eta[:, 0] = 0.0
        eta[:, jm_end] = 0.0
    else:
        # Open boundary: Neumann ghost = inner
        eta[0, :] = eta[1, :]
        eta[im_end, :] = eta[im_m1, :]
        eta[:, 0] = eta[:, 1]
        eta[:, jm_end] = eta[:, jm_m1]

These boundary conditions are enforced at every iteration of the CG solve and also applied to the gradient direction $p$ to keep the search within the correct subspace.

After the Barotropic Step

Once the barotropic solver completes, the implicit core performs a Shapiro filter on ${\eta }^{n+1}$ to suppress grid-scale noise:

def _apply_barotropic_shapiro(self, elf):
    gamma_shapiro = 0.05
    elf_smooth[s_i, s_j] = elf[s_i, s_j] + (gamma_shapiro / 4.0) * (
        elf[s_i_p1, s_j] + elf[s_i_m1, s_j] +
        elf[s_i, s_j_p1] + elf[s_i, s_j_m1] -
        4.0 * elf[s_i, s_j]
    )
    return elf_smooth * self.fsm

The standard V2-Align configuration uses gamma_shapiro=0.05. For controlled operator ablations only, NEURALPOM_IMPLICIT_SPATIAL_SHAPIRO=off sets the barotropic, velocity, and tracer Shapiro coefficients to zero together while leaving the Robert-Asselin time filter unchanged.

Then the baroclinic correction is applied. Since the implicit core uses a single barotropic step, the vertical integrals egf, utf, vtf are computed via trapezoidal averaging rather than accumulation:

# Trapezoidal: (start + end) / 2
self.egf = (self.elb + self.el) * 0.5
self.utf[...] = 0.5 * (self.uab * d_baro + self.ua * d_new)
self.vtf[...] = 0.5 * (self.vab * d_baro + self.va * d_new)

This contrasts with the explicit core's approach of accumulating $N$ substep contributions.

Time Loop Structure

The implicit time loop follows this high-level structure:

# From src/neuralpom/cores/implicit/time_loop.py (simplified)
for iint in range(1, iend + 1):
    # --- Baroclinic preparations ---
    # Wind stress, surface fluxes
    # Smagorinsky viscosity, baroclinic pressure gradient
    # Depth-averaged advection

    # --- Implicit barotropic step ---
    u_star, v_star = _barotropic_predictor_with_bc()
    fluxua, fluxva, rhs_eta = _barotropic_transport_rhs_full(u_star, v_star)
    self.elf = solve_implicit_eta_cg(rhs_eta, self.d)   # ELLIPTIC SOLVE
    self.elf = _surface_bc_after_solver(self.elf)
    self.d = self.h + self.elf
    self.uaf, self.vaf = _barotropic_corrector_with_bc(...)

    # --- Robert-Asselin filter ---
    # ...

    # --- Baroclinic correction ---
    # Fix U, V (repartition depth-average)
    # Vertical velocity
    # Turbulence (Mellor-Yamada)
    # Tracer advection and diffusion
    # 3D momentum (advu, advv, profu, profv)

Use Case: Data Generation

The accelerated implicit core provides a dedicated run_generation_day method for offline training data generation. It temporarily enables generation_fast_path mode, which bypasses the leaf-variable cloning step for maximum throughput, and captures snapshots at regular intervals:

# From src/neuralpom/cores/implicit/core.py
def run_generation_day(self, steps_per_day, capture_interval_steps=8,
                       capture_count=12):
    self.generation_fast_path = True
    self._generation_capture_interval_steps = capture_interval_steps
    self._generation_capture_limit_steps = \
        capture_interval_steps * capture_count
    self._generation_reference_records = []
    self.iend = int(steps_per_day)
    try:
        self.run_time_loop()
        return list(self._generation_reference_records)
    finally:
        self.generation_fast_path = False
        # Clean up capture state

Each snapshot captures the model state for supervised learning:

def _capture_generation_reference_snapshot(self):
    return {
        "time": float(self.time),
        "u": self.u.detach().clone(),
        "v": self.v.detach().clone(),
        "t": self.t.detach().clone(),
        "s": self.s.detach().clone(),
        "el": self.el.detach().clone(),
        "ua": self.ua.detach().clone(),
        "va": self.va.detach().clone(),
        # ...
    }

Usage Example

from neuralpom.cores.implicit import POM2K_Core

# Create an implicit core
core = POM2K_Core(
    im=250, jm=250,
    problem_id=5
)

# Initialize
core.run_initialization()

# Run time integration (implicit barotropic)
core.run_time_loop()

# Check solver statistics
print(f"CG iterations: {core._last_eta_solver_iterations}")
print(f"Final residual: {core._last_eta_solver_residual:.2e}")
print(f"Surface elevation: {core.el.min():.3f} to {core.el.max():.3f} m")

Data Generation Mode

# Generate one day of training data with 12 snapshots
records = core.run_generation_day(
    steps_per_day=96,           # dti=900s -> 96 steps/day
    capture_interval_steps=8,   # snapshot every 8 steps
    capture_count=12            # 12 snapshots per day
)

for rec in records:
    print(f"Day {rec['time']:.3f}: "
          f"u_max={rec['u'].abs().max():.3f} m/s")

Solver Tuning Parameters

Parameter	Default	Description
max_iter	200	Maximum CG/PCG iterations
tol	${10}^{-12}$ (accel) / ${10}^{-7}$ (ref)	Absolute convergence tolerance (squared residual)
eta_solver_rel_tol	${10}^{-5}$	Relative tolerance (active V2-Align core)
alpha	0.333	Semi-implicit pressure-gradient time weight
smoth	0.15	Robert-Asselin filter strength