Data Generation and Pipeline

NeuralPOM's training data pipeline consists of three stages: (1) ensemble simulation at high resolution using the physics core, (2) conservative downsampling to the coarse training grid, and (3) loading and batching via the POMCheckpointDataset class.

Pipeline Overview

┌─────────────────────────────────────────────────────────────────┐
│  Stage 1: Ensemble Generation (High Resolution)                 │
│  ┌──────────┐  ┌──────────┐        ┌──────────┐                │
│  │ Member 0 │  │ Member 1 │  ...   │ Member 7 │  8 GPUs         │
│  │ GPU 0    │  │ GPU 1    │        │ GPU 7    │  multiprocess   │
│  └────┬─────┘  └────┬─────┘        └────┬─────┘                │
│       │              │                   │                       │
│       ▼              ▼                   ▼                       │
│  354×322 grid   354×322 grid       354×322 grid                  │
│  20-year sim    20-year sim        20-year sim                   │
│       │              │                   │                       │
│       ▼              ▼                   ▼                       │
│  checkpoints/    checkpoints/        checkpoints/                │
│  reference_      reference_          reference_                  │
│  data_packed/    data_packed/        data_packed/                │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│  Stage 2: Conservative Downsampling                             │
│                                                                  │
│  354×322 ──── area-weighted average (factor=4) ───► 90×82       │
│                                                                  │
│  checkpoints/      ──► checkpoints_lowres_consistent/            │
│  reference_data_   ──► reference_data_lowres_consistent/         │
│  packed/                                                         │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│  Stage 3: Dataset Loading (Training)                            │
│                                                                  │
│  POMCheckpointDataset                                            │
│  ├── checkpoint_Day{N}.npz    (input state)                      │
│  └── ref_Day{N+1}_seg{XX}.npz (target sequence)                  │
│                                                                  │
│  DistributedSampler ──► DataLoader (DDP)                         │
└─────────────────────────────────────────────────────────────────┘

Stage 1: Ensemble Generation

The script scripts/data_generation/generate_double_gyre_ensemble.py drives the ensemble simulation. It spawns 8 worker processes, each bound to a GPU, and runs an independent Double Gyre simulation (problem_id = 4) with small ensemble perturbations.

Configuration

Key environment variables:

Variable	Default	Description
NEURALPOM_OUTPUT_ROOT	./simulation_ensemble_20y	Output directory
NEURALPOM_TOTAL_DAYS	7300 (20 years)	Total simulation days
NEURALPOM_NUM_MEMBERS	8	Number of ensemble members
NEURALPOM_GENERATION_CORE	implicit	Core variant (implicit or explicit)
NEURALPOM_REF_LAYOUT	packed	Reference file layout (packed or legacy)
NEURALPOM_ENABLE_DIAGNOSTICS	0	Enable verbose runtime diagnostics

High-Resolution Grid

The ensemble runs on a fine grid of im=354, jm=322 with kb=21 vertical levels. This resolution resolves mesoscale features that the coarse 90 x 82 training grid cannot capture directly, providing high-quality reference trajectories for supervision.

Ensemble Perturbation

To generate diverse training data while sharing the same large-scale statistics, each member receives a small random perturbation at initialization:

noise_scale = 1e-4
noise_u = torch.randn_like(model.u) * noise_scale
noise_v = torch.randn_like(model.v) * noise_scale

model.u = model.u + noise_u * model.dum.unsqueeze(2)  # Mask over land
model.v = model.v + noise_v * model.dvm.unsqueeze(2)
model.ub = model.ub + noise_u * model.dum.unsqueeze(2)  # Update history too
model.vb = model.vb + noise_v * model.dvm.unsqueeze(2)

The random seed is derived from the member ID: rng_seed = 42 + rank * 999.

Daily Schedule

Each simulated day follows a fixed schedule:

