Explicit Core

The explicit core integrates the Princeton Ocean Model using a classical leapfrog time-stepping scheme with explicit free-surface treatment. It faithfully reproduces the original POM2K dynamics while making every operation differentiable through PyTorch's autograd engine.

Architecture

The explicit core is assembled via multiple inheritance from mixin classes in src/neuralpom/cores/explicit/:

# From src/neuralpom/cores/explicit/core.py
class POM2K_Core(
    GridMixin, CasesMixin, InitializationMixin,
    BaroclinicMixin, TracersMixin, MomentumMixin,
    DiagnosticsMixin, CheckpointingMixin,
    TimeLoopMixin, BaseCore
):
    def __init__(self, im, jm, problem_id=4,
                 settings_overrides=None, case_name=None):
        self._requested_problem_id = problem_id
        self._settings_overrides = dict(settings_overrides or {})
        self.case_name = case_name
        super().__init__(im, jm)
        self._apply_settings_overrides()

Each mixin encapsulates a distinct physical or numerical component:

Mixin	Responsibility
GridMixin	Vertical sigma-coordinate generation ($z$, $zz$, $dz$, $dzz$)
CasesMixin	Problem-specific setup (double gyre, seamount, box, Bire 2025)
InitializationMixin	Top-level initialization orchestrator
BaroclinicMixin	Baroclinic pressure gradient, boundary conditions, advection
TracersMixin	Tracer advection (advt1, advt2) and vertical diffusion (proft)
MomentumMixin	3D momentum equations, vertical friction (profu, profv)
DiagnosticsMixin	Runtime diagnostics and checkpoint output
CheckpointingMixin	Model state save/restore
TimeLoopMixin	Main integration loop
BaseCore	grid dimension setup, parameter arrays, derived indices

Arakawa-C Grid

The model uses an Arakawa-C staggered grid in the horizontal and a sigma coordinate in the vertical, following the standard POM convention. The staggering is as follows:

Variable	I-index	J-index	Staggering
$\eta$ (free surface)	center	center	T-point
$u$ (zonal velocity)	face	center	U-point (east-west faces)
$v$ (meridional velocity)	center	face	V-point (north-south faces)
$T,S,\rho$	center	center	T-point
$w$ (sigma velocity)	center	center	W-point (vertical interfaces)
$K_M,K_H,q^2,q^{2}l$	center	center	T-point

The staggering leads to careful averaging in the code. For instance, when computing depth at U-points:

$$D_u(i,j)=\frac{1}{2}[D(i,j)+D(i-1,j)]$$

This is reflected in the momentum code where depth averaging appears throughout:

# Depth at U-points (from explicit momentum.py advu)
h_sum = h[s_i, s_j] + h[s_im1, s_j]
depth_new = (h_sum + etf[s_i, s_j] + etf[s_im1, s_j]) * aru[s_i, s_j]

Metric Factors

Grid cell dimensions are stored in arrays that encode the spherical geometry:

• dx: zonal grid spacing at T-points (meters)
• dy: meridional grid spacing at T-points (meters)
• art: cell area at T-points, $dx\cdot dy$
• aru: cell area at U-points, $0.25(dx(i,j)+dx(i-1,j))\cdot 0.5(dy(i,j)+dy(i-1,j))$
• arv: cell area at V-points
• dum, dvm: land-sea masks for U and V points
• fsm: land-sea mask for T-points

Vertical Coordinate

The vertical grid $\sigma$ ranges from $\sigma =0$ at the surface to $\sigma =-1$ at the bottom. The grid generator depth() in GridMixin distributes layers logarithmically near the surface and bottom, with linear spacing in between:

# From src/neuralpom/cores/explicit/grid.py
class GridMixin:
    def depth(self):
        # KDZ: layer thickness ratios (powers of 2)
        kdz_list = [1.0, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0,
                    64.0, 128.0, 256.0, 512.0, 1024.0]
        # ... generates z (layer centers), zz (interfaces),
        #     dz (layer thicknesses), dzz (interface spacing)

The physical depth of each layer center is $H\cdot {\sigma }_k$, and the layer thickness is $H\cdot d{\sigma }_k$. These are stored in the 1D arrays:

Array	Shape	Description
z	(kb,)	Sigma coordinate at layer centers (${\sigma }_k$)
zz	(kb,)	Sigma coordinate at layer interfaces
dz	(kb,)	Layer thickness in $\sigma$ space ($d{\sigma }_k$)
dzz	(kb,)	Distance between adjacent layer centers

Time Stepping: Leapfrog with Asselin Filter

The explicit core uses a three-time-level leapfrog scheme for all prognostic variables, stabilized by the Robert-Asselin time filter.

Time Level Convention

Three time levels are maintained for each prognostic variable:

Variable	Past ($n-1$)	Present ($n$)	Future ($n+1$)
Free surface	elb	el	elf
Total depth	--	d = h + el	--
Surface elevation (T-points)	etb	et	etf
Barotropic U	uab	ua	uaf
Barotropic V	vab	va	vaf
3D U	ub	u	uf
3D V	vb	v	vf
Temperature	tb	t	uf (*)
Salinity	sb	s	vf (*)
TKE ($q^2$)	q2b	q2	uf (*)
TKE $\times$ length ($q^{2}l$)	q2lb	q2l	vf (*)

(*) Note: the uf and vf arrays are reused as temporary storage for the future values of tracers and turbulence quantities, since the 3D velocity and tracer updates occur in separate phases of the time loop.

Mode Splitting

Each baroclinic step of size $\mathrm{\Delta }{t}_i$ contains $N$ barotropic substeps of size $\mathrm{\Delta }{t}_e$:

self.dti = self.dte * self.isplit    # Internal time step
self.dte2 = self.dte * 2.0           # 2 * external time step
self.dti2 = self.dti * 2.0           # 2 * internal time step
self.ispi = 1.0 / self.isplit        # 1/N
self.isp2i = 1.0 / (2.0 * self.isplit)  # 1/(2N)

The external mode loop integrates the depth-averaged equations $N$ times per baroclinic step:

# From src/neuralpom/cores/explicit/time_loop.py
for iext in range(1, self.isplit + 1):
    # Compute transport fluxes
    self.fluxua = d_avg_u * dy_avg_u * self.ua
    self.fluxva = d_avg_v * dx_avg_v * self.va

    # Explicit free-surface update
    flux_div = (fluxua[i+1] - fluxua[i] +
                fluxva[j+1] - fluxva[j])
    self.elf = self.elb + self.dte2 * (-flux_div/art - vfluxf)

    # Update barotropic momentum
    self.uaf = compute_uaf(...)
    self.vaf = compute_vaf(...)

    # Accumulate vertical integrals for baroclinic correction
    self.egf += self.el * self.ispi
    self.utf += self.ua * d_sum * self.isp2i
    self.vtf += self.va * d_sum * self.isp2i

Robert-Asselin Filter

After each barotropic substep, the Robert-Asselin filter is applied to suppress the computational mode of the leapfrog scheme:

# From src/neuralpom/cores/explicit/time_loop.py
# Robert-Asselin Time Filter
self.ua = self.ua + 0.5 * self.smoth * (
    self.uab - 2.0 * self.ua + self.uaf
)
self.va = self.va + 0.5 * self.smoth * (
    self.vab - 2.0 * self.va + self.vaf
)
self.el = self.el + 0.5 * self.smoth * (
    self.elb - 2.0 * self.el + self.elf
)

# Shift time levels
self.elb = self.el.clone()
self.el = self.elf.clone()
self.uab = self.ua.clone()
self.ua = self.uaf.clone()
self.vab = self.va.clone()
self.va = self.vaf.clone()

The filter coefficient smoth (default 0.10 for the generation phase, 0.225 per the POM standard) controls the strength of the filter. The filter strength $\nu$ modifies the leapfrog solution by:

$${\varphi }^n\leftarrow {\varphi }^n+\frac{\nu }{2}({\varphi }^{n-1}-2{\varphi }^n+{\varphi }^{n+1})$$

The cloned assignments (self.elb = self.el.clone()) break reference sharing to ensure the time-level shift is compatible with autograd (see the gradient safety page).

Momentum Equations

