NeuralPOMA Differentiable Princeton Ocean Model

Hybrid AI-Physics ocean simulation with end-to-end differentiable training

Dual Dynamic Cores

Full-featured explicit and implicit free-surface POM cores, implemented in PyTorch with end-to-end differentiability for gradient-based optimization.

🤖

Hybrid AI-Physics Training

U-Net CorrectorNet learns to compensate for numerical truncation errors. DDP training across 8×A100 GPUs with multi-step physics-in-the-loop.

🧪

Rich Case Library

Double Gyre, Seamount, and Box configurations spanning both explicit and implicit cores, with ensemble simulation data generation pipelines.

📐

Rigorous Numerics

Arakawa-C staggered grid, leapfrog time stepping with Asselin filter, Mellor-Yamada 2.5 turbulence closure, and conjugate gradient elliptic solvers.

🔬

Multi-Scale Supervision

Pointwise, multiscale-lowpass, and horizon-weighted supervision strategies with configurable loss component weighting for 12 prognostic fields.

🚀

Production Ready

Config-driven execution via JSON, DDP training with NCCL backend, checkpoint save/resume, and comprehensive inference pipelines.

What is NeuralPOM?

NeuralPOM (Neural Princeton Ocean Model) is a research framework that integrates deep learning with classical ocean modeling. It wraps the Princeton Ocean Model (POM) — a three-dimensional, primitive-equation, Boussinesq ocean circulation model — in a fully differentiable PyTorch implementation with two complementary dynamic cores:

Explicit Core: Explicit free-surface formulation with leapfrog time stepping, suitable for high-frequency variability and process studies.
Implicit Core: Implicit free-surface formulation using conjugate gradient / preconditioned conjugate gradient (CG/PCG) solvers, enabling longer time steps and stable barotropic mode integration.

On top of these differentiable physics cores, NeuralPOM trains a CorrectorNet (a U-Net architecture) to predict and compensate for numerical truncation errors accumulated during multi-step integration. This hybrid AI-physics paradigm enables accurate long-term forecasts while maintaining physical consistency.

Quick Start

bash

git clone https://github.com/ChiyodaMomo01/NeuralPOM.git
cd NeuralPOM
pip install -r requirements.txt

# Run a Double Gyre simulation with the explicit core
python -m neuralpom.cases.explicit.double_gyre

For detailed installation and training instructions, see the Quick Start guide.

Citation

If you use NeuralPOM in your research, please cite:

bibtex

@software{neuralpom2025,
  author = {Shu, Ruiqi and Liu, Guangliang},
  title = {NeuralPOM: A Differentiable Princeton Ocean Model with Hybrid AI-Physics Training},
  year = {2025},
  url = {https://github.com/ChiyodaMomo01/NeuralPOM}
}