1. First 4 hours (reference period): Integrate in 20-minute sub-steps, saving high-frequency reference data at each sub-step (12 segments per day).
2. Remaining 20 hours (bulk integration): Integrate in a single block without reference output.
3. End of day: Save a checkpoint containing the full model state.

Day N:
├── 00:00-04:00: 12 x 20-min segments → ref_Day{N}_seg00.npz ... seg11.npz
├── 04:00-24:00: Bulk integration (no reference output)
└── 24:00:         checkpoint_Day{N}.npz

Output Structure

For each member, the output directory is organized as:

simulation_ensemble_20y/
└── member_{id}_gpu{id}/
    ├── checkpoints/
    │   ├── checkpoint_Day0000.npz
    │   ├── checkpoint_Day0001.npz
    │   └── ...
    ├── reference_data_packed/       # if NEURALPOM_REF_LAYOUT=packed
    │   ├── ref_Day0000.npz          # (time=12, ...) packed file
    │   └── ...
    └── reference_data/              # if NEURALPOM_REF_LAYOUT=legacy
        ├── ref_Day0000_seg00.npz
        ├── ref_Day0000_seg01.npz
        └── ...

Stage 2: Conservative Downsampling

The script scripts/data_generation/downsample_reference.py downsamples the high-resolution (354 x 322) outputs to the coarse (90 x 82) training grid using conservative area-weighted averaging with factor 4.

Downsampling Operators

Five specialized operators handle different field types:

Operator	Fields	Method
downsample_scalar_consistent	T, S, el, tb, sb, elb, et, etb, rho, q2, q2l, km, kh, kq, w, aam	Area-weighted averaging (conservative for scalars)
downsample_vector_u_consistent	u, ub, ua, uab	Flux-conservative averaging over U-velocity faces, weighted by depth and dy
downsample_vector_v_consistent	v, vb, va, vab	Flux-conservative averaging over V-velocity faces, weighted by depth and dx
downsample_transport_u_consistent	utf, utb	Flux-conservative averaging for U-face advective transport
downsample_transport_v_consistent	vtf, vtb	Flux-conservative averaging for V-face advective transport

The scalar downsampling uses area weighting:

$${\overline{\varphi }}_{I,J}=\frac{\sum _{i,j\in {\mathrm{\Omega }}_{I,J}}{\varphi }_{i,j}\cdot A_{i,j}}{\sum _{i,j\in {\mathrm{\Omega }}_{I,J}}{A}_{i,j}}$$

where ${\mathrm{\Omega }}_{I,J}$ is the set of fine-grid cells mapping to coarse cell $(I,J)$, and $A_{i,j}$ are the cell areas from the fine grid.

Vector-Specific Downsampling

For U-velocity (defined on east-west cell faces), the downsampling accounts for the staggered grid geometry:

$${\overline{u}}_{I,J}=\frac{\sum _{j\in J\text{-block}}{u}_{i_{\text{face}},j}\cdot D_{i_{\text{face}},j}\cdot \mathrm{\Delta }{y}_{i_{\text{face}},j}}{D_{I,J}\cdot \mathrm{\Delta }{y}_{I,J}}$$

where $D$ is the total water depth ($h+\eta$) interpolated to the U-face and $\mathrm{\Delta }y$ is the meridional grid spacing at the U-face. An analogous formula applies to V-velocity.

Execution

NEURALPOM_DATA_BASE_DIR=./data/simulation_ensemble_20y \
python scripts/data_generation/downsample_reference.py

This processes all 8 members, producing checkpoints_lowres_consistent/ and reference_data_lowres_consistent/ subdirectories within each member directory.

Output structure after downsampling:

member_{id}_gpu{id}/
├── checkpoints_lowres_consistent/
│   └── checkpoint_Day{N}.npz    # 90×82 grid
└── reference_data_lowres_consistent/
    └── ref_Day{N}_seg{XX}.npz   # 90×82 grid, individual segments

