Dynamic Core Architecture Overview

NeuralPOM now organizes dynamics as a small registry of core families rather than a single monolithic core. The production hydrostatic POM families remain explicit and implicit, while specialized families isolate coastal-tide, non-hydrostatic, and two-dimensional geostrophic benchmarks.

This keeps new scientific cases from accumulating case-specific patches inside the main implicit core. A case should enter through the registry and a narrow adapter unless it truly changes shared hydrostatic POM numerics.

Core Families

Core	Module	Main State	Primary Use Case
Explicit	neuralpom.cores.explicit	3D hydrostatic POM	Small external time step, standard POM-style free surface, fast-path/training consistency checks
Implicit	neuralpom.cores.implicit	3D hydrostatic POM	Large barotropic time step, implicit free-surface solve, default training and long-rollout workflows
Implicit JZB	neuralpom.cores.implicit_jzb	3D hydrostatic POM with coastal tide adapters	Jiaozhou Bay realistic-topography M2 tide reproduction
Non-hydrostatic	neuralpom.cores.nonhydrostatic	3D non-hydrostatic benchmark state	Taylor-Green, Rayleigh-Benard, Kelvin-Helmholtz, and lock-exchange validation
Geostrophic 2D	neuralpom.cores.geostrophic_2d	2D periodic u, v, p	Double-periodic geostrophic turbulence and 2D hybrid experiments

Hydrostatic POM Design

The explicit and implicit hydrostatic cores are organized as Python mixin classes that compose into a single POM2K_Core class via multiple inheritance:

Core	Module	Barotropic Integrator	Time Step	Primary Use Case
Explicit	neuralpom.cores.explicit	Leapfrog with Asselin filter	$\mathrm{\Delta }{t}_e\approx 30\text{--}120\text{s}$	High-frequency dynamics, CFL-sensitive regimes
Implicit	neuralpom.cores.implicit	Semi-implicit with CG/PCG elliptic solver	$\mathrm{\Delta }{t}_i\approx 900\text{s}$	Large time steps, training workflows, stable long integrations

Both cores trace their lineage to the classic POM2K Fortran code, translated into PyTorch with careful attention to automatic differentiation compatibility. The code uses a mixin architecture where each numerical component (momentum, tracers, baroclinic pressure gradient, turbulence, grid, time loop) is a separate class, and the final POM2K_Core inherits from all of them:

# 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):
        ...

Core Factory

Core selection is handled at runtime through a lightweight registry in core_factory.py. This design allows training scripts and inference pipelines to switch between cores by changing a single string argument:

# From src/neuralpom/core_factory.py
from neuralpom.cores.explicit import POM2K_Core as ExplicitPOM2KCore
from neuralpom.cores.geostrophic_2d import POM2K_Core as Geostrophic2DCore
from neuralpom.cores.implicit import POM2K_Core as ImplicitPOM2KCore
from neuralpom.cores.implicit_jzb import POM2K_Core as ImplicitJZBPOM2KCore
from neuralpom.cores.nonhydrostatic import POM2K_Core as NonhydrostaticPOM2KCore

CORE_REGISTRY = {
    "implicit": ImplicitPOM2KCore,
    "implicit_jzb": ImplicitJZBPOM2KCore,
    "nonhydrostatic": NonhydrostaticPOM2KCore,
    "geostrophic_2d": Geostrophic2DCore,
    "explicit": ExplicitPOM2KCore,
}

def resolve_core_class(name: str):
    if name not in CORE_REGISTRY:
        raise ValueError(
            f"Unknown core '{name}'. "
            f"Available: {sorted(CORE_REGISTRY)}"
        )
    return CORE_REGISTRY[name]

Usage in training and inference:

from neuralpom.core_factory import resolve_core_class

CoreClass = resolve_core_class("implicit")  # or "explicit", "implicit_jzb", "nonhydrostatic", "geostrophic_2d"
core = CoreClass(im=250, jm=250, problem_id=5)
core.run_initialization()

Shared Physics

Both cores solve the same system of Boussinesq hydrostatic primitive equations. The governing equations in sigma coordinates are:

Momentum equations:

$$\frac{\partial (Du)}{\partial t}+\mathrm{\nabla }\cdot (D\mathbf{v}u)-fDv=-D\frac{\partial }{\partial x}\left(\frac{p}{{\rho }_0}\right)-gD\frac{\partial \eta }{\partial x}+\frac{\partial }{\partial \sigma }\left(\frac{K_M}{D}\frac{\partial u}{\partial \sigma }\right)+{\mathcal{F}}_u$$$$\frac{\partial (Dv)}{\partial t}+\mathrm{\nabla }\cdot (D\mathbf{v}v)+fDu=-D\frac{\partial }{\partial y}\left(\frac{p}{{\rho }_0}\right)-gD\frac{\partial \eta }{\partial y}+\frac{\partial }{\partial \sigma }\left(\frac{K_M}{D}\frac{\partial v}{\partial \sigma }\right)+{\mathcal{F}}_v$$

