NPZ Data Format Specification

NeuralPOM uses the NumPy .npz format (compressed NumPy archives) for all checkpoint and reference data files. This page specifies the exact array keys, shapes, dimensions, and conventions used throughout the data pipeline.

General Conventions

Array Dimensions

All 3D prognostic fields follow the spatial ordering:

(time, j, i, k) -- for multi-time reference files
(j, i, k)       -- for single-time checkpoint files

where:

j is the meridional index (y-direction, 0 to jm-1)
i is the zonal index (x-direction, 0 to im-1)
k is the vertical level index (0 to kb-1, where k=0 is the surface)

2D fields use the shape (j, i) for single-time or (time, j, i) for multi-time files.

Data Types

All arrays are stored as float32 (single precision) unless otherwise noted. The time key may be stored as float64 for precise temporal tracking.

Masking Convention

Velocity fields on the C-grid (u, ub, v, vb, etc.) use the corresponding velocity masks:

U-face fields (u, ub, ua, uab, utf, utb): Values are valid only where dum[j, i] = 1 at interior U-points.
V-face fields (v, vb, va, vab, vtf, vtb): Values are valid only where dvm[j, i] = 1 at interior V-points.
T-point fields (T, S, el, etc.): Values are valid only where fsm[j, i] = 1 at interior T-points.

Land-point values are set to zero and should be excluded from diagnostics.

Checkpoint File Format

File name pattern: checkpoint_Day{dddd}.npz

Daily checkpoint files contain the full model state at the end of a simulation day. All arrays have shape (jm, im) for 2D and (jm, im, kb) for 3D fields.

Prognostic State Variables

Key	Shape	Grid	Description
`u`	`(jm, im, kb)`	U-face	Eastward (zonal) velocity component
`v`	`(jm, im, kb)`	V-face	Northward (meridional) velocity component
`t`	`(jm, im, kb)`	T-point	Potential temperature
`s`	`(jm, im, kb)`	T-point	Salinity
`el`	`(jm, im)`	T-point	Free-surface elevation $η$

Time-History Variables (Leapfrog)

The leapfrog time-stepping scheme requires previous time-level variables (suffix b for "backward"):

Key	Shape	Grid	Description
`ub`	`(jm, im, kb)`	U-face	u at previous time level
`vb`	`(jm, im, kb)`	V-face	v at previous time level
`tb`	`(jm, im, kb)`	T-point	T at previous time level
`sb`	`(jm, im, kb)`	T-point	S at previous time level
`elb`	`(jm, im)`	T-point	el at previous time level
`etb`	`(jm, im)`	T-point	Time-filtered el at previous time level

Time-Filtered Variables

The Asselin-Robert time filter produces filtered fields (suffix f):

Key	Shape	Grid	Description
`et`	`(jm, im)`	T-point	Time-filtered free-surface elevation
`etf`	`(jm, im)`	T-point	Filtered et at previous time level
`elf`	`(jm, im)`	T-point	Filtered free-surface elevation (for implicit core)

Depth-Averaged Velocity

Key	Shape	Grid	Description
`ua`	`(jm, im)`	U-face	Depth-averaged zonal velocity $\bar{u}$
`va`	`(jm, im)`	V-face	Depth-averaged meridional velocity $\bar{v}$
`uab`	`(jm, im)`	U-face	ua at previous time level
`vab`	`(jm, im)`	V-face	va at previous time level
`uaf`	`(jm, im)`	U-face	Filtered ua (implicit core)
`vaf`	`(jm, im)`	V-face	Filtered va (implicit core)

Transport Terms

Key	Shape	Grid	Description
`utf`	`(jm, im, kb)`	U-face	Zonal advective flux (filtered)
`utb`	`(jm, im, kb)`	U-face	Zonal advective flux (previous)
`vtf`	`(jm, im, kb)`	V-face	Meridional advective flux (filtered)
`vtb`	`(jm, im, kb)`	V-face	Meridional advective flux (previous)

Turbulence Closure Variables (Mellor-Yamada 2.5)

Key	Shape	Grid	Description
`q2`	`(jm, im, kb)`	T-point	Twice the turbulent kinetic energy ( $2 \bar{q^{2}}$ )
`q2b`	`(jm, im, kb)`	T-point	q2 at previous time level
`q2l`	`(jm, im, kb)`	T-point	q2 $\times$ turbulent length scale ( $2 \bar{q^{2}} l$ )
`q2lb`	`(jm, im, kb)`	T-point	q2l at previous time level
`km`	`(jm, im, kb)`	T-point	Vertical eddy viscosity $K_{M}$
`kh`	`(jm, im, kb)`	T-point	Vertical eddy diffusivity $K_{H}$
`kq`	`(jm, im, kb)`	T-point	Vertical diffusivity for TKE
`w`	`(jm, im, kb)`	T-point (W-points)	Vertical velocity component $ω$

Optional Forcing and Diagnostic Fields

These fields may or may not be present depending on the simulation configuration:

