Case System Overview

NeuralPOM supports multiple benchmark cases for ocean simulation. Legacy POM cases can be run through the explicit or implicit hydrostatic cores, while newer benchmarks may target a specific core family such as implicit_jzb, nonhydrostatic, or geostrophic_2d. The case system is designed around a centralized registry and a shared settings dataclass, keeping configuration separate from core implementation details.

Architecture

CaseSettings Dataclass

All case configuration flows through the CaseSettings dataclass defined in src/neuralpom/cases/common.py:

from dataclasses import dataclass, field
from typing import Any

@dataclass(frozen=True)
class CaseSettings:
    name: str
    problem_id: int
    im: int
    jm: int
    description: str
    overrides: dict[str, Any] = field(default_factory=dict)

Field	Type	Description
name	str	Human-readable case identifier
problem_id	int	Numeric ID used internally by the core to select case setup logic
im	int	Number of grid points in the i (zonal / x) direction
jm	int	Number of grid points in the j (meridional / y) direction
description	str	Short description of the case
overrides	dict	Optional key-value overrides for core-level parameters

catalog.py Case Registry

The central registry is in src/neuralpom/cases/catalog.py and maps case names to CaseSettings objects:

from neuralpom.cases.catalog import CASE_SETTINGS_REGISTRY, get_case_settings

# Access the full registry
for name, settings in CASE_SETTINGS_REGISTRY.items():
    print(f"{name}: {settings.description}")

# Look up a specific case
dg_settings = get_case_settings("double_gyre")
print(f"Grid: {dg_settings.im} x {dg_settings.jm}")

The get_case_settings(name) function raises KeyError if the case name is not found, and returns the CaseSettings object otherwise.

Explicit vs. Implicit Variants

Each case has two implementations, located in separate subdirectories:

src/neuralpom/cases/
├── catalog.py          # Central registry
├── common.py           # CaseSettings dataclass
├── explicit/
│   ├── double_gyre.py  # Explicit free-surface variant
│   ├── seamount.py
│   └── box.py
└── implicit/
    ├── double_gyre.py  # Implicit free-surface variant
    ├── seamount.py
    └── box.py

Each variant file exposes a module-level CASE_SETTINGS object that the corresponding core reads during initialization. The same problem_id is used across both variants so that the core's internal cases.py dispatches to the correct setup routine.

Available Cases

Case	problem_id	Core Family	Typical Grid	Description
Seamount	1	explicit, implicit	65 x 49	Flow over an isolated Gaussian seamount. Classic POM benchmark.
Box	2	explicit, implicit	129 x 129	Minimal closed basin for debugging and conservation checks.
File Init	3	explicit, implicit	65 x 49	External file-driven initialization for custom datasets.
Double Gyre	4	explicit, implicit	90 x 82	Wind-driven circulation in a closed basin. Primary training benchmark.
Bire 2025	5	implicit	250 x 250	High-resolution double gyre following Bire et al. (2025).
Baroclinic Adjustment	6	implicit	194 x 194 x 33	Oceananigans-style baroclinic front and channel adjustment benchmark.
JZB Tide	8	implicit_jzb	312 x 213	Jiaozhou Bay realistic-topography M2 tide reproduction.
2D Geostrophic Turbulence	20	geostrophic_2d	100 x 100	Double-periodic 2D geostrophic turbulence benchmark.
Taylor-Green NH	101	nonhydrostatic	64 x 64 x 32	Non-hydrostatic Taylor-Green vortex regression case.
Rayleigh-Benard NH	102	nonhydrostatic	64 x 64 x 32	Non-hydrostatic convection benchmark.
Kelvin-Helmholtz NH	103	nonhydrostatic	128 x 64 x 64	Non-hydrostatic shear-instability benchmark.
Lock Exchange NH	104	nonhydrostatic	128 x 64 x 64	Non-hydrostatic density-current benchmark.

Note: The typical NeuralPOM grid sizes for training are 90 x 82 (coarse) and 354 x 322 (fine). The catalog entries record nominal grid dimensions; the per-variant files may override these for specific workflows.

How Cases Are Configured

When a core is instantiated with a grid size and a problem_id, the initialization sequence is:

1. Grid allocation -- Arrays are allocated on the specified (im, jm) grid with kb vertical levels (typically 21).
2. Depth and topography -- Bathymetry is set from the case setup (e.g., flat bottom for double gyre, Gaussian seamount for the seamount case).
3. Initial conditions -- Velocity (u, v), temperature (T), salinity (S), and free-surface elevation (el) are initialized according to the case.
4. Masks and metrics -- Land-sea masks (fsm, dum, dvm) and geometric factors (dx, dy, art, h) are computed.
5. Boundary conditions -- Lateral boundary conditions are applied based on the case configuration.

Usage Pattern

The typical usage pattern for running a case simulation:

# Option 1: Import and run from the case module directly
from neuralpom.cases.explicit.double_gyre import CASE_SETTINGS
from neuralpom.cores.explicit import ExplicitPOM2KCore

core = ExplicitPOM2KCore(
    im=CASE_SETTINGS.im,
    jm=CASE_SETTINGS.jm
)
core.iproblem = CASE_SETTINGS.problem_id
core.run_initialization()

# Run time loop
core.time0 = core.time
core.iend = 1000
core.run_time_loop()

# Option 2: Use the catalog for dynamic case selection
from neuralpom.cases.catalog import get_case_settings
from neuralpom.core_factory import resolve_core_class

settings = get_case_settings("double_gyre")
core_cls = resolve_core_class("implicit")
core = core_cls(im=settings.im, jm=settings.jm)
core.iproblem = settings.problem_id
core.run_initialization()

The iproblem attribute is the critical link -- it tells the core's internal cases.py module which setup routine to invoke inside run_initialization(). Each problem_id corresponds to a dedicated setup method that configures domain geometry, forcing, bathymetry, and initial conditions.

Core-Level Case Setup

The actual case implementation lives inside each core's cases.py module (e.g., src/neuralpom/cores/implicit/cases.py). When run_initialization() is called, the core dispatches to the appropriate setup method based on self.iproblem:

problem_id 1 -> seamount_setup()
problem_id 2 -> box_setup()
problem_id 3 -> file2ic_setup()
problem_id 4 -> double_gyre_setup()
problem_id 5 -> bire_2025_setup()
problem_id 6 -> baroclinic_adjustment_setup()

This separation means that case configuration (the CaseSettings dataclass and registry) is decoupled from case implementation (the cases.py setup routines or equivalent core-family adapters). New cases can be added by:

1. Adding a CaseSettings entry to catalog.py
2. Choosing the owning core family through overrides["core_family"] when the case is not a legacy hydrostatic POM case
3. Implementing the setup routine or adapter in the owning core family