ocean_data.py

# SPDX-License-Identifier: Apache-2.0
"""Streaming access to the public 1° OM4 dataset used by Samudra 2.

The processed OM4 zarr stores live on the NYU OSN pod and are public-read:

    https://nyu1.osn.mghpcc.org/m2lines-pubs/Samudra/v2026-07/om4_onedeg/

Only the handful of (time, y, x) chunks a rollout actually needs are pulled,
so a request moves tens of MB rather than the 92 GiB of the full store.

Channel layout, normalization and masking follow
`samudra.datasets.InferenceDataset` upstream:

* prognostic input  = (hist+1=2 timesteps) x 77 variables = 154 channels
* boundary  input   = 2 timesteps x 4 variables            =   8 channels
* model output      = 154 channels = the next 2 timesteps of the 77 variables

so one model step advances the ocean state by 2 x 5 days = 10 days.
"""

from __future__ import annotations

import os
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache

import numpy as np

OSN_ENDPOINT = "https://nyu1.osn.mghpcc.org"
BUCKET_ROOT = "m2lines-pubs/Samudra/v2026-07/om4_onedeg"

LEVELS = 19
DEPTHS = (
    2.5, 10.0, 22.5, 40.0, 65.0, 105.0, 165.0, 250.0, 375.0, 550.0,
    775.0, 1050.0, 1400.0, 1850.0, 2400.0, 3100.0, 4000.0, 5000.0, 6000.0,
)
# `thermo_dynamic_all` prognostic variables, in upstream order.
PROG_VARS: list[str] = (
    [f"uo_{i}" for i in range(LEVELS)]
    + [f"vo_{i}" for i in range(LEVELS)]
    + [f"thetao_{i}" for i in range(LEVELS)]
    + [f"so_{i}" for i in range(LEVELS)]
    + ["zos"]
)
# `tau_hfds_hfds_anom` boundary (forcing) variables.
BOUNDARY_VARS = ["tauuo", "tauvo", "hfds", "hfds_anomalies"]

N_PROG = len(PROG_VARS)          # 77
HIST = 1                          # samudra_om4_v2 default
STEP_DAYS = 10                    # (hist + 1) * 5-day timesteps

_HERE = os.path.dirname(os.path.abspath(__file__))
HFDS_ANOM_STATS = os.path.join(_HERE, "hfds_anom_stats.npz")


def _level_of(var: str) -> int:
    tail = var.rsplit("_", 1)[-1]
    return int(tail) if tail.isdigit() else 0