Key	Shape	Grid	Description
`adx2d`	`(jm, im)`	U-face	Depth-averaged zonal advection term
`ady2d`	`(jm, im)`	V-face	Depth-averaged meridional advection term
`drx2d`	`(jm, im)`	U-face	Depth-averaged zonal diffusion term
`dry2d`	`(jm, im)`	V-face	Depth-averaged meridional diffusion term
`aam`	`(jm, im)`	T-point	Spatially varying horizontal viscosity $A_{M}$
`vfluxf`	`(jm, im)`	T-point	Surface volume flux (filtered)
`vfluxb`	`(jm, im)`	T-point	Bottom volume flux
`wubot`	`(jm, im)`	U-face	Bottom stress zonal component
`wvbot`	`(jm, im)`	V-face	Bottom stress meridional component
`egb`	`(jm, im)`	T-point	Total kinetic energy at bottom
`rho`	`(jm, im, kb)`	T-point	In-situ density $ρ$

Time Variable

Key	Shape	Grid	Description
`time`	scalar or `(1,)`	--	Model elapsed time in seconds

Reference File Format (Legacy Layout)

File name pattern: ref_Day{dddd}_seg{ss}.npz

When NEURALPOM_REF_LAYOUT=legacy, each 20-minute reference segment is saved as a separate file. Each file contains a subset of the keys at a single time snapshot:

All arrays have shape (jm, im) for 2D and (jm, im, kb) for 3D.
Typically includes: time, u, ub, v, vb, t, s, el, ua, va, uab, vab.
No time dimension is stacked -- each file contains one instant.

Reference File Format (Packed Layout)

File name pattern: ref_Day{dddd}.npz

When NEURALPOM_REF_LAYOUT=packed, an entire day's reference data is saved in a single file. The arrays have an additional time dimension stacked along axis 0:

Key	Shape	Description
`time`	`(12,)`	Timestamps for each segment (in seconds)
`u`	`(12, jm, im, kb)`	Velocity at each 20-min segment
`v`	`(12, jm, im, kb)`	Velocity at each 20-min segment
...	...	All other keys follow the same pattern

The packed layout is the default and recommended format as it reduces file system overhead for large ensemble runs.

Reference File Format (Downsampled)

File name pattern: ref_Day{dddd}_seg{ss}.npz

These live in the reference_data_lowres_consistent/ subdirectory within each member. The downsampling script unpacks the packed reference files (if present) and outputs individual segment files at the coarse 90 x 82 resolution.

Coarse-grid array shapes:

Key	Shape	Grid
`u`, `v`, `t`, `s`	`(jm=82, im=90, kb=21)`	As above
`el`, `ua`, `uab`, `va`, `vab`	`(jm=82, im=90)`	As above
`time`	scalar	--

Initial Condition Format

Initial conditions are typically set programmatically via the core's case setup routines (double_gyre_setup(), seamount_setup(), box_setup()). However, for custom datasets (problem_id = 3, file_init), the initial conditions are read from external NPZ files.

Mean Field File Format

For the StateNormalizer, mean and standard deviation fields are stored in NPZ files for each prognostic variable:

File name: mean_{variable}.npz and std_{variable}.npz

Each file contains:

A single array with the same spatial shape as the variable itself: (jm, im, kb) for 3D fields, (jm, im) for 2D fields.
Computed as the time-mean over a specified training period.

Usage:

\tilde{ϕ} = \frac{ϕ - μ_{ϕ}}{σ_{ϕ} + ϵ}

where $μ_{ϕ}$ is from mean_phi.npz, $σ_{ϕ}$ is from std_phi.npz, and $ϵ$ is a small constant for numerical stability.

Dataset Loading

The POMCheckpointDataset class reads all arrays from checkpoint and reference files, converts them to PyTorch tensors, and returns (inputs, targets) tuple. The loading code in __getitem__ handles all the required keys transparently:

python

def to_tensor(arr):
    return torch.from_numpy(arr).float()

# Inputs from checkpoint
inputs['u']  = to_tensor(data_in['u'])
inputs['v']  = to_tensor(data_in['v'])
inputs['t']  = to_tensor(data_in['t'])
inputs['s']  = to_tensor(data_in['s'])
inputs['el'] = to_tensor(data_in['el'])
# ... and all other keys

# Targets from reference segments, stacked along T dimension
u_list = [ref_tensor(d, 'u') for d in ref_data_list]
targets['u'] = torch.stack(u_list, dim=0)  # (T, jm, im, kb)

Array Shape Summary

Category	3D Fields	2D Fields
Checkpoint (coarse)	`(82, 90, 21)`	`(82, 90)`
Reference legacy	`(82, 90, 21)`	`(82, 90)`
Reference packed	`(12, 82, 90, 21)`	`(12, 82, 90)`
Mean/std fields	`(82, 90, 21)`	`(82, 90)`
Dataset input (PyTorch)	`(jm, im, kb)`	`(jm, im)`
Dataset target (PyTorch)	`(T, jm, im, kb)`	`(T, jm, im)`

Where T = target_seg_idx + 1 and (jm, im) = (82, 90) for the standard training grid.

NPZ Data Format Specification ​

General Conventions ​

Array Dimensions ​

Data Types ​

Masking Convention ​

Checkpoint File Format ​

Prognostic State Variables ​

Time-History Variables (Leapfrog) ​

Time-Filtered Variables ​

Depth-Averaged Velocity ​

Transport Terms ​

Turbulence Closure Variables (Mellor-Yamada 2.5) ​

Optional Forcing and Diagnostic Fields ​

Time Variable ​

Reference File Format (Legacy Layout) ​

Reference File Format (Packed Layout) ​

Reference File Format (Downsampled) ​

Initial Condition Format ​

Mean Field File Format ​

Dataset Loading ​

Array Shape Summary ​