Stage 3: POMCheckpointDataset

The POMCheckpointDataset class (in src/neuralpom/data/dataset.py) loads the downsampled data and pairs checkpoint states with reference targets for supervised training.

Train/Test Split

The 8 ensemble members are split into training and test sets:

Split	Members	Description
Train	0, 1, 2, 3, 4, 5, 6	Used for training the CorrectorNet
Test	7	Held-out member for evaluation

This split ensures the test data comes from an ensemble member never seen during training, providing a clean measure of generalization.

target_seg_idx Parameter

The target_seg_idx parameter controls how many high-frequency reference segments are loaded per sample:

• target_seg_idx=0: Load only seg00 (20-minute forecast). The model predicts one step ahead.
• target_seg_idx=3: Load seg00 through seg03 (80-minute rollout). The model predicts 4 steps ahead, enabling multi-step training.

The targets are stacked along the first dimension, producing a tensor of shape (T, jm, im, kb) for 3D fields, where T = target_seg_idx + 1.

Sample Pairing Logic

Each sample pairs a checkpoint at Day N with reference segments at Day N+1:

Input:  checkpoint_Day0100.npz    (state at end of day 100)
Target: ref_Day0101_seg00.npz     (20 min into day 101)
        ref_Day0101_seg01.npz     (40 min into day 101)
        ...
        ref_Day0101_seg{target_seg_idx}.npz

This pairing ensures the model is trained to forecast from a state that the dynamical core would naturally produce.

Sample Filtering

The dataset automatically filters samples by:

• Day range: Only days in [start_day, end_day) are included.
• Segment completeness: A sample is included only if ALL required reference segments (seg00 through seg{target_seg_idx}) exist.
• Temporal subsampling: The sample_freq parameter keeps every N-th valid sample for reducing dataset size.

DDP and DistributedSampler

For distributed training across multiple GPUs, the dataset is used with PyTorch's DistributedSampler:

from torch.utils.data import DataLoader, DistributedSampler

dataset = POMCheckpointDataset(
    mode='train',
    target_seg_idx=3,
    start_day=100,
    end_day=720
)

sampler = DistributedSampler(
    dataset,
    num_replicas=world_size,
    rank=local_rank,
    shuffle=True
)

loader = DataLoader(
    dataset,
    batch_size=batch_size,
    sampler=sampler,
    num_workers=4,
    pin_memory=True
)

Member-Based Data Organization

Each member directory is self-contained with its own checkpoints and reference data. This design allows:

• Independent generation of each member on different machines or at different times.
• Easy addition or removal of members without affecting others.
• Clean separation of training and test sets at the member level.

The dataset class iterates over the configured member directories and merges all valid samples into a flat list for the DataLoader.

Usage Examples

Generation

# Generate ensemble data
python scripts/data_generation/generate_double_gyre_ensemble.py

# Downsample to training resolution
python scripts/data_generation/downsample_reference.py

Loading Data

from neuralpom.data.dataset import POMCheckpointDataset

# Training dataset with 4-step rollout (80 minutes)
train_ds = POMCheckpointDataset(
    mode='train',
    target_seg_idx=3,
    start_day=100,
    end_day=720,
    sample_freq=1
)

# Test dataset with single-step
test_ds = POMCheckpointDataset(
    mode='test',
    target_seg_idx=0,
    start_day=100,
    end_day=1000
)

# Single-day inference slice
inference_ds = POMCheckpointDataset(
    mode='test',
    start_day=700,
    end_day=701,
    target_seg_idx=0
)

print(f"Train samples: {len(train_ds)}")
print(f"Test samples:  {len(test_ds)}")

inputs, targets = train_ds[0]
print(f"Input keys: {list(inputs.keys())}")
print(f"Target u shape: {targets['u'].shape}")   # (4, jm, im, kb)
print(f"Target eta shape: {targets['eta'].shape}")  # (4, jm, im)