The 3D momentum equations are computed in advu and advv (horizontal + vertical advection terms) followed by profu and profv (vertical friction via tridiagonal solve).

U-Momentum (advu)

The right-hand side of the U-momentum equation is assembled from five terms:

$$\begin{array}{rl}\frac{\partial u}{\partial t}{|}_{\text{RHS}}=& \underset{\text{horizontal advection}}{\underbrace{\text{advx}}}+\underset{\text{vertical advection}}{\underbrace{\frac{\partial }{\partial \sigma }(wu)}}+\underset{\text{Coriolis}}{\underbrace{f\cdot {\bar{v}}^{xy}}}\\ & +\underset{\text{surface slope}}{\underbrace{g\cdot \frac{\partial \eta }{\partial x}}}+\underset{\text{baroclinic pressure gradient}}{\underbrace{\text{drhox}}}\end{array}$$

Implemented in the code as:

# From src/neuralpom/cores/explicit/momentum.py, advu()
uf[s_i, s_j, k_range] = (
    advx[s_i, s_j, k_range]      # horizontal advection
    + vert_div                     # vertical advection flux divergence
    + coriolis                     # Coriolis force
    + press_grad                   # surface pressure gradient
    + drhox[s_i, s_j, k_range]    # baroclinic pressure gradient
)

The time advancement is then:

$$u_f^{n+1}=\frac{D_u^n\cdot u_b^n-2\mathrm{\Delta }{t}_i\cdot \text{RHS}}{D_u^{n+1}}$$

where $D_u^n=\frac{1}{2}[D^n(i,j)+D^n(i-1,j)+{\eta }_b(i,j)+{\eta }_b(i-1,j)]\cdot \text{aru}$. The factor $\frac{1}{D_u^{n+1}}$ corrects for the change in water column thickness between time levels.

V-Momentum (advv)

The V-momentum equation is analogous but uses $j$-direction staggerings and the opposite Coriolis sign:

$$v_f^{n+1}=\frac{D_v^n\cdot v_b^n-2\mathrm{\Delta }{t}_i\cdot {\text{RHS}}_v}{D_v^{n+1}}$$

Vertical Friction (profu, profv)

Vertical momentum diffusion is solved implicitly via a tridiagonal system, solved with a forward sweep (Thomas algorithm) and backward substitution. The system arises from the vertical diffusion term:

$$\frac{u^{n+1}-u^n}{2\mathrm{\Delta }{t}_i}=\frac{\partial }{\partial \sigma }\left(\frac{K_M}{D^2}\frac{\partial u^{n+1}}{\partial \sigma }\right)$$

The tridiagonal coefficients are constructed as:

# From src/neuralpom/cores/explicit/momentum.py, profu()
for k in range(1, kbm1):
    num = -dti2 * (km_u[s_i, s_j, k] + umol)
    a[s_i, s_j, k-1] = num / (dz[k-1] * dzz[k-1] * dh2)
    c[s_i, s_j, k]   = num / (dz[k]   * dzz[k-1] * dh2)

Surface and bottom boundary conditions are applied as additional terms in the tridiagonal system:

• Surface: wind stress ${\tau }_s/{\rho }_0$ applied as a momentum flux
• Bottom: quadratic drag ${\tau }_b/{\rho }_0=C_D|{\mathbf{u}}_b|{\mathbf{u}}_b$

Baroclinic Pressure Gradient

The baroclinic pressure gradient $\mathrm{\nabla }{p}_{\text{bc}}$ is computed from the density field using the density Jacobian formulation in sigma coordinates. For the $x$-component:

$$\text{drhox}(k)=\text{drhox}(k-1)+\frac{g}{4}[({\sigma }_{k-1}-{\sigma }_k)(D_i+D_{i-1})(\mathrm{\Delta }{\rho }_k+\mathrm{\Delta }{\rho }_{k-1})+({\sigma }_{k-1}+{\sigma }_k)(D_i-D_{i-1})(\mathrm{\Sigma }{\rho }_k-\mathrm{\Sigma }{\rho }_{k-1})]$$