Continuity:

$$\frac{\partial \eta }{\partial t}+\mathrm{\nabla }\cdot (D\overline{\mathbf{v}})=0$$

Tracer transport:

$$\frac{\partial (DT)}{\partial t}+\mathrm{\nabla }\cdot (D\mathbf{v}T)=\frac{\partial }{\partial \sigma }\left(\frac{K_H}{D}\frac{\partial T}{\partial \sigma }\right)+{\mathcal{F}}_T$$

where $D=h+\eta$ is the total water depth, $h$ is the resting depth, $\eta$ is the free-surface elevation, $f$ is the Coriolis parameter, $K_M$ and $K_H$ are vertical eddy viscosity and diffusivity from the Mellor-Yamada scheme, ${\mathcal{F}}_u,{\mathcal{F}}_v,{\mathcal{F}}_T$ represent horizontal diffusion terms, and $\overline{\mathbf{v}}$ is the depth-averaged velocity.

Both cores share:

• Arakawa-C staggered grid: $u$ on east-west faces, $v$ on north-south faces, $\eta$, $T$, $S$, $\rho$ at cell centers
• Sigma vertical coordinate: terrain-following layers with logarithmic stretching near surface and bottom
• Mellor-Yamada 2.5 turbulence closure: prognostic $q^2$ (twice TKE) and $q^{2}l$ (TKE times length scale)
• Mellor (1991) equation of state: nonlinear $\rho (T,S,p)$
• Smagorinsky lateral viscosity: grid-scale-dependent horizontal mixing
• Robert-Asselin time filter: controls the leapfrog computational mode with filter coefficient $\alpha =0.225$

Key Architectural Difference: Mode Splitting

The fundamental difference between the two cores lies in how they handle the barotropic-baroclinic mode splitting.

Explicit Core

The external (barotropic) mode is integrated explicitly using small time steps $\mathrm{\Delta }{t}_e=\mathrm{\Delta }{t}_i/N$, where $N$ is the splitting ratio (isplit, typically 30):

• The free-surface equation is stepped forward explicitly using leapfrog time differencing
• $N$ external substeps are taken for each baroclinic step
• The pressure gradient $\mathrm{\nabla }\eta$ is computed using a weighted average of time levels (controlled by $\alpha$)
• Vertical integrals of 3D fields are accumulated over substeps for the baroclinic correction

$${\eta }^{n+1}={\eta }^{n-1}+2\mathrm{\Delta }{t}_e[-\mathrm{\nabla }\cdot (\overline{D}\overline{\mathbf{v}})-Q_f]$$

Implicit Core

The external mode uses a semi-implicit fractional step method with a single large step per baroclinic timestep:

1. Predictor step: compute provisional barotropic velocities ${\mathbf{u}}^{\ast }$, ${\mathbf{v}}^{\ast }$ using explicit forcing terms

2. Elliptic solve: compute ${\eta }^{n+1}$ by solving a 2D elliptic equation

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

3. Corrector step: correct velocities using the implicit pressure gradient from ${\eta }^{n+1}$

The elliptic system is solved using Conjugate Gradient (CG) or Jacobi-preconditioned PCG. See the implicit core page for details.

When to Use Each Core

Use the Explicit Core when:

• You need to resolve high-frequency barotropic dynamics (tides, storm surges)
• The CFL condition from surface gravity waves is not a bottleneck
• You want a direct, transparent time-stepping scheme with minimal numerical dispersion
• You are performing short integrations (days to weeks) at fine resolution
• You need to benchmark results against the standard POM

Use the Implicit Core when:

• You are training neural network subgrid parameterizations with backpropagation through multi-year integrations
• You need stable long integrations (years to decades) without restrictive CFL limits
• You want to amortize the cost of the elliptic solver through the CG/PCG warm-start mechanism
• You are using the run_generation_day pipeline for offline training data generation
• You want the performance benefits of the accelerated implicit core (cached geometry, Jacobi-preconditioned PCG)

References

• Blumberg, A. F., & Mellor, G. L. (1987). A description of a three-dimensional coastal ocean circulation model. Three-Dimensional Coastal Ocean Models, 4, 1-16.
• Mellor, G. L. (1991). An equation of state for numerical models of oceans and estuaries. Journal of Atmospheric and Oceanic Technology, 8(4), 609-611.
• Mellor, G. L., & Yamada, T. (1982). Development of a turbulence closure model for geophysical fluid problems. Reviews of Geophysics, 20(4), 851-875.