class OM4Store:
    """Lazily-opened handle on the public 1° OM4 zarr store."""

    def __init__(self, root: str = BUCKET_ROOT, endpoint: str = OSN_ENDPOINT):
        import s3fs
        import xarray as xr

        fs = s3fs.S3FileSystem(anon=True, endpoint_url=endpoint)
        self.ds = xr.open_zarr(s3fs.S3Map(root=f"{root}/OM4.zarr", s3=fs, check=False))
        means = xr.open_zarr(
            s3fs.S3Map(root=f"{root}/OM4_means.zarr", s3=fs, check=False)
        ).load()
        stds = xr.open_zarr(
            s3fs.S3Map(root=f"{root}/OM4_stds.zarr", s3=fs, check=False)
        ).load()

        # hfds_anomalies is a derived channel (hfds minus its day-of-year
        # climatology); the climatology + its normalization stats are shipped
        # with the Space, precomputed from this very store.
        stats = np.load(HFDS_ANOM_STATS)
        self._hfds_clim = stats["clim"].astype(np.float32)             # (73, y, x)
        self._clim_doy = {int(d): i for i, d in enumerate(stats["dayofyear"])}
        anom_mean, anom_std = float(stats["mean"]), float(stats["std"])

        self.means = {v: float(means[v].values) for v in means.data_vars}
        self.stds = {v: float(stds[v].values) for v in stds.data_vars}
        self.means["hfds_anomalies"] = anom_mean
        self.stds["hfds_anomalies"] = anom_std

        self.time = self.ds.time.values
        self.dayofyear = self.ds.time.dt.dayofyear.values
        # `lat`/`lon` in the store are 2-D (y, x) curvilinear coords; the plot
        # axes are the 1-D `y` / `x` cell centers.
        self.lat = np.asarray(self.ds.y.values, np.float64)
        self.lon = np.asarray(self.ds.x.values, np.float64)
        self.masks = np.stack(
            [self.ds[f"mask_{i}"].values.astype(bool) for i in range(LEVELS)]
        )  # (19, y, x)
        self.shape = self.masks.shape[1:]

        self.prog_mask = np.stack([self.masks[_level_of(v)] for v in PROG_VARS])
        self.prog_means = np.array([self.means[v] for v in PROG_VARS], dtype=np.float32)
        self.prog_stds = np.array([self.stds[v] for v in PROG_VARS], dtype=np.float32)

    # ---------------------------------------------------------------- helpers
    def date_str(self, index: int) -> str:
        return str(self.time[index])[:10]

    def _read(self, var: str, t0: int, n: int) -> np.ndarray:
        """(n, y, x) raw values for `var` over times [t0, t0+n)."""
        return np.asarray(self.ds[var].isel(time=slice(t0, t0 + n)).values, np.float32)

    def _read_many(self, variables: list[str], t0: int, n: int) -> np.ndarray:
        """(n, len(variables), y, x), fetched in parallel."""
        with ThreadPoolExecutor(max_workers=16) as pool:
            arrays = list(pool.map(lambda v: self._read(v, t0, n), variables))
        return np.stack(arrays, axis=1)

    def _hfds_anomalies(self, hfds: np.ndarray, t0: int, n: int) -> np.ndarray:
        idx = [self._clim_doy[int(d)] for d in self.dayofyear[t0 : t0 + n]]
        return hfds - self._hfds_clim[idx]

    # ------------------------------------------------------------- public API
    def initial_prognostic(self, t0: int) -> np.ndarray:
        """Normalized, masked (1, 154, y, x) initial state at times [t0, t0+1]."""
        raw = self._read_many(PROG_VARS, t0, HIST + 1)          # (2, 77, y, x)
        norm = (raw - self.prog_means[None, :, None, None]) / self.prog_stds[
            None, :, None, None
        ]
        norm = np.nan_to_num(norm, nan=0.0)
        norm = np.where(self.prog_mask[None], norm, 0.0)
        return norm.reshape(1, (HIST + 1) * N_PROG, *self.shape).astype(np.float32)

    def boundary_sequence(self, t0: int, n_steps: int) -> np.ndarray:
        """Normalized, masked (n_steps, 8, y, x) forcing for `n_steps` model steps."""
        n_times = (HIST + 1) * n_steps
        raw = self._read_many(["tauuo", "tauvo", "hfds"], t0, n_times)   # (T, 3, y, x)
        anom = self._hfds_anomalies(raw[:, 2], t0, n_times)[:, None]
        raw = np.concatenate([raw, anom], axis=1)                        # (T, 4, y, x)

        means = np.array([self.means[v] for v in BOUNDARY_VARS], np.float32)
        stds = np.array([self.stds[v] for v in BOUNDARY_VARS], np.float32)
        norm = (raw - means[None, :, None, None]) / stds[None, :, None, None]
        norm = np.nan_to_num(norm, nan=0.0)
        norm = np.where(self.masks[0][None, None], norm, 0.0)
        return norm.reshape(n_steps, (HIST + 1) * len(BOUNDARY_VARS), *self.shape).astype(
            np.float32
        )

    def truth(self, var: str, t0: int, n_steps: int) -> np.ndarray:
        """Raw (physical-unit) ground truth for `var` over the predicted times."""
        n_times = (HIST + 1) * n_steps
        raw = self._read(var, t0 + HIST + 1, n_times)
        return np.where(self.masks[_level_of(var)][None], raw, np.nan)

    def denormalize(self, channels: np.ndarray, var: str) -> np.ndarray:
        """Turn normalized model output for one variable into physical units."""
        v = PROG_VARS.index(var)
        out = channels * self.prog_stds[v] + self.prog_means[v]
        return np.where(self.masks[_level_of(var)][None], out, np.nan)


@lru_cache(maxsize=1)
def get_store() -> OM4Store:
    return OM4Store()