Hybrid AI-Physics Framework

Overview

NeuralPOM implements a hybrid AI-physics paradigm for ocean simulation. The core idea is simple: a classical physics-based ocean model (the "physics core") provides the fundamental dynamical structure, while a learned neural corrector compensates for the inevitable truncation errors that arise from discretization and subgrid-scale processes. This approach combines the physical consistency of equation-based solvers with the flexibility and expressiveness of deep neural networks.

In contrast to pure-ML approaches (which must learn all of fluid dynamics from scratch) and pure-physics approaches (which are limited by their fixed resolution and parameterizations), the hybrid paradigm lets each component do what it is best at:

Approach	Strengths	Weaknesses
Pure Physics	Physically consistent, conserves invariants	Fixed truncation error, expensive subgrid parameterizations
Pure ML	Fast inference, learns from data	No physical guarantees, can violate conservation laws
Hybrid (NeuralPOM)	Physics provides structure; ML corrects errors	Requires differentiable physics core, more complex training

Architecture

The NeuralPOM hybrid pipeline processes initial conditions through alternating physics and AI stages:

Initial Condition (12 prognostic fields)
        │
        ▼
┌───────────────────────┐
│  pack_pom_batch()     │  ← Pack 12 fields into a 101-channel 2D tensor
│  StateNormalizer      │  ← Subtract T/S climatological mean
└───────────────────────┘
        │
        ▼
┌───────────────────────┐
│  Encoder CorrectorNet │  ← U-Net: correct initial state before physics
│  (optional)           │
└───────────────────────┘
        │
        ▼
┌───────────────────────┐
│  Physics Core         │  ← Multi-step POM integration
│  (run_time_loop)      │    Gradient flows through each step
└───────────────────────┘
        │
        ▼
┌───────────────────────┐
│  Decoder CorrectorNet │  ← U-Net: correct state after physics
│  (optional)           │
└───────────────────────┘
        │
        ▼
   Prediction (12 prognostic fields)
        │
        ▼
┌───────────────────────┐
│  Supervision Loss     │  ← Compare against reference trajectory
│  (MSE / multiscale)   │
└───────────────────────┘

The pipeline runs for n_steps correction cycles. Within each cycle, the encoder corrects the state before 2 physics sub-steps (each 600 s), then the decoder corrects the state after. The total forward horizon is n_steps * 1200 s.

Encoder and Decoder Design

NeuralPOM optionally uses two independent CorrectorNet instances:

• Encoder (enc): Applied before the physics step. Its role is to adjust the initial state so that the physics model can produce a more accurate trajectory. Think of it as learning an optimal "pre-processing" nudging.

• Decoder (dec): Applied after the physics step. Its role is to correct the physics output, compensating for the model's truncation error. Think of it as a learned "post-processing" debiasing.

Both networks are architecturally identical U-Nets but are trained with separate parameters. You can enable either, both, or neither via the --use_encoder and --use_decoder flags.

Prognostic Fields as a 101-Channel Input

The ocean state is represented by 12 prognostic fields, which are packed into a single 101-channel 2D tensor for input to the CorrectorNet. The packing is defined by PACK_ORDER in constants.py:

Field	Description	Channels per sample	Dimensionality
u	Zonal velocity (baroclinic)	16 (levels 0-15)	3D
ub	Zonal velocity (barotropic)	16 (levels 0-15)	3D
v	Meridional velocity (baroclinic)	16 (levels 0-15)	3D
vb	Meridional velocity (barotropic)	16 (levels 0-15)	3D
t	Temperature	16 (levels 0-15)	3D
s	Salinity	16 (levels 0-15)	3D
el	Sea surface height	1	2D
ua	Depth-averaged zonal velocity	1	2D
uab	Depth-averaged barotropic zonal velocity	1	2D
va	Depth-averaged meridional velocity	1	2D
vab	Depth-averaged barotropic meridional velocity	1	2D

Total channels: $6\times 16+5\times 1=101$.

Each 3D field is reshaped from (H, W, K) to (K, H, W) so that the vertical dimension becomes the channel dimension. Each 2D field contributes 1 channel. All fields are concatenated along the channel axis to form a (B, 101, H, W) tensor.

StateNormalizer

The StateNormalizer handles preprocessing and postprocessing of the state tensors. It is initialized from a precomputed mean state file containing climatological means for temperature t and salinity s.

normalizer = StateNormalizer.from_path(mean_path, device)

# Normalize: subtracts T/S mean; passes other fields through unchanged
normalizer.normalize(tensor, "t")    # → tensor - t_mean
normalizer.normalize(tensor, "s")    # → tensor - s_mean
normalizer.normalize(tensor, "u")    # → tensor (no change)

# Denormalize: identity operation (no inverse transform needed)
normalizer.denormalize(tensor, "t")  # → tensor

Currently, only T and S are mean-subtracted. Velocity fields and sea surface height pass through unchanged. The denormalization is an identity operation, meaning the network learns to predict corrections in the original physical units.

pack_pom_batch and apply_correction_to_pom

These two functions form the interface between the POM physics objects and the neural network tensors:

pack_pom_batch(poms, normalizer, kb=16)

Collects prognostic fields from a batch of POM core instances, normalizes them via StateNormalizer, and packs them into a (B, 101, H, W) tensor:

# Input: list of POM core objects
poms = [...]

# Pack into network-ready tensor
network_input = pack_pom_batch(poms, normalizer)  # shape: (B, 101, H, W)

apply_correction_to_pom(pom, correction, normalizer, scale_factor=1e-3, kb=16)

Unpacks a (101, H, W) correction tensor and applies it back to the POM core fields, scaled by scale_factor:

# Apply network correction to physics state
apply_correction_to_pom(pom, correction, normalizer, scale_factor=1e-3)

The correction tensor is unpacked in the same order as PACK_ORDER. Each 3D field correction has shape (kb, H, W) (permuted to (H, W, kb) before addition), and each 2D field correction has shape (1, H, W) (squeezed to (H, W)). The correction is multiplied by scale_factor before addition, providing a hyperparameter to control the magnitude of the learned corrections.

Note that some fields share corrections:

• t correction is also applied to tb
• s correction is also applied to sb
• el correction is also applied to elb, et, etb, etf
• ua correction is also applied to uaf
• va correction is also applied to vaf

Physics-in-the-Loop Training

NeuralPOM's training is differentiable end-to-end: the gradient of the supervision loss flows backward through all physics steps and both CorrectorNets. This means the encoder and decoder networks learn to produce corrections that are optimal for the full multi-step trajectory, not just the immediate next step.

The training loop consists of:

1. Borrow physics cores: Reuse cores from a CorePool or instantiate fresh ones for each batch
2. Inject initial conditions: Set the initial ocean state via inject_core_state
3. Run multi-step physics: For each correction cycle:
   • Optionally apply encoder CorrectorNet
   • Run 2 physics sub-steps (600 s each, total 1200 s per cycle)
   • Optionally apply decoder CorrectorNet
4. Compute supervision loss: Compare the predicted trajectory against the reference
5. Backward pass: Gradient flows through all physics operations and both networks
6. Update: Gradient clipping + Adam optimizer step

This "physics-in-the-loop" approach is what distinguishes NeuralPOM from simple post-processing ML methods. The networks learn to cooperate with the physics model, not just to fix its output in isolation.