where $\mathrm{\Delta }{\rho }_k=\rho (i,k)-\rho (i-1,k)$ and $\mathrm{\Sigma }{\rho }_k=\rho (i,k)+\rho (i-1,k)$. The implementation uses torch.cumsum for efficient vectorized integration:

# From src/neuralpom/cores/explicit/momentum.py, baropg()
integral = torch.cumsum(increment, dim=2)
drhox[s_i, s_j, k_sl] = _b(drhox[s_i, s_j, 0]) + integral

The density is computed using the full Mellor (1991) equation of state, which includes pressure-dependent terms:

# From src/neuralpom/cores/explicit/momentum.py, dens()
# Seawater EOS at atmospheric pressure (Mellor 1991)
rhor = -0.157406 + 6.793952e-2 * tr - 9.095290e-3 * tr2 + \
       1.001685e-4 * tr3 - 1.120083e-6 * tr4 + 6.536332e-9 * tr5

# Salinity terms
rhor += (0.824493 - 4.0899e-3 * tr + 7.6438e-5 * tr2 -
         8.2467e-7 * tr3 + 5.3875e-9 * tr4) * sr

# Pressure correction
cr = 1449.1 + 0.0821 * p + 4.55 * tr - 0.045 * tr2 + 1.34 * (sr - 35.0)
rhor += 1.0e5 * p / (cr**2) * (1.0 - 2.0 * p / (cr**2))

Tracer Advection

Tracers (temperature $T$ and salinity $S$) are advected using one of two schemes controlled by the nadv parameter:

• nadv = 1: Centered differencing (default)
• nadv = 2: Smolarkiewicz iterative upstream scheme

Vertical diffusion of tracers is solved using the same tridiagonal approach as momentum, calling proft with the vertical diffusivity $K_H$ from the Mellor-Yamada scheme.

Initialization Procedure

The initialization sequence is:

# From src/neuralpom/cores/explicit/initialization.py
class InitializationMixin:
    def run_initialization(self):
        # 1. Generate vertical grid (unless double-gyre)
        if self.iproblem not in (3, 4, 5):
            self.dz, self.z, self.zz, self.dz, self.dzz = self.depth()

        # 2. Set up problem-specific bathymetry, forcing, and ICs
        if self.iproblem == 1:    self.seamount()
        elif self.iproblem == 2:  self.box()
        elif self.iproblem == 3:  self.file2ic()
        elif self.iproblem == 4:  self.double_gyre()
        elif self.iproblem == 5:  self.bire_2025_setup()

        # 3. Compute inertial period for ramp function
        self.period = (2.0 * pi) / abs(cor(mid)) / 86400.0

        # 4. Initialize time and copy ICs to working arrays
        self.ua = self.uab.clone()
        self.va = self.vab.clone()
        self.el = self.elb.clone()
        # ...

        # 5. Initialize turbulence fields
        self.q2b = torch.full_like(self.q2b, self.small)
        self.q2lb = self.l * self.q2b
        # ...

        # 6. Call DENS to compute initial density
        self.s, self.t, self.rho = self.dens()

Usage Example

import torch
from neuralpom.cores.explicit import POM2K_Core

# Create an explicit core for the Bire et al. (2025) double-gyre
core = POM2K_Core(
    im=250, jm=250,
    problem_id=5  # Bire et al. (2025) Double Gyre
)

# Initialize grid, bathymetry, initial conditions
core.run_initialization()

# Run the time integration
# (Uses dte=120s, isplit=30 -> dti=3600s)
core.run_time_loop()

# Access model state
print(f"Surface elevation range: "
      f"{core.el.min():.3f} to {core.el.max():.3f} m")
print(f"Max velocity: {core.u.abs().max():.3f} m/s")

With Autograd (Training)

When using for training with neural network parameterizations:

# Enable gradient tracking on initial conditions
core.u.requires_grad_(True)
core.v.requires_grad_(True)
core.t.requires_grad_(True)
core.s.requires_grad_(True)

# Or inject a neural network correction
from neuralpom.models import SubgridModel
model = SubgridModel()

# The time loop will propagate gradients through
# all physics operations
core.run_time_loop()

# Compute a loss and backpropagate
loss = compute_loss(core)
loss.backward()