Implicit Core

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 Jacobi-preconditioned PCG, cached geometric terms, and reusable scratch buffers:

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
              ├── 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,
    Jacobi-preconditioned PCG, 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-6

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()

Jacobi-Preconditioned PCG (Accelerated Core)

The accelerated implicit core stores pre-computed grid geometry in a cache and uses a Jacobi (diagonal) preconditioner:

# 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 (e.g., during data generation), the accelerated Jacobi-preconditioned PCG is used 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

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}^{-6}$	Relative tolerance (accelerated core)
alpha	0.225	Implicitness parameter (same as Asselin filter coefficient)
smoth	0.10	Robert-Asselin filter strength