Time Stepping

Overview

NeuralPOM employs a split-explicit time-stepping scheme that separates the fast barotropic (external) mode from the slower baroclinic (internal) 3D mode. This approach, inherited from the Princeton Ocean Model (POM), uses a leapfrog time discretization for the internal mode with an Asselin-Robert time filter to control the computational mode.

Time Scales and Step Splitting

The model operates with two distinct time steps:

Parameter	Symbol	Description	Typical Value (Double Gyre)
External time step	dte	$\mathrm{\Delta }{t}_e$	30 s
Internal time step	dti	$\mathrm{\Delta }{t}_i$	900 s
Splitting ratio	isplit	$N_{split}=\mathrm{\Delta }{t}_i/\mathrm{\Delta }{t}_e$	30

The splitting ratio is defined as:

$$\mathrm{\Delta }{t}_i=N_{split}\cdot \mathrm{\Delta }{t}_e$$

where $N_{split}$ is stored as self.isplit in the code. Derived quantities include:

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   # inverse splitting ratio
self.isp2i = 1.0 / (2 * self.isplit)  # half-inverse splitting ratio

Leapfrog Time Discretization

The basic time integration uses a three-time-level leapfrog scheme centered at time level $n$:

$$\frac{\partial \varphi }{\partial t}\approx \frac{{\varphi }^{n+1}-{\varphi }^{n-1}}{2\mathrm{\Delta }t}$$

For a generic prognostic equation ${\partial }_t\varphi =F(\varphi )$, the leapfrog discretization is:

$${\varphi }^{n+1}={\varphi }^{n-1}+2\mathrm{\Delta }t\cdot F({\varphi }^n)$$

The leapfrog scheme is formally second-order accurate in time and is computationally efficient since it is explicit and requires only one evaluation of the right-hand side per time step.

Three Time-Level Storage

NeuralPOM maintains three time levels for each prognostic variable:

Time Level	Variable Suffix	Meaning
$n-1$ (backward)	b (e.g., ub, vb, tb, elb)	Previous time step
$n$ (current)	none (e.g., u, v, t, el)	Current time step
$n+1$ (forward)	f (e.g., uf, vf, tf, elf)	Next time step

Computational Mode and the Asselin-Robert Filter

The leapfrog scheme has a well-known weakness: it supports a spurious computational mode that oscillates with period $2\mathrm{\Delta }t$ and can grow over time. To suppress this mode, NeuralPOM applies the Asselin-Robert time filter after each time step:

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

where $\nu$ is the filter coefficient stored as self.smoth. Typical values are:

• Explicit core: smoth = 0.10
• Implicit core: smoth = 0.15 (stronger damping at larger internal time steps)

The filter is applied in the code as:

# Barotropic variables (depth-averaged)
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)

# Baroclinic variables (3D fields)
u_filtered = self.u + 0.5 * self.smoth * (self.uf + self.ub - 2.0 * self.u - tps_u)
v_filtered = self.v + 0.5 * self.smoth * (self.vf + self.vb - 2.0 * self.v - tps_v)

For 3D velocity fields, the depth-mean is subtracted before filtering and restored afterward, following standard POM practice, to ensure that the depth-averaged transport is conserved during the filter step.

After filtering, the three time levels are shifted:

$${\varphi }^{n-1}\leftarrow {\varphi }^n\text{ (filtered)},{\varphi }^n\leftarrow {\varphi }^{n+1}$$

CFL Constraints

External Mode CFL

The fast barotropic mode involves surface gravity waves with phase speed $c=\sqrt{gH}$. The CFL condition for the external time step is:

$$\mathrm{\Delta }{t}_e<\frac{min(\mathrm{\Delta }x,\mathrm{\Delta }y)}{\sqrt{g{H}_{max}}}$$

For the Double Gyre configuration ($\mathrm{\Delta }x\approx 7$ km, $H_{max}=2000$ m):

$$\sqrt{g{H}_{max}}\approx \sqrt{9.81\times 2000}\approx 140\text{ m/s},\mathrm{\Delta }{t}_e<\frac{7000}{140}\approx 50\text{ s}$$

The explicit core uses $\mathrm{\Delta }{t}_e=30$ s, providing a safety margin. The implicit core can use larger $\mathrm{\Delta }{t}_e$ values (e.g., 120 s for the Bire et al. configuration) since the free-surface is treated semi-implicitly, removing the external gravity wave CFL constraint.

Internal Mode CFL

The 3D internal mode is constrained by the maximum advective velocity and internal wave speed:

$$\mathrm{\Delta }{t}_i<\frac{min(\mathrm{\Delta }x,\mathrm{\Delta }y)}{|u_{max}|+c_{int}}$$

where $c_{int}\approx 2-3$ m/s is the first-mode baroclinic wave speed. For velocities up to 1 m/s:

$$\mathrm{\Delta }{t}_i<\frac{7000}{3}\approx 2333\text{ s}$$

The standard internal time step of $\mathrm{\Delta }{t}_i=900$ s (30 external steps at 30 s each) is well within this limit.

Barotropic-Baroclinic Time Step Splitting

Explicit Core: Sub-cycled External Mode

In the explicit core (src/neuralpom/cores/explicit/time_loop.py), the external mode is sub-cycled isplit times for each internal time step. The structure is:

for iint in range(1, iend + 1):           # Internal mode loop
    # 1. Apply surface forcing, boundary conditions
    # 2. Compute baroclinic terms (advection, pressure gradient, viscosity)

    for iext in range(1, isplit + 1):      # External mode sub-cycle
        # 2a. Compute depth-averaged fluxes
        # 2b. Update free surface (elf) using external time step

        # Every ispadv-th step: compute depth-averaged advection

        # 2c. Update depth-averaged velocity (uaf, vaf)
        # 2d. Apply boundary conditions
        # 2e. Time filter at final sub-steps

    # 3. Accumulate egf, utf, vtf (trapezoidal averages over external steps)
    # 4. Correct 3D velocity (remove barotropic, add averaged barotropic)
    # 5. Compute vertical velocity
    # 6. Solve turbulence (q2, q2l, km, kh)
    # 7. Advect and diffuse T/S
    # 8. Update density
    # 9. Solve 3D momentum with vertical diffusion
    # 10. Apply Asselin filter to all fields

The free-surface elevation elf is updated at each external sub-step using the 2D continuity equation:

$${\eta }_f(i,j)={\eta }_b(i,j)-\frac{2\mathrm{\Delta }{t}_e}{art(i,j)}[{\mathrm{\Phi }}_u(i+1,j)-{\mathrm{\Phi }}_u(i,j)+{\mathrm{\Phi }}_v(i,j+1)-{\mathrm{\Phi }}_v(i,j)]-2\mathrm{\Delta }{t}_e\cdot vfluxf(i,j)$$

where ${\mathrm{\Phi }}_u$, ${\mathrm{\Phi }}_v$ are the depth-integrated transports at U- and V-faces.

The accumulated external mode contributions egf, utf, vtf are computed as trapezoidal averages over all isplit sub-steps:

if iext != self.isplit:
    self.egf = self.egf + self.el * self.ispi
    self.utf[1:im, :] = self.utf[1:im, :] + self.ua[1:im, :] * d_sum_u * self.isp2i
    self.vtf[:, 1:jm] = self.vtf[:, 1:jm] + self.va[:, 1:jm] * d_sum_v * self.isp2i

These averages are then used to correct the 3D velocity:

$$u_{corrected}(i,j,k)=u(i,j,k)-⟨u⟩+\frac{utb+utf}{D(i,j)+D(i-1,j)}$$

Implicit Core: Single-Step Semi-Implicit External Mode

In the implicit core (src/neuralpom/cores/implicit/time_loop.py), the external mode is treated in a single step using a fractional step (predictor-corrector) approach:

for iint in range(1, iend + 1):           # Internal mode loop
    # 1. Apply surface forcing, boundary conditions
    # 2. Compute baroclinic terms

    # 3. Predictor step: compute provisional velocities U*, V* (explicit momentum
    #    minus terms depending on unknown future free surface)

    # 4. Form RHS of elliptic equation for eta^{n+1}

    # 5. Solve elliptic equation: L(eta) = rhs using CG/PCG

    # 6. Corrector step: update uaf, vaf using solved eta^{n+1}

    # 7. Apply Asselin filter (single step)

    # 8. Compute trapezoidal averages (egf, utf, vtf)

    # 9-14. Same baroclinic sequence as explicit core

The predictor step computes provisional depth-averaged velocities $U^{\ast }$, $V^{\ast }$ that exclude the implicit pressure gradient correction:

$$U^{\ast }=\frac{D_b\cdot U_b-4\mathrm{\Delta }{t}_i\cdot F_{expl}}{D^n+D^{n-1}}$$

where $F_{expl}$ includes all explicit forcing terms (advection, Coriolis, explicit part of pressure gradient, wind stress, bottom friction) but excludes the $\alpha$-weighted implicit gradient of ${\eta }^{n+1}$.

The free surface is then solved implicitly (see Elliptic Solver), and the corrector applies the missing pressure gradient:

$$U^{n+1}=\frac{D_b\cdot U^{\ast }-4\mathrm{\Delta }{t}_i\cdot g\cdot \alpha \cdot \mathrm{\nabla }{\eta }^{n+1}\cdot D\cdot A}{D^{n+1}}$$

The iint Counter and Step Management

The simulation is driven by the iint loop counter:

for iint in range(1, self.iend + 1):
    self.iint = iint
    self.time = self.dti * iint * 1.0 / 86400.0 + self.time0

Key parameters:

• iend: Total number of internal time steps, computed as:

  $$iend=max(\lfloor \frac{\mathrm{days}\times 86400}{\mathrm{\Delta }{t}_i}+0.5\rfloor ,2)$$

• time0: Initial model time offset in days (for restart runs).

• time: Current model time in days since the simulation start.

• iprint: Print interval in internal steps.

Ramp Function

A smooth spin-up is implemented using a linear ramp for the first period days:

if self.lramp:
    self.ramp = self.time / self.period
    if self.ramp > 1.0:
        self.ramp = 1.0

The ramp is used to gradually introduce surface forcing, reducing inertial shock during initialization.

Time Stepping in the Implicit Core for GPU/Autograd

The implicit core is designed to be compatible with PyTorch autograd for neural network training. Unlike the explicit core, which sub-cycles the external mode, the implicit core solves the barotropic mode in a single algebraic step. This avoids the need to unroll isplit external iterations through autograd, significantly reducing memory usage during backpropagation.

Key differences in the implicit core time loop:

1. No explicit sub-cycling: The external mode is solved as one linear system.
2. Clone-once pattern: All state updates use .clone() to create new tensors, avoiding in-place operations that break autograd.
3. Pre-allocated scratch buffers: The accelerated implicit core (POM2K_Core) reuses work arrays (_cg_work) to reduce memory allocations.
4. Shapiro spatial filter: A 3D horizontal Shapiro filter (strength $\gamma =0.05$) is applied to velocity and tracer fields to suppress 2$\mathrm{\Delta }x$ checkerboard noise that can arise from the implicit W velocity under large internal time steps.