diff --git a/README.md b/README.md index ef81356ccdeccab85def4e9a358cf6fed9e2d210..f76e05b209efcb02a72184bf971305cb70aafcef 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,40 @@ --- -title: Samudra2 -emoji: 🦀 -colorFrom: pink -colorTo: indigo +title: Samudra 2 +emoji: 🌊 +colorFrom: green +colorTo: gray sdk: gradio sdk_version: 6.25.0 -python_version: '3.12' app_file: app.py +short_description: Neural ocean emulator rollouts scored against GFDL OM4 +python_version: "3.12" +startup_duration_timeout: 1h pinned: false +license: apache-2.0 +models: + - M2LInES/Samudra2 --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# Samudra 2 — a neural ocean emulator + +[Samudra 2](https://huggingface.co/M2LInES/Samudra2) (M2LInES) is an autoregressive +deep-learning emulator of the **GFDL OM4 global ocean model**. Given one ocean state plus +atmospheric forcing it advances temperature, salinity and horizontal currents across 19 +depth levels together with sea surface height — 77 prognostic fields at once — 10 days per +autoregressive step. + +This Space runs the released **1° checkpoint** (`onedeg/ema_ckpt.pt`, 84M parameters) on +genuinely held-out years (2014–2022). The initial condition and the atmospheric forcing are +streamed live from the public OM4 archive (`m2lines-pubs/Samudra` on the Open Storage +Network), the emulator is rolled forward on ZeroGPU, and the emulated ocean is scored +against the OM4 simulation it emulates — including a persistence baseline. + +## Notes + +- The vendored `samudra/` package is the upstream + [`m2lines/Samudra`](https://github.com/m2lines/Samudra) library (Apache-2.0), included + directly because its full dependency set (SkyPilot, xESMF) cannot be pip-resolved inside + a Space. Model configuration matches the upstream `samudra_om4_v2` config. +- Only the 1° checkpoint is demoed: the half- and quarter-degree checkpoints require the + `hfds_anomalies` boundary field, which is not present in the corresponding public stores. +- Each emulator step emits two 5-day snapshots, so *n* steps ⇒ 10·*n* days of forecast. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..dd4534a10f2016ee21415994b9fa04f6d2b02066 --- /dev/null +++ b/app.py @@ -0,0 +1,400 @@ +"""Samudra 2 — an autoregressive neural ocean emulator (M2LInES). + +The app streams real GFDL-OM4 ocean state from public S3 storage, rolls the +released 1-degree Samudra 2 checkpoint forward autoregressively on ZeroGPU, +and compares the emulated ocean against the OM4 simulation it emulates. +""" + +import spaces # noqa: F401 (must precede any CUDA-touching import) + +import dataclasses +import functools +import os +import tempfile +import time + +import matplotlib + +matplotlib.use("Agg") + +import gradio as gr +import imageio.v2 as imageio +import matplotlib.pyplot as plt +import numpy as np +import s3fs +import torch +import xarray as xr +from huggingface_hub import hf_hub_download + +import samudra.utils.device as sdev +from samudra.config import BlockConfig, SamudraConfig, UNetBackboneConfig +from samudra.constants import build_om4_spec +from samudra.datasets import InferenceDataset +from samudra.utils.data import DataSource, Normalize + +# Data pipeline work happens on CPU in the main process; the GPU worker flips +# this to "cuda" for the forward passes only. +sdev.set_device(torch.device("cpu")) + +# -------------------------------------------------------------------------- +# Data: public 1-degree OM4 store (Open Storage Network, anonymous access) +# -------------------------------------------------------------------------- +ENDPOINT = "https://nyu1.osn.mghpcc.org" +ROOT = "m2lines-pubs/Samudra" +HIST = 1 +MAX_STEPS = 12 + +_fs = s3fs.S3FileSystem(anon=True, client_kwargs={"endpoint_url": ENDPOINT}) + + +def _open(name: str) -> xr.Dataset: + return xr.open_zarr( + s3fs.S3Map(root=f"{ROOT}/{name}", s3=_fs, check=False), + consolidated=True, + chunks={}, + ) + + +print("[samudra] opening public OM4 zarr stores ...", flush=True) +_t0 = time.time() +SPEC = build_om4_spec("thermo_dynamic_all", "tau_hfds_hfds_anom") +SRC = DataSource.from_datasets( + _open("OM4"), + _open("OM4_means"), + _open("OM4_stds"), + dataset_spec=SPEC, + prognostic_var_names=SPEC.prognostic_var_names, + boundary_var_names=SPEC.boundary_var_names, + name="om4-1deg", +) +PROG_NAMES = list(SPEC.prognostic_var_names) +BND_NAMES = list(SPEC.boundary_var_names) +N_PROG = len(PROG_NAMES) +_norm = Normalize(SRC, SPEC.prognostic_var_names, SPEC.boundary_var_names) +PROG_MEAN = np.asarray(_norm._prognostic_mean_np, dtype=np.float64) +PROG_STD = np.asarray(_norm._prognostic_std_np, dtype=np.float64) +TIMES = [str(t)[:10] for t in SRC.data.time.values] +print(f"[samudra] data source ready in {time.time() - _t0:.1f}s " + f"({len(TIMES)} snapshots, {N_PROG} prognostic vars)", flush=True) + +# Samudra 2 was trained on 1975-2013 and validated to 2014-10; everything from +# 2014-10-10 onward is held out, so restrict the demo to genuine forecasts. +FIRST_HELDOUT = TIMES.index("2014-10-10") +LAST_START = len(TIMES) - (2 * MAX_STEPS + 2) +START_DATES = TIMES[FIRST_HELDOUT : LAST_START + 1] + +# -------------------------------------------------------------------------- +# Model: the released 1-degree Samudra 2 checkpoint +# -------------------------------------------------------------------------- +print("[samudra] building model ...", flush=True) +_model_cfg = SamudraConfig( + checkpointing=None, + pred_residuals=False, + last_kernel_size=3, + pad="circular", + unet=UNetBackboneConfig( + ch_width=[280, 380, 480, 520], + dilation=[1, 2, 4, 8], + n_layers=[1, 1, 1, 1], + core_block=BlockConfig( + block_type="conv_next_block", + kernel_size=3, + activation="capped_gelu", + upscale_factor=2, + norm="batch", + ), + down_sampling_block="avg_pool", + up_sampling_block="zonally_periodic_upsample", + ), +) +_grid_src = dataclasses.replace(SRC, data=SRC.data.isel(time=slice(0, 4))) +MODEL = _model_cfg.build( + prog_channels=(HIST + 1) * N_PROG, + boundary_channels=(HIST + 1) * len(BND_NAMES), + out_channels=(HIST + 1) * N_PROG, + hist=HIST, + srcs=[_grid_src], +) +_ckpt_path = hf_hub_download("M2LInES/Samudra2", "onedeg/ema_ckpt.pt") +_ck = torch.load(_ckpt_path, map_location="cpu", weights_only=False) +MODEL.load_state_dict( + {k.removeprefix("module."): v for k, v in _ck["model"].items()}, strict=True +) +del _ck +DEVICE = os.environ.get("SAMUDRA_DEVICE", "cuda") +MODEL.eval().to(DEVICE) +N_PARAMS = sum(p.numel() for p in MODEL.parameters()) +print(f"[samudra] model ready ({N_PARAMS / 1e6:.1f}M params)", flush=True) + +# -------------------------------------------------------------------------- +# Variable menu +# -------------------------------------------------------------------------- +LEVELS = [2.5, 10, 22.5, 40, 65, 105, 165, 250, 375, 550, 775, 1050, 1400, + 1850, 2400, 3100, 4000, 5000, 6000] + +VARIABLES: dict[str, dict] = { + "Sea surface temperature (2.5 m)": dict(key="thetao_0", unit="°C", cmap="RdYlBu_r"), + "Ocean temperature at 105 m": dict(key="thetao_5", unit="°C", cmap="RdYlBu_r"), + "Ocean temperature at 550 m": dict(key="thetao_9", unit="°C", cmap="RdYlBu_r"), + "Sea surface salinity (2.5 m)": dict(key="so_0", unit="psu", cmap="viridis"), + "Sea surface height": dict(key="zos", unit="m", cmap="Spectral_r"), + "Zonal velocity (2.5 m)": dict(key="uo_0", unit="m/s", cmap="RdBu_r"), + "Meridional velocity (2.5 m)": dict(key="vo_0", unit="m/s", cmap="RdBu_r"), +} +for _meta in VARIABLES.values(): + _meta["idx"] = PROG_NAMES.index(_meta["key"]) + + +# -------------------------------------------------------------------------- +# Rollout +# -------------------------------------------------------------------------- +@functools.lru_cache(maxsize=4) +def _prepare(i0: int, n_steps: int): + """Stream the initial ocean state and atmospheric forcing for a window.""" + window = SRC.data.isel(time=slice(i0, i0 + 2 * n_steps + 2)) + windowed_src = dataclasses.replace(SRC, data=window) + dataset = InferenceDataset( + src=windowed_src, + prognostic_var_names=SPEC.prognostic_var_names, + boundary_var_names=SPEC.boundary_var_names, + hist=HIST, + normalize_before_mask=True, + masked_fill_value=0.0, + long_rollout=False, + ) + initial = dataset.initial_prognostic + boundaries = tuple(dataset.get_boundary(k) for k in range(n_steps)) + target_times = dataset.get_target_time(0, n_steps) + return initial, boundaries, dataset.ctx, target_times, window + + +def _gpu_duration(initial, boundaries, ctx, var_idx) -> int: + return int(12 + 2 * len(boundaries)) + + +@spaces.GPU(duration=_gpu_duration) +def _rollout(initial, boundaries, ctx, var_idx: int): + """Autoregressively step the emulator; returns one normalised field on CPU.""" + sdev.set_device(torch.device(DEVICE)) + ctx = ctx.to(torch.device(DEVICE)) + state = initial.to(DEVICE) + outputs = [] + with torch.no_grad(): + for boundary in boundaries: + state = MODEL.forward_once(state, boundary.to(DEVICE), ctx) + outputs.append(state) + stacked = torch.cat(outputs, dim=0) + stacked = stacked.reshape(-1, N_PROG, stacked.shape[-2], stacked.shape[-1]) + return stacked[:, var_idx].float().cpu().numpy() + + +# -------------------------------------------------------------------------- +# Rendering +# -------------------------------------------------------------------------- +def _roll(field: np.ndarray) -> np.ndarray: + """Shift longitudes from 0..360 to -180..180 for a familiar world map.""" + return np.roll(field, field.shape[-1] // 2, axis=-1) + + +def _make_video(pred, truth, dates, label, unit, cmap, fps): + diff = pred - truth + finite = np.isfinite(truth) + vmin, vmax = np.nanpercentile(truth[finite], [1, 99]) + dmax = float(np.nanpercentile(np.abs(diff[np.isfinite(diff)]), 99.5)) + dmax = max(dmax, 1e-6) + extent = [-180, 180, -90, 90] + + fig, axes = plt.subplots(3, 1, figsize=(7.2, 9.6), dpi=100) + artists = [] + titles = [f"Samudra 2 emulator — {label}", f"GFDL OM4 simulation — {label}", + "Difference (emulator − OM4)"] + for ax, title, cm, lo, hi in zip( + axes, titles, [cmap, cmap, "RdBu_r"], [vmin, vmin, -dmax], [vmax, vmax, dmax] + ): + colormap = matplotlib.colormaps[cm].copy() + colormap.set_bad("0.85") + im = ax.imshow( + np.zeros_like(pred[0]), origin="lower", extent=extent, + cmap=colormap, vmin=lo, vmax=hi, interpolation="nearest", + ) + ax.set_title(title, fontsize=10) + ax.set_xticks([-180, -90, 0, 90, 180]) + ax.set_yticks([-90, -45, 0, 45, 90]) + ax.tick_params(labelsize=7) + fig.colorbar(im, ax=ax, fraction=0.025, pad=0.02).ax.tick_params(labelsize=7) + artists.append(im) + suptitle = fig.suptitle("", fontsize=12) + fig.tight_layout(rect=[0, 0, 1, 0.96]) + + frames = [] + for k, date in enumerate(dates): + artists[0].set_data(_roll(pred[k])) + artists[1].set_data(_roll(truth[k])) + artists[2].set_data(_roll(diff[k])) + suptitle.set_text(f"{date} · lead time +{5 * (k + 1)} days · [{unit}]") + fig.canvas.draw() + frames.append(np.asarray(fig.canvas.buffer_rgba())[..., :3].copy()) + plt.close(fig) + + frames = frames + [frames[-1]] * max(1, int(fps)) + path = os.path.join(tempfile.mkdtemp(), "samudra2_rollout.mp4") + imageio.mimwrite( + path, frames, fps=fps, codec="libx264", quality=8, macro_block_size=16 + ) + return path + + +def _make_skill_plot(rmse, persistence, unit, label): + leads = 5 * (np.arange(len(rmse)) + 1) + fig, ax = plt.subplots(figsize=(7.2, 3.6), dpi=100) + ax.plot(leads, persistence, "o--", color="0.55", label="Persistence baseline") + ax.plot(leads, rmse, "o-", color="#1f77b4", label="Samudra 2 emulator") + ax.set_xlabel("Lead time (days)") + ax.set_ylabel(f"RMSE vs OM4 [{unit}]") + ax.set_title(f"Emulator error growth — {label}", fontsize=11) + ax.grid(alpha=0.3) + ax.legend(fontsize=9) + ax.set_ylim(bottom=0) + fig.tight_layout() + path = os.path.join(tempfile.mkdtemp(), "samudra2_skill.png") + fig.savefig(path) + plt.close(fig) + return path + + +# -------------------------------------------------------------------------- +# Handler +# -------------------------------------------------------------------------- +def forecast( + variable: str, + start_date: str, + n_steps: int = 6, + fps: int = 3, + progress=gr.Progress(track_tqdm=True), +): + """Roll Samudra 2 forward from a held-out ocean state and score it against OM4. + + Args: + variable: Which ocean field to visualise (the emulator always advances + the full 77-channel ocean state). + start_date: Initial condition date, from the held-out 2014-2022 period. + n_steps: Number of autoregressive emulator steps; each step advances the + ocean by 10 days (two 5-day snapshots). + fps: Frames per second of the returned animation. + + Returns: + An MP4 animation (emulator / OM4 / difference), an error-growth chart, + and a short text summary. + """ + meta = VARIABLES[variable] + n_steps = int(n_steps) + i0 = TIMES.index(start_date) + + progress(0.05, desc="Streaming ocean state from public OM4 archive…") + initial, boundaries, ctx, target_times, window = _prepare(i0, n_steps) + + key = meta["key"] + idx = meta["idx"] + + progress(0.45, desc="Running the emulator on GPU…") + t_gpu = time.time() + pred_norm = _rollout(initial, boundaries, ctx, idx) + gpu_seconds = time.time() - t_gpu + + progress(0.7, desc="Fetching the OM4 reference…") + wet = SRC.masks.prognostic[idx].numpy().astype(bool) + pred = pred_norm.astype(np.float64) * PROG_STD[idx] + PROG_MEAN[idx] + pred = np.where(wet[None], pred, np.nan) + truth = np.asarray(window[key].sel(time=target_times).values, dtype=np.float32) + truth = np.where(wet[None], truth, np.nan) + initial_field = np.where(wet, np.asarray(window[key].isel(time=1).values), np.nan) + + rmse = np.sqrt(np.nanmean((pred - truth) ** 2, axis=(1, 2))) + persistence = np.sqrt(np.nanmean((initial_field[None] - truth) ** 2, axis=(1, 2))) + + dates = [str(t)[:10] for t in np.asarray(target_times)] + progress(0.8, desc="Rendering animation…") + video = _make_video(pred, truth, dates, variable, meta["unit"], meta["cmap"], int(fps)) + plot = _make_skill_plot(rmse, persistence, meta["unit"], variable) + + summary = ( + f"**Initial condition** {TIMES[i0 + 1]} → **forecast** {dates[0]} … {dates[-1]} \n" + f"**{n_steps} autoregressive steps** = {10 * n_steps} days " + f"({len(dates)} 5-day snapshots) in **{gpu_seconds:.1f} s** of GPU time. \n" + f"**{variable}** RMSE vs OM4: {rmse[0]:.3f} {meta['unit']} at +5 days → " + f"{rmse[-1]:.3f} {meta['unit']} at +{5 * len(dates)} days " + f"(persistence: {persistence[-1]:.3f} {meta['unit']}). \n" + f"Every step advances all {N_PROG} prognostic channels " + f"(temperature, salinity, currents over 19 depth levels, plus sea surface height)." + ) + progress(1.0, desc="Done") + return video, plot, summary + + +DESCRIPTION = """ +# 🌊 Samudra 2 — a neural ocean emulator + +[**Samudra 2**](https://huggingface.co/M2LInES/Samudra2) (M2LInES) is an autoregressive +deep-learning emulator of the **GFDL OM4 global ocean model**. From one ocean state plus +the atmospheric forcing, it advances *temperature, salinity, and horizontal currents across +19 depth levels together with sea surface height* — 77 prognostic fields at once — 10 days +per step, orders of magnitude faster than the ocean model it replaces. + +This Space runs the released **1° checkpoint** on genuinely held-out years (2014–2022): +the initial condition and forcing are streamed live from the public OM4 archive, the +emulator is rolled forward on ZeroGPU, and the result is scored against the OM4 +simulation itself. +""" + +with gr.Blocks(title="Samudra 2 ocean emulator") as demo: + gr.Markdown(DESCRIPTION) + + with gr.Row(): + with gr.Column(scale=1): + variable = gr.Dropdown( + choices=list(VARIABLES), value="Sea surface temperature (2.5 m)", + label="Field to visualise", + info="The emulator always advances the full ocean state.", + ) + start_date = gr.Dropdown( + choices=START_DATES, value="2015-06-02", + label="Initial condition (held-out period)", + ) + n_steps = gr.Slider( + 1, MAX_STEPS, value=6, step=1, + label="Autoregressive steps (10 days each)", + ) + run = gr.Button("Run the ocean emulator", variant="primary") + with gr.Accordion("Advanced", open=False): + fps = gr.Slider(1, 8, value=3, step=1, label="Animation frames per second") + summary = gr.Markdown() + with gr.Column(scale=2): + video = gr.Video(label="Emulator vs OM4", autoplay=True, loop=True) + plot = gr.Image(label="Error growth with lead time", type="filepath") + + inputs = [variable, start_date, n_steps, fps] + outputs = [video, plot, summary] + run.click(forecast, inputs=inputs, outputs=outputs, api_name="forecast") + + gr.Examples( + examples=[ + ["Sea surface temperature (2.5 m)", "2015-06-02", 6], + ["Sea surface height", "2016-02-07", 8], + ["Sea surface salinity (2.5 m)", "2018-11-04", 6], + ["Ocean temperature at 105 m", "2020-03-18", 10], + ], + inputs=[variable, start_date, n_steps], + outputs=outputs, + fn=forecast, + cache_examples=True, + cache_mode="lazy", + ) + + gr.Markdown( + "Data: public 1° GFDL-OM4 archive (`m2lines-pubs/Samudra`, Open Storage Network). " + "Model: [M2LInES/Samudra2](https://huggingface.co/M2LInES/Samudra2) · " + "code: [m2lines/Samudra](https://github.com/m2lines/Samudra) (Apache-2.0)." + ) + +if __name__ == "__main__": + demo.queue().launch(theme=gr.themes.Citrus(), mcp_server=True) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..84f91a4383cca2c25f6d37971ed50c284624c0d9 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,23 @@ +numpy>=1.26.4,<2 +pandas>=2.2 +scipy>=1.15.2 +matplotlib>=3.10.1 +imageio>=2.34 +imageio-ffmpeg>=0.5 +cftime>=1.6.4.post1 +dacite>=1.9.2 +dask>=2025.2,<2026 +einops>=0.7 +jaxtyping>=0.3 +microsoft-aurora>=1.8 +perceiver-pytorch>=0.8.8 +pydantic-settings>=2.8.1 +pyyaml>=6.0.2 +s3fs==2025.5.1 +torchinfo>=1.8 +tqdm>=4.67.1 +typing-extensions>=4.15 +wandb>=0.19.8 +xarray>=2025.1.2 +xarray-einstats>=0.8 +zarr<3 diff --git a/samudra/__init__.py b/samudra/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f38775e2ab78d90c201e65888708701088c08970 --- /dev/null +++ b/samudra/__init__.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("samudra") +except PackageNotFoundError: # not installed (e.g. running from a bare source tree) + __version__ = "0.0.0+unknown" + +__all__ = ["__version__"] diff --git a/samudra/aggregator/LICENSE-ACE.txt b/samudra/aggregator/LICENSE-ACE.txt new file mode 100644 index 0000000000000000000000000000000000000000..f49a4e16e68b128803cc2dcea614603632b04eac --- /dev/null +++ b/samudra/aggregator/LICENSE-ACE.txt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/samudra/aggregator/__init__.py b/samudra/aggregator/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..57e09016f926c41b855f55898c13af2f5bf66f52 --- /dev/null +++ b/samudra/aggregator/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from .inference import InferenceEvaluatorAggregator +from .loss import ( + get_channel_loss_dict, + get_channel_loss_scale_dict, + get_depth_loss_dict, + get_variable_loss_dict, +) +from .main import Aggregator +from .validate import ValidateAggregator diff --git a/samudra/aggregator/inference/__init__.py b/samudra/aggregator/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3e6c163b893bde5b5eacd192caa84045617ad6eb --- /dev/null +++ b/samudra/aggregator/inference/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from .main import InferenceEvaluatorAggregator diff --git a/samudra/aggregator/inference/main.py b/samudra/aggregator/inference/main.py new file mode 100644 index 0000000000000000000000000000000000000000..57d8a1448e4b67a6ec10798953f5676ed00de20f --- /dev/null +++ b/samudra/aggregator/inference/main.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch +import wandb +import xarray as xr + +from samudra.constants import TensorMap +from samudra.utils.data import Normalize, get_aggregator_dicts +from samudra.utils.output import ModelInferenceOutput +from samudra.utils.wandb import Metrics, MetricsDict + +from ..validate.reduced import MeanAggregator as OneStepMeanAggregator +from .reduced import MeanAggregator +from .time_mean import TimeMeanEvaluatorAggregator + + +class InferenceEvaluatorAggregator: + """ + Aggregates statistics for inference comparing a generated and target series. + """ + + def __init__( + self, + n_timesteps: int, + metadata: dict[str, dict[str, str]], + hist: int, + area_weights: torch.Tensor, + wet: torch.Tensor, + num_prognostic_channels: int, + normalize: Normalize, + tensor_map: TensorMap, + record_step_20: bool = True, + log_global_mean_time_series: bool = True, + log_global_mean_norm_time_series: bool = True, + time_mean_reference_data: xr.Dataset | None = None, + channel_mean_names: list[str] | None = None, + ): + """ + Args: + n_timesteps: Number of timesteps of inference that will be run. + metadata: Mapping of variable names their metadata that will + used in generating logged image captions. + hist: Number of timesteps of history. + area_weights: Area weights for the data. + wet: Wet mask for the data. + num_prognostic_channels: Number of prognostic channels in the data. + normalize: Normalization helper for prognostic channels. + tensor_map: Mapping from prognostic variables to tensor channels. + record_step_20: Whether to record the mean of the 20th steps. + log_global_mean_time_series: Whether to log global mean time series metrics. + log_global_mean_norm_time_series: Whether to log the normalized global mean + time series metrics. + time_mean_reference_data: Reference time means for computing bias stats. + channel_mean_names: List of channel names to compute the mean of. + """ + self._aggregators: dict[ + str, MeanAggregator | OneStepMeanAggregator | TimeMeanEvaluatorAggregator + ] = {} + self._time_dependent_aggregators: dict[str, TimeMeanEvaluatorAggregator] = {} + self._log_time_series = ( + log_global_mean_time_series or log_global_mean_norm_time_series + ) + if log_global_mean_time_series: + self._aggregators["mean"] = MeanAggregator( + target="denorm", + n_timesteps=n_timesteps, + metadata=metadata, + area_weights=area_weights, + ) + if log_global_mean_norm_time_series: + self._aggregators["mean_norm"] = MeanAggregator( + target="norm", + n_timesteps=n_timesteps, + metadata=metadata, + area_weights=area_weights, + ) + if record_step_20: + self._aggregators["mean_step_20"] = OneStepMeanAggregator( + target_time=20, + area_weights=area_weights, + ) + self._aggregators["time_mean"] = TimeMeanEvaluatorAggregator( + metadata=metadata, + area_weights=area_weights, + reference_means=time_mean_reference_data, + channel_mean_names=channel_mean_names, + ) + self._aggregators["time_mean_norm"] = TimeMeanEvaluatorAggregator( + metadata=metadata, + area_weights=area_weights, + target="norm", + reference_means=time_mean_reference_data, + channel_mean_names=channel_mean_names, + ) + + self._summary_aggregators = { + name: agg + for name, agg in list(self._aggregators.items()) + + list(self._time_dependent_aggregators.items()) + if name not in ["mean", "mean_norm"] + } + self._n_timesteps_seen = 0 + self._normalize = normalize + self._tensor_map = tensor_map + self.num_prognostic_channels = num_prognostic_channels + self.hist = hist + self.wet = wet + + @property + def log_time_series(self) -> bool: + return self._log_time_series + + @torch.no_grad() + def record_batch(self, data: ModelInferenceOutput): + if len(data.prediction) == 0: + raise ValueError("No prediction values in data") + if len(data.target) == 0: + raise ValueError("No target values in data") + total_len = len(data.time) + assert data.prediction.shape[0] == total_len // (self.hist + 1) + target_norm_dict, target_unnorm_dict = get_aggregator_dicts( + data.target, + normalize=self._normalize, + tensor_map=self._tensor_map, + wet=self.wet, + long_rollout=True, + input_type="prognostic", + num_prognostic_channels=self.num_prognostic_channels, + hist=self.hist, + ) + gen_norm_dict, gen_unnorm_dict = get_aggregator_dicts( + data.prediction, + normalize=self._normalize, + tensor_map=self._tensor_map, + wet=self.wet, + long_rollout=True, + input_type="prognostic", + num_prognostic_channels=self.num_prognostic_channels, + hist=self.hist, + ) + + for aggregator in self._aggregators.values(): + aggregator.record_batch( + target_data=target_unnorm_dict, + gen_data=gen_unnorm_dict, + target_data_norm=target_norm_dict, + gen_data_norm=gen_norm_dict, + i_time_start=self._n_timesteps_seen, + ) + # TODO: Add time-dependent aggregators + # for time_dependent_aggregator in self._time_dependent_aggregators.values(): + # time_dependent_aggregator.record_batch( + # time=data.time, + # target_data=target_unnorm_dict, + # gen_data=gen_unnorm_dict, + # ) + key = list(target_unnorm_dict.keys())[0] + n_times = target_unnorm_dict[key].shape[1] + logs = self._get_inference_logs_slice( + step_slice=slice(self._n_timesteps_seen, self._n_timesteps_seen + n_times), + ) + self._n_timesteps_seen += n_times + return logs + + def record_initial_prognostic( + self, + initial_prognostic: torch.Tensor, + ): + if self._n_timesteps_seen != 0: + raise RuntimeError( + "record_initial_condition may only be called once, " + "before recording any batches" + ) + + data_norm_dict, data_unnorm_dict = get_aggregator_dicts( + initial_prognostic, + normalize=self._normalize, + tensor_map=self._tensor_map, + wet=self.wet, + long_rollout=True, + input_type="input", + num_prognostic_channels=self.num_prognostic_channels, + hist=self.hist, + ) + for aggregator_name in ["mean", "mean_norm"]: + aggregator = self._aggregators.get(aggregator_name) + if aggregator is not None: + aggregator.record_batch( + target_data=data_unnorm_dict, + gen_data=data_unnorm_dict, + target_data_norm=data_norm_dict, + gen_data_norm=data_norm_dict, + i_time_start=0, + ) + key = list(data_unnorm_dict.keys())[0] + n_times = data_unnorm_dict[key].shape[1] + logs = self._get_inference_logs_slice( + step_slice=slice(0, n_times), + ) + self._n_timesteps_seen = n_times + return logs + + def get_summary_logs(self): + # These aggregators require a full timeseries of + # data and thus wandb logged at the end + logs: MetricsDict = {} + for name, aggregator in self._summary_aggregators.items(): + logs.update(aggregator.get_logs(label=name)) + return logs + + @torch.no_grad() + def _get_logs(self) -> Metrics: + """ + Returns logs as can be reported to WandB. + """ + logs: MetricsDict = {} + for name, aggregator in self._aggregators.items(): + logs.update(aggregator.get_logs(label=name)) + for name, time_dependent_aggregator in self._time_dependent_aggregators.items(): + logs.update(time_dependent_aggregator.get_logs(label=name)) + return logs + + @torch.no_grad() + def _get_inference_logs_slice(self, step_slice: slice): + """ + Returns a subset of the time series for applicable metrics + for a specific slice of as can be reported to WandB. + + Args: + step_slice: Timestep slice to determine the time series subset. + + Returns: + Tuple of start index and list of logs. + """ + logs = {} + for name, aggregator in self._aggregators.items(): + if isinstance(aggregator, MeanAggregator): + logs.update(aggregator.get_logs(label=name, step_slice=step_slice)) + return to_inference_logs(logs) + + +def to_inference_logs(log): + # we have a dictionary which contains WandB tables + # which we will convert to a list of dictionaries, one for each + # row in the tables. Any scalar values will be reported in the last + # dictionary. + n_rows = 0 + for val in log.values(): + if isinstance(val, wandb.Table): + n_rows = max(n_rows, len(val.data)) + logs: list[dict[str, float | int]] = [] + for i in range(max(1, n_rows)): # need at least one for non-series values + logs.append({}) + for key, val in log.items(): + if isinstance(val, wandb.Table): + for i, row in enumerate(val.data): + for j, col in enumerate(val.columns): + key_without_table_name = key[: key.rfind("/")] + logs[i][f"{key_without_table_name}/{col}"] = row[j] + else: + logs[-1][key] = val + return logs diff --git a/samudra/aggregator/inference/reduced.py b/samudra/aggregator/inference/reduced.py new file mode 100644 index 0000000000000000000000000000000000000000..a4226992b6efe2004658601cc7483d69cdeb3ca7 --- /dev/null +++ b/samudra/aggregator/inference/reduced.py @@ -0,0 +1,475 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +from collections import defaultdict +from functools import partial +from typing import Literal, Protocol + +import numpy as np +import torch +import wandb + +from samudra.aggregator.metrics import ( + area_weighted_gradient_magnitude_percent_diff, + area_weighted_mean, + area_weighted_mean_bias, + area_weighted_rmse, + area_weighted_std, +) +from samudra.utils.device import get_device +from samudra.utils.distributed import all_reduce_mean +from samudra.utils.wandb import Metrics + + +@dataclasses.dataclass +class _SeriesData: + metric_name: str + var_name: str + data: np.ndarray + + def get_wandb_key(self) -> str: + return f"{self.metric_name}/{self.var_name}" + + +def get_gen_shape(gen_data: dict[str, torch.Tensor]): + for name in gen_data: + return gen_data[name].shape + + +class MeanMetric(Protocol): + def record(self, target: torch.Tensor, gen: torch.Tensor, i_time_start: int): + """ + Update metric for a batch of data. + """ + ... + + def get(self) -> torch.Tensor: + """ + Get the total metric value, not divided by number of recorded batches. + """ + ... + + +class SingleTargetMeanMetric(Protocol): + def record(self, tensor: torch.Tensor, i_time_start: int): + """ + Update metric for a batch of data. + """ + ... + + def get(self) -> torch.Tensor: + """ + Get the total metric value, not divided by number of recorded batches. + """ + ... + + +class AreaWeightedFunction(Protocol): + """ + A function that computes a metric on the true and predicted values, + weighted by area. + """ + + def __call__( + self, + target: torch.Tensor, + gen: torch.Tensor, + ) -> torch.Tensor: ... + + +class AreaWeightedSingleTargetFunction(Protocol): + """ + A function that computes a metric on a single value, weighted by area. + """ + + def __call__( + self, + tensor: torch.Tensor, + ) -> torch.Tensor: ... + + +def compute_metric_on( + source: Literal["gen", "target"], metric: AreaWeightedSingleTargetFunction +) -> AreaWeightedFunction: + """Turns a single-target metric function + (computed on only the generated or target data) into a function that takes in + both the generated and target data as arguments, as required for the APIs + which call generic metric functions. + """ + + def metric_wrapper( + target: torch.Tensor, + gen: torch.Tensor, + ) -> torch.Tensor: + if source == "gen": + return metric(gen) + elif source == "target": + return metric(target) + + return metric_wrapper + + +class AreaWeightedReducedMetric: + """ + A wrapper around an area-weighted metric function. + """ + + def __init__( + self, + device: torch.device, + compute_metric: AreaWeightedFunction, + n_timesteps: int, + ): + self._compute_metric = compute_metric + self._total: torch.Tensor | None = None + self._n_batches = torch.zeros( + n_timesteps, dtype=torch.int32, device=get_device() + ) + self._device = device + self._n_timesteps = n_timesteps + + def record(self, target: torch.Tensor, gen: torch.Tensor, i_time_start: int): + """Add a batch of data to the metric. + + Args: + target: Target data. Should have shape [batch, time, height, width]. + gen: Generated data. Should have shape [batch, time, height, width]. + i_time_start: The index of the first timestep in the batch. + """ + time_dim = 1 + if target.shape != gen.shape: + raise RuntimeError( + "target and gen must have the same shape, got " + f"{target.shape} and {gen.shape}" + ) + new_value = ( + self._compute_metric(target=target, gen=gen).mean(dim=0).to(self._device) + ) + if self._total is None: + self._total = torch.zeros( + [self._n_timesteps], dtype=new_value.dtype, device=self._device + ) + time_slice = slice(i_time_start, i_time_start + gen.shape[time_dim]) + self._total[time_slice] += new_value + self._n_batches[time_slice] += 1 + + def get(self) -> torch.Tensor: + """Returns the mean metric across recorded batches.""" + if self._total is None: + return torch.tensor(torch.nan) + return self._total / self._n_batches + + +class MeanAggregator: + def __init__( + self, + target: Literal["norm", "denorm"], + n_timesteps: int, + area_weights: torch.Tensor, + metadata: dict[str, dict[str, str]] | None = None, + ): + self._variable_metrics: dict[str, dict[str, MeanMetric]] | None = None + self._shape_x = None + self._shape_y = None + self._target = target + self._n_timesteps = n_timesteps + self._area_weights = area_weights + + if metadata is None: + self._metadata: dict[str, dict[str, str]] = {} + else: + self._metadata = metadata + + def _get_variable_metrics(self, gen_data: dict[str, torch.Tensor]): + if self._variable_metrics is None: + self._variable_metrics = {} + + device = get_device() + self._variable_metrics["weighted_rmse"] = defaultdict( + lambda: AreaWeightedReducedMetric( + device=device, + compute_metric=partial( + area_weighted_rmse, area_weights=self._area_weights + ), + n_timesteps=self._n_timesteps, + ) + ) + if self._target == "denorm": + self._variable_metrics["weighted_grad_mag_percent_diff"] = defaultdict( + lambda: AreaWeightedReducedMetric( + device=device, + compute_metric=partial( + area_weighted_gradient_magnitude_percent_diff, + area_weights=self._area_weights, + ), # noqa: E501 + n_timesteps=self._n_timesteps, + ) + ) + self._variable_metrics["weighted_mean_gen"] = defaultdict( + lambda: AreaWeightedReducedMetric( + device=device, + compute_metric=compute_metric_on( + source="gen", + metric=lambda tensor: area_weighted_mean( + tensor, area_weights=self._area_weights + ), + ), + n_timesteps=self._n_timesteps, + ) + ) + self._variable_metrics["weighted_mean_target"] = defaultdict( + lambda: AreaWeightedReducedMetric( + device=device, + compute_metric=compute_metric_on( + source="target", + metric=lambda tensor: area_weighted_mean( + tensor, area_weights=self._area_weights + ), + ), + n_timesteps=self._n_timesteps, + ) + ) + self._variable_metrics["weighted_bias"] = defaultdict( + lambda: AreaWeightedReducedMetric( + device=device, + compute_metric=partial( + area_weighted_mean_bias, area_weights=self._area_weights + ), + n_timesteps=self._n_timesteps, + ) + ) + self._variable_metrics["weighted_std_gen"] = defaultdict( + lambda: AreaWeightedReducedMetric( + device=device, + compute_metric=compute_metric_on( + source="gen", + metric=( + lambda tensor: area_weighted_std( + tensor, area_weights=self._area_weights + ) + ), + ), + n_timesteps=self._n_timesteps, + ) + ) + return self._variable_metrics + + @torch.no_grad() + def record_batch( + self, + target_data: dict[str, torch.Tensor], + gen_data: dict[str, torch.Tensor], + target_data_norm: dict[str, torch.Tensor], + gen_data_norm: dict[str, torch.Tensor], + i_time_start: int = 0, + ): + if self._target == "norm": + target_data = target_data_norm + gen_data = gen_data_norm + variable_metrics = self._get_variable_metrics(gen_data) + for name in gen_data.keys(): + for metric in variable_metrics: + variable_metrics[metric][name].record( + target=target_data[name], + gen=gen_data[name], + i_time_start=i_time_start, + ) + + def _get_series_data(self, step_slice: slice | None = None) -> list[_SeriesData]: + """Converts internally stored variable_metrics to a list.""" + if self._variable_metrics is None: + raise ValueError("No batches have been recorded.") + data: list[_SeriesData] = [] + for metric in self._variable_metrics: + sorted_keys = sorted(list(self._variable_metrics[metric].keys())) + for key in sorted_keys: + arr = self._variable_metrics[metric][key].get().detach() + if step_slice is not None: + arr = arr[step_slice] + datum = _SeriesData( + metric_name=metric, + var_name=key, + data=all_reduce_mean(arr).cpu().numpy(), + ) + data.append(datum) + return data + + @torch.no_grad() + def get_logs(self, label: str, step_slice: slice | None = None): + """ + Returns logs as can be reported to WandB. + + Args: + label: Label to prepend to all log keys. + step_slice: Slice of forecast steps to log. + """ + logs = {} + series_data: dict[str, np.ndarray] = { + datum.get_wandb_key(): datum.data + for datum in self._get_series_data(step_slice) + } + init_step = 0 if step_slice is None else step_slice.start + table = data_to_table(series_data, init_step) + logs[f"{label}/series"] = table + return logs + + +def data_to_table(data: dict[str, np.ndarray], init_step: int = 0) -> wandb.Table: + """ + Convert a dictionary of 1-dimensional timeseries data to a wandb Table. + + Args: + data: dictionary of timeseries data. + init_step: initial step corresponding to the first row's "forecast_step" + """ + keys = sorted(list(data.keys())) + table = wandb.Table(columns=["forecast_step"] + keys) + if len(keys) > 0: + for i in range(len(data[keys[0]])): + row = [init_step + i] + for key in keys: + row.append(data[key][i]) + table.add_data(*row) + return table + + +class AreaWeightedSingleTargetReducedMetric: + """ + A wrapper around an area-weighted metric function on a single data source. + """ + + def __init__( + self, + device: torch.device, + compute_metric: AreaWeightedSingleTargetFunction, + n_timesteps: int, + ): + self._compute_metric = compute_metric + self._total: torch.Tensor | None = None + self._n_batches = torch.zeros( + n_timesteps, dtype=torch.int32, device=get_device() + ) + self._device = device + self._n_timesteps = n_timesteps + + def record(self, tensor: torch.Tensor, i_time_start: int): + """Add a batch of data to the metric. + + Args: + tensor: batch data. Should have shape [batch, time, height, width]. + i_time_start: The index of the first timestep in the batch. + """ + new_value = self._compute_metric(tensor).mean(dim=0).to(self._device) + if self._total is None: + self._total = torch.zeros( + [self._n_timesteps], dtype=new_value.dtype, device=self._device + ) + time_slice = slice(i_time_start, i_time_start + tensor.shape[1]) + self._total[time_slice] += new_value + self._n_batches[time_slice] += 1 + + def get(self) -> torch.Tensor: + """Returns the mean metric across recorded batches.""" + if self._total is None: + return torch.tensor(torch.nan) + return self._total / self._n_batches + + +class SingleTargetMeanAggregator: + def __init__( + self, + n_timesteps: int, + area_weights: torch.Tensor, + metadata: dict[str, dict[str, str]] | None = None, + ): + self._variable_metrics: dict[str, dict[str, SingleTargetMeanMetric]] | None = ( + None + ) + self._shape_x = None + self._shape_y = None + self._n_timesteps = n_timesteps + self._area_weights = area_weights + if metadata is None: + self._metadata: dict[str, dict[str, str]] = {} + else: + self._metadata = metadata + + def _get_variable_metrics(self, gen_data: dict[str, torch.Tensor]): + if self._variable_metrics is None: + self._variable_metrics = { + "weighted_mean_gen": {}, + "weighted_std_gen": {}, + } + + device = get_device() + self._variable_metrics["weighted_mean_gen"] = defaultdict( + lambda: AreaWeightedSingleTargetReducedMetric( + device=device, + compute_metric=lambda x: area_weighted_mean(x, self._area_weights), + n_timesteps=self._n_timesteps, + ) + ) + self._variable_metrics["weighted_std_gen"] = defaultdict( + lambda: AreaWeightedSingleTargetReducedMetric( + device=device, + compute_metric=lambda x: area_weighted_std(x, self._area_weights), + n_timesteps=self._n_timesteps, + ) + ) + + return self._variable_metrics + + @torch.no_grad() + def record_batch( + self, + data: dict[str, torch.Tensor], + i_time_start: int = 0, + ): + variable_metrics = self._get_variable_metrics(data) + for name in data.keys(): + for metric in variable_metrics: + variable_metrics[metric][name].record( + tensor=data[name], + i_time_start=i_time_start, + ) + + def _get_series_data(self, step_slice: slice | None = None) -> list[_SeriesData]: + """Converts internally stored variable_metrics to a list.""" + if self._variable_metrics is None: + raise ValueError("No batches have been recorded.") + data: list[_SeriesData] = [] + for metric in self._variable_metrics: + sorted_keys = sorted(list(self._variable_metrics[metric].keys())) + for key in sorted_keys: + arr = self._variable_metrics[metric][key].get().detach() + if step_slice is not None: + arr = arr[step_slice] + datum = _SeriesData( + metric_name=metric, + var_name=key, + data=all_reduce_mean(arr).cpu().numpy(), + ) + data.append(datum) + return data + + @torch.no_grad() + def get_logs(self, label: str, step_slice: slice | None = None) -> Metrics: + """ + Returns logs as can be reported to WandB. + + Args: + label: Label to prepend to all log keys. + step_slice: Slice of forecast steps to log. + """ + logs = {} + series_data: dict[str, np.ndarray] = { + datum.get_wandb_key(): datum.data + for datum in self._get_series_data(step_slice) + } + init_step = 0 if step_slice is None else step_slice.start + table = data_to_table(series_data, init_step) + logs[f"{label}/series"] = table + return logs diff --git a/samudra/aggregator/inference/time_mean.py b/samudra/aggregator/inference/time_mean.py new file mode 100644 index 0000000000000000000000000000000000000000..f703cb98cdf990c5a62627a5228a9ffa201d3a5f --- /dev/null +++ b/samudra/aggregator/inference/time_mean.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +from typing import Literal + +import matplotlib.pyplot as plt +import torch +import wandb +import xarray as xr + +from samudra.aggregator.metrics import area_weighted_mean_bias, area_weighted_rmse +from samudra.aggregator.plotting import get_cmap_limits, plot_imshow +from samudra.utils.distributed import all_reduce_mean +from samudra.utils.wandb import Metrics, MetricsDict + + +@dataclasses.dataclass +class _TargetGenPair: + name: str + target: torch.Tensor + gen: torch.Tensor + area_weights: torch.Tensor + + def bias(self): + return self.gen - self.target + + def rmse(self) -> float: + ret = float( + area_weighted_rmse( + gen=self.gen, + target=self.target, + area_weights=self.area_weights, + ) + .cpu() + .numpy() + ) + return ret + + def weighted_mean_bias(self) -> float: + return float( + area_weighted_mean_bias( + gen=self.gen, + target=self.target, + area_weights=self.area_weights, + ) + .cpu() + .numpy() + ) + + +def get_gen_shape(gen_data: dict[str, torch.Tensor]): + for name in gen_data: + return gen_data[name].shape + + +class TimeMeanAggregator: + _image_captions = { + "bias_map": "{name} time-mean bias (generated - reference) [{units}]", + "gen_map": "{name} time-mean generated [{units}]", + } + + def __init__( + self, + area_weights: torch.Tensor, + target: Literal["norm", "denorm"] = "denorm", + metadata: dict[str, dict[str, str]] | None = None, + reference_means: xr.Dataset | None = None, + ): + """ + Args: + area_weights: Tensor of shape [n_lat, n_lon] representing the area weights. + target: Whether to compute metrics on the normalized or denormalized data, + defaults to "denorm". + metadata: Mapping of variable names their metadata that will + used in generating logged image captions. + reference_means: Dataset containing reference time-mean values + for bias computation. + """ + self._target = target + if metadata is None: + self._metadata: dict[str, dict[str, str]] = {} + else: + self._metadata = metadata + # Dictionaries of tensors of shape [n_lat, n_lon] represnting time means + self._data: dict[str, torch.Tensor] | None = None + self._n_timesteps = 0 + self._n_samples: int | None = None + self._reference_means = reference_means + self._reference_validated = False + self._area_weights = area_weights + + @staticmethod + def _add_or_initialize_time_mean( + maybe_dict: dict[str, torch.Tensor] | None, + new_data: dict[str, torch.Tensor], + ignore_initial: bool = False, + ) -> dict[str, torch.Tensor]: + sample_dim = 0 + time_dim = 1 + if ignore_initial: + time_slice = slice(1, None) + else: + time_slice = slice(0, None) + if maybe_dict is None: + d: dict[str, torch.Tensor] = { + name: tensor[:, time_slice].sum(dim=time_dim).sum(dim=sample_dim) + for name, tensor in new_data.items() + } + else: + d = dict(maybe_dict) + for name, tensor in new_data.items(): + d[name] += tensor[:, time_slice].sum(dim=time_dim).sum(dim=sample_dim) + return d + + @torch.no_grad() + def record_batch( + self, + data: dict[str, torch.Tensor], + i_time_start: int = 0, + ): + ignore_initial = i_time_start == 0 + if ignore_initial: + raise Exception( + "This is wrong! Make sure this isnt called while recording " + "initial prognostic" + ) + self._data = self._add_or_initialize_time_mean(self._data, data, ignore_initial) + if self._n_samples is None: + self._n_samples = data[list(data)[0]].size(0) + if ignore_initial: + self._n_timesteps = data[list(data)[0]].size(1) - 1 + else: + self._n_timesteps += data[list(data)[0]].size(1) + if not self._reference_validated: + if self._reference_means is not None: + self.get_logs(label="") + self._reference_validated = True + + def get_data(self) -> dict[str, torch.Tensor]: + if self._n_timesteps == 0 or self._data is None: + raise ValueError("No data recorded.") + + assert self._n_samples is not None # for type checker + + ret = {} + names = sorted(list(self._data.keys())) # sort for rank-consistent order + for name in names: + value = self._data[name] + gen = all_reduce_mean(value / self._n_timesteps / self._n_samples) + ret[name] = gen + return ret + + @torch.no_grad() + def get_logs(self, label: str) -> Metrics: + logs: MetricsDict = {} + data = self.get_data() + gen_map_key = "gen_map" + for name, pred in data.items(): + vmin_pred, vmax_pred = get_cmap_limits(pred.cpu().numpy()) + prediction_image = wandb.Image( + plot_imshow(pred.cpu().numpy()), + caption=self._get_caption(gen_map_key, name, vmin_pred, vmax_pred), + ) + plt.close("all") + logs.update( + { + f"{gen_map_key}/{name}": prediction_image, + } + ) + if self._reference_means is not None and name in self._reference_means: + pair = _TargetGenPair( + name=name, + target=torch.as_tensor( + self._reference_means[name].values, device=pred.device + ), + gen=pred, + area_weights=self._area_weights, + ) + bias_map = pair.bias().cpu().numpy() + vmin_bias, vmax_bias = get_cmap_limits(bias_map, diverging=True) + cmap = plt.colormaps.get_cmap("RdBu_r") + cmap.set_bad(color=(0.7, 0.7, 0.7)) + + bias_fig: plt.Figure = plot_imshow( + bias_map, + vmin=vmin_bias, + vmax=vmax_bias, + cmap=cmap, + nan_padding=False, + ) + bias_image = wandb.Image( + bias_fig, + caption=self._get_caption("bias_map", name, vmin_bias, vmax_bias), + ) + logs.update({f"ref_bias/{name}": pair.weighted_mean_bias()}) + logs.update({f"ref_rmse/{name}": pair.rmse()}) + logs.update({f"ref_bias_map/{name}": bias_image}) + + if len(label) != 0: + return {f"{label}/{key}": logs[key] for key in logs} + return logs + + def _get_caption(self, key: str, name: str, vmin: float, vmax: float) -> str: + if name in self._metadata: + caption_name = self._metadata[name]["long_name"] + units = self._metadata[name]["units"] + else: + caption_name, units = name, "unknown_units" + caption = self._image_captions[key].format(name=caption_name, units=units) + caption += f" vmin={vmin:.4g}, vmax={vmax:.4g}." + return caption + + +class TimeMeanEvaluatorAggregator: + """Statistics and images on the time-mean state. + + This aggregator keeps track of the time-mean state, then computes + statistics and images on that time-mean state when logs are retrieved. + """ + + _image_captions = { + "bias_map": "{name} time-mean bias (generated - target) [{units}]", + "gen_map": "{name} time-mean generated [{units}]", + } + + def __init__( + self, + area_weights: torch.Tensor, + target: Literal["norm", "denorm"] = "denorm", + metadata: dict[str, dict[str, str]] | None = None, + reference_means: xr.Dataset | None = None, + channel_mean_names: list[str] | None = None, + ): + """ + Args: + area_weights: Tensor of shape [n_lat, n_lon] representing the area weights. + target: Whether to compute metrics on the normalized or denormalized data, + defaults to "denorm". + metadata: Mapping of variable names their metadata that will + used in generating logged image captions. + reference_means: Dataset containing reference time-mean values + for bias computation. + channel_mean_names: Names of variables whose RMSE will be averaged. If + not provided, all available variables will be used. + """ + self._target = target + if metadata is None: + self._metadata: dict[str, dict[str, str]] = {} + else: + self._metadata = metadata + self._area_weights = area_weights + # Dictionaries of tensors of shape [n_lat, n_lon] represnting time means + self._target_agg = TimeMeanAggregator( + target=target, metadata=metadata, area_weights=self._area_weights + ) + self._gen_agg = TimeMeanAggregator( + target=target, + metadata=metadata, + reference_means=reference_means, + area_weights=self._area_weights, + ) + self._channel_mean_names = channel_mean_names + + @torch.no_grad() + def record_batch( + self, + target_data: dict[str, torch.Tensor], + gen_data: dict[str, torch.Tensor], + target_data_norm: dict[str, torch.Tensor], + gen_data_norm: dict[str, torch.Tensor], + i_time_start: int = 0, + ): + if self._target == "norm": + target_data = target_data_norm + gen_data = gen_data_norm + self._target_agg.record_batch(target_data, i_time_start) + self._gen_agg.record_batch(gen_data, i_time_start) + + def _get_target_gen_pairs(self) -> list[_TargetGenPair]: + target_data = self._target_agg.get_data() + gen_data = self._gen_agg.get_data() + + ret = [] + for name in gen_data.keys(): + ret.append( + _TargetGenPair( + gen=gen_data[name], + target=target_data[name], + name=name, + area_weights=self._area_weights, + ) + ) + return ret + + @torch.no_grad() + def get_logs(self, label: str) -> Metrics: + logs = dict(self._gen_agg.get_logs("")) + preds = self._get_target_gen_pairs() + bias_map_key = "bias_map" + rmse_all_channels = {} + for pred in preds: + bias_data = pred.bias().cpu().numpy() + vmin_bias, vmax_bias = get_cmap_limits(bias_data, diverging=True) + cmap = plt.colormaps.get_cmap("RdBu_r") + cmap.set_bad(color=(0.7, 0.7, 0.7)) + bias_fig = plot_imshow( + bias_data, vmin=vmin_bias, vmax=vmax_bias, cmap=cmap, nan_padding=False + ) + bias_image = wandb.Image( + bias_fig, + caption=self._get_caption( + bias_map_key, pred.name, vmin_bias, vmax_bias + ), + ) + plt.close("all") + rmse_all_channels[pred.name] = pred.rmse() + logs.update({f"rmse/{pred.name}": rmse_all_channels[pred.name]}) + if self._target == "denorm": + logs.update( + { + f"{bias_map_key}/{pred.name}": bias_image, + f"bias/{pred.name}": pred.weighted_mean_bias(), + } + ) + if self._target == "norm": + metric_name = "rmse/channel_mean" + if self._channel_mean_names is None: + values_to_average = list(rmse_all_channels.values()) + else: + values_to_average = [ + rmse_all_channels[name] for name in self._channel_mean_names + ] + logs.update({metric_name: sum(values_to_average) / len(values_to_average)}) + + if len(label) != 0: + return {f"{label}/{key}": logs[key] for key in logs} + return logs + + def _get_caption(self, key: str, name: str, vmin: float, vmax: float) -> str: + if name in self._metadata: + caption_name = self._metadata[name]["long_name"] + units = self._metadata[name]["units"] + else: + caption_name, units = name, "unknown_units" + caption = self._image_captions[key].format(name=caption_name, units=units) + caption += f" vmin={vmin:.4g}, vmax={vmax:.4g}." + return caption diff --git a/samudra/aggregator/loss.py b/samudra/aggregator/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..6bca1e5198ac7b1d743a0bba4f9b992e9376313c --- /dev/null +++ b/samudra/aggregator/loss.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from samudra.constants import TensorMap + + +def get_depth_loss_dict( + label: str, + loss_per_channel: torch.Tensor, + *, + tensor_map: TensorMap, +) -> dict[str, torch.Tensor]: + metrics = {} + for depth in tensor_map.DEPTH_SET: + metrics[f"{label}/loss/depth/depth_{depth}_loss"] = loss_per_channel[ + tensor_map.DP_3D_IDX[depth] + ].mean() + return metrics + + +def get_variable_loss_dict( + label: str, + loss_per_channel: torch.Tensor, + *, + tensor_map: TensorMap, +) -> dict[str, torch.Tensor]: + metrics = {} + for variable in tensor_map.VAR_SET: + metrics[f"{label}/loss/variable/{variable}_loss"] = loss_per_channel[ + tensor_map.VAR_3D_IDX[variable] + ].mean() + return metrics + + +def get_channel_loss_dict( + label: str, + loss_per_channel: torch.Tensor, + *, + tensor_map: TensorMap, + loss_name: str = "loss", +) -> dict[str, torch.Tensor]: + return get_channel_dict(label, loss_name, loss_per_channel, tensor_map=tensor_map) + + +def get_channel_loss_scale_dict( + label: str, + loss_scale_per_channel: torch.Tensor, + *, + tensor_map: TensorMap, +) -> dict[str, torch.Tensor]: + return get_channel_dict( + label, + "loss_scale", + loss_scale_per_channel, + tensor_map=tensor_map, + ) + + +def get_channel_dict( + prefix: str, + measure: str, + per_channel: torch.Tensor, + *, + tensor_map: TensorMap, +) -> dict[str, torch.Tensor]: + metrics = {} + for i, channel in enumerate(tensor_map.prognostic_var_names): + metrics[f"{prefix}/{measure}/channel/{channel}_{measure}"] = per_channel[i] + return metrics diff --git a/samudra/aggregator/main.py b/samudra/aggregator/main.py new file mode 100644 index 0000000000000000000000000000000000000000..b09cb465bafbf2434a6caa13ca7c66335c6ca7dc --- /dev/null +++ b/samudra/aggregator/main.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +""" +The code in this directory is directly inspired by the ACE project by AI2. + +See the original repository at: https://github.com/ai2cm/ace/tree/39133c18524cda85965486ecdc8cb64aac06f4c3/fme/fme/ace/aggregator +""" + +import torch + +from samudra.aggregator.inference import InferenceEvaluatorAggregator +from samudra.aggregator.train import TrainAggregator +from samudra.aggregator.validate import ValidateAggregator +from samudra.aggregator.validate.map import MapAggregator +from samudra.aggregator.validate.reduced import MeanAggregator +from samudra.aggregator.validate.snapshot import SnapshotAggregator +from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator +from samudra.constants import TensorMap +from samudra.utils.data import Normalize + + +class Aggregator: + @staticmethod + def get_train_aggregator(tensor_map: TensorMap) -> TrainAggregator: + return TrainAggregator(tensor_map) + + @staticmethod + def get_validation_aggregator( + metadata: dict[str, dict[str, str]], + hist: int, + area_weights: torch.Tensor, + num_prognostic_channels: int, + tensor_map: TensorMap, + normalize: Normalize, + *, + include_image_aggregators: bool = True, + ) -> ValidateAggregator: + val_aggregators: dict[str, ValidateSubAggregator] = { + "reduced": MeanAggregator(area_weights, hist), + } + if include_image_aggregators: + val_aggregators.update( + { + "snapshot": SnapshotAggregator(metadata, hist), + "mean_map": MapAggregator(metadata, hist), + } + ) + + return ValidateAggregator( + val_aggregators, + hist=hist, + num_prognostic_channels=num_prognostic_channels, + tensor_map=tensor_map, + normalize=normalize, + ) + + @staticmethod + def get_inline_inference_aggregator( + n_timesteps: int, + metadata: dict[str, dict[str, str]], + hist: int, + area_weights: torch.Tensor, + wet: torch.Tensor, + num_prognostic_channels: int, + tensor_map: TensorMap, + normalize: Normalize, + channel_mean_names: list[str] | None = None, + ) -> InferenceEvaluatorAggregator: + return InferenceEvaluatorAggregator( + n_timesteps=n_timesteps, + metadata=metadata, + hist=hist, + area_weights=area_weights, + wet=wet, + num_prognostic_channels=num_prognostic_channels, + normalize=normalize, + tensor_map=tensor_map, + record_step_20=(n_timesteps > 20), + log_global_mean_time_series=False, + log_global_mean_norm_time_series=False, + channel_mean_names=channel_mean_names, + ) + + @staticmethod + def get_standalone_inference_aggregator( + n_timesteps: int, + metadata: dict[str, dict[str, str]], + hist: int, + area_weights: torch.Tensor, + wet: torch.Tensor, + num_prognostic_channels: int, + tensor_map: TensorMap, + normalize: Normalize, + channel_mean_names: list[str] | None = None, + ) -> InferenceEvaluatorAggregator: + return InferenceEvaluatorAggregator( + n_timesteps=n_timesteps, + metadata=metadata, + hist=hist, + area_weights=area_weights, + wet=wet, + num_prognostic_channels=num_prognostic_channels, + normalize=normalize, + tensor_map=tensor_map, + record_step_20=(n_timesteps > 20), + log_global_mean_time_series=True, + log_global_mean_norm_time_series=True, + channel_mean_names=channel_mean_names, + ) diff --git a/samudra/aggregator/metrics.py b/samudra/aggregator/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..f54afe5ead547cafd8616b45f86e420675e58b4c --- /dev/null +++ b/samudra/aggregator/metrics.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def weighted_sum( + tensor: torch.Tensor, + weights: torch.Tensor | None = None, + dim: tuple[int, ...] = (-2, -1), + keepdim: bool = False, +) -> torch.Tensor: + """Computes the weighted sum across the specified list of dimensions. + + Args: + tensor: torch.Tensor + weights: Weights to apply to the sum. + dim: Dimensions to compute the sum over. + keepdim: Whether the output tensor has `dim` retained or not. + + Returns: + a tensor of the weighted sum averaged over the specified dimensions `dim`. + """ + if weights is None: + return tensor.nansum(dim=dim, keepdim=keepdim) + + weights = weights.to(tensor.device) + return (tensor * weights).nansum(dim=dim, keepdim=keepdim) + + +def area_weighted_sum( + data: torch.Tensor, + area_weights: torch.Tensor, + dim: tuple[int, ...] = (-2, -1), + keepdim: bool = False, +) -> torch.Tensor: + return weighted_sum(data, area_weights, dim=dim, keepdim=keepdim) + + +def weighted_mean( + tensor: torch.Tensor, + weights: torch.Tensor | None = None, + dim: tuple[int, ...] = (-2, -1), + keepdim: bool = False, +) -> torch.Tensor: + """Computes the weighted mean across the specified list of dimensions. + + This function treats NaNs as land cells and does not include them in the + computation. + + Args: + tensor: torch.Tensor + weights: Weights to apply to the mean. + dim: Dimensions to compute the mean over. + keepdim: Whether the output tensor has `dim` retained or not. + + Returns: + a tensor of the weighted mean averaged over the specified dimensions `dim`. + """ + if weights is None: + return tensor.nanmean(dim=dim, keepdim=keepdim) + + weights = weights.to(tensor.device) + denom = torch.where(torch.isnan(tensor), 0.0, weights.expand(tensor.shape)).sum( + dim=dim, keepdim=keepdim + ) + return (tensor * weights).nansum(dim=dim, keepdim=keepdim) / denom + + +def area_weighted_mean( + data: torch.Tensor, + area_weights: torch.Tensor, + dim: tuple[int, ...] = (-2, -1), + keepdim: bool = False, +) -> torch.Tensor: + return weighted_mean(data, area_weights, dim=dim, keepdim=keepdim) + + +def area_weighted_std( + data: torch.Tensor, + area_weights: torch.Tensor, + keepdim: bool = False, +): + return area_weighted_mean( + (data - area_weighted_mean(data, area_weights, keepdim=True)) ** 2, + area_weights, + keepdim=keepdim, + ).sqrt() + + +def area_weighted_rmse( + target: torch.Tensor, gen: torch.Tensor, area_weights: torch.Tensor +) -> torch.Tensor: + area_weights = area_weights.to(target.device) + return torch.sqrt(area_weighted_mean((gen - target) ** 2, area_weights)) + + +def area_weighted_mean_bias( + target: torch.Tensor, gen: torch.Tensor, area_weights: torch.Tensor +) -> torch.Tensor: + area_weights = area_weights.to(target.device) + return area_weighted_mean(gen - target, area_weights) + + +def gradient_magnitude( + tensor: torch.Tensor, dim: tuple[int, ...] = (-2, -1) +) -> torch.Tensor: + """Compute the magnitude of gradient across the specified dimensions.""" + gradients = torch.gradient(tensor, dim=dim) + return torch.sqrt(sum([g**2 for g in gradients])) # type: ignore + + +def weighted_mean_gradient_magnitude( + tensor: torch.Tensor, + weights: torch.Tensor | None = None, + dim: tuple[int, ...] = (-2, -1), +) -> torch.Tensor: + """Compute weighted mean of gradient magnitude across the specified dimensions.""" + return weighted_mean(gradient_magnitude(tensor, dim), weights=weights, dim=dim) + + +def gradient_magnitude_percent_diff( + target: torch.Tensor, + gen: torch.Tensor, + weights: torch.Tensor | None = None, + dim: tuple[int, ...] = (-2, -1), +) -> torch.Tensor: + """Compute the percent difference of the weighted mean gradient magnitude across + the specified dimensions. + """ + target_grad_mag = weighted_mean_gradient_magnitude(target, weights, dim) + gen_grad_mag = weighted_mean_gradient_magnitude(gen, weights, dim) + return 100 * (gen_grad_mag - target_grad_mag) / target_grad_mag + + +def area_weighted_gradient_magnitude_percent_diff( + target: torch.Tensor, gen: torch.Tensor, area_weights: torch.Tensor +): + area_weights = area_weights.to(target.device) + return gradient_magnitude_percent_diff(target, gen, weights=area_weights) diff --git a/samudra/aggregator/plotting.py b/samudra/aggregator/plotting.py new file mode 100644 index 0000000000000000000000000000000000000000..b9dc4cd7b232ce0ccf71f8bed85f94175a3d6453 --- /dev/null +++ b/samudra/aggregator/plotting.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import gc + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import Colormap +from matplotlib.figure import Figure + +from samudra.utils.wandb import WandBLogger + + +def get_cmap_limits(data: np.ndarray, diverging=False) -> tuple[float, float]: + vmin = np.nanmin(data) + vmax = np.nanmax(data) + if diverging: + vmax = max(abs(vmin), abs(vmax)) + vmin = -vmax + return vmin, vmax + + +def plot_imshow( + data: np.ndarray, + vmin: float | None = None, + vmax: float | None = None, + cmap: Colormap | None = None, + flip_lat: bool = True, + use_colorbar: bool = True, + nan_padding: bool = True, +) -> Figure: + """Plot a 2D array using imshow, ensuring figure size is same as array size.""" + min_ = np.nanmin(data) if vmin is None else vmin + max_ = np.nanmax(data) if vmax is None else vmax + + if flip_lat: + lat_dim = -2 + data = np.flip(data, axis=lat_dim) + + if use_colorbar: + height, width = data.shape + colorbar_width = max(1, int(0.025 * width)) + range_ = np.linspace(min_, max_, height) + range_ = np.repeat(range_[:, np.newaxis], repeats=colorbar_width, axis=1) + range_ = np.flipud(range_) # wandb images start from top (and left) + padding = np.zeros((height, colorbar_width)) + if nan_padding: + padding = padding + np.nan # Set when using non-diverging map + data = np.concatenate((data, padding, range_), axis=1) + + # make figure size (in pixels) be the same as array size + figsize = np.array(data.T.shape) / plt.rcParams["figure.dpi"] + fig = Figure(figsize=figsize) # create directly for cleanup when it leaves scope + ax = fig.add_axes((0, 0, 1, 1)) + ax.imshow(data, cmap=cmap, vmin=vmin, vmax=vmax) + ax.set_axis_off() + return fig + + +def plot_paneled_data( + data: list[list[np.ndarray]], + diverging: bool, + caption: str | None = None, +): + """Plot a list of 2D data arrays in a paneled plot.""" + if diverging: + cmap = plt.colormaps.get_cmap("RdBu_r") + cmap.set_bad(color=(0.7, 0.7, 0.7)) + else: + cmap = plt.colormaps.get_cmap("viridis") + cmap.set_bad(color="white") + vmin = np.inf + vmax = -np.inf + for row in data: + for arr in row: + vmin = min(vmin, np.nanmin(arr)) + vmax = max(vmax, np.nanmax(arr)) + if diverging: + vmax = max(abs(vmin), abs(vmax)) + vmin = -vmax + if caption is not None: + caption += " " + else: + caption = "" + + caption += f"vmin={vmin:.4g}, vmax={vmax:.4g}." + + if diverging: + fill_value = 0.5 * (vmin + vmax) + else: + fill_value = vmin + all_data = _stitch_data_panels(data, fill_value=fill_value) + + fig = plot_imshow( + all_data, vmin=vmin, vmax=vmax, cmap=cmap, nan_padding=not diverging + ) + wandb = WandBLogger.get_instance() + wandb_image = wandb.Image(fig, caption=caption) + plt.close(fig) + + # necessary to avoid CUDA error in some contexts + # see https://github.com/ai2cm/full-model/issues/740#issuecomment-2086546187 + gc.collect() + + return wandb_image + + +def _stitch_data_panels(data: list[list[np.ndarray]], fill_value) -> np.ndarray: + for row in data: + if len(row) != len(data[0]): + raise ValueError("All rows must have the same number of panels.") + + n_rows = len(data) + n_cols = len(data[0]) + for row in data: + for arr in row: + if arr.shape != data[0][0].shape: + raise ValueError("All panels must have the same shape.") + + stitched_data = np.full( + ( + n_rows * data[0][0].shape[0] + n_rows - 1, + n_cols * data[0][0].shape[1] + n_cols - 1, + ), + fill_value=fill_value, + ) + + # iterate over rows backwards, as the image starts in the bottom left + # and moves upwards + for i, row in enumerate(reversed(data)): + for j, arr in enumerate(row): + start_row = i * (arr.shape[0] + 1) + end_row = start_row + arr.shape[0] + start_col = j * (arr.shape[1] + 1) + end_col = start_col + arr.shape[1] + stitched_data[start_row:end_row, start_col:end_col] = arr + + return stitched_data diff --git a/samudra/aggregator/train.py b/samudra/aggregator/train.py new file mode 100644 index 0000000000000000000000000000000000000000..70c894fa918c5a7265a0ca0ef75f1fd0f9108b20 --- /dev/null +++ b/samudra/aggregator/train.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from samudra.aggregator.loss import ( + get_channel_loss_dict, + get_depth_loss_dict, + get_variable_loss_dict, +) +from samudra.constants import TensorMap +from samudra.utils.distributed import all_reduce_mean +from samudra.utils.output import TrainBatchOutput +from samudra.utils.wandb import Metrics + + +class TrainAggregator: + """Aggregates train statistics for an epoch.""" + + def __init__(self, tensor_map: TensorMap): + self.tensor_map = tensor_map + self._n_batches = 0 + self._loss = torch.tensor(torch.nan) + self._loss_per_channel = torch.tensor(torch.nan) + + @torch.no_grad() + def record_batch(self, batch: TrainBatchOutput): + if self._n_batches == 0: + self._loss = batch.loss + self._loss_per_channel = batch.loss_per_channel + else: + self._loss += batch.loss + self._loss_per_channel += batch.loss_per_channel + self._n_batches += 1 + + @torch.no_grad() + def get_logs(self, label: str = "train") -> Metrics: + loss = self._loss / self._n_batches + + loss_per_channel = self._loss_per_channel / self._n_batches + depth_loss_dict = get_depth_loss_dict( + label, loss_per_channel, tensor_map=self.tensor_map + ) + var_loss_dict = get_variable_loss_dict( + label, loss_per_channel, tensor_map=self.tensor_map + ) + channel_loss_dict = get_channel_loss_dict( + label, loss_per_channel, tensor_map=self.tensor_map + ) + logs = { + f"{label}/mean/loss": loss, + **depth_loss_dict, + **var_loss_dict, + **channel_loss_dict, + } + meaned_logs = {} + for key in sorted(logs.keys()): + meaned_logs[key] = float(all_reduce_mean(logs[key].detach()).cpu().numpy()) + + return meaned_logs diff --git a/samudra/aggregator/validate/__init__.py b/samudra/aggregator/validate/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..50da82846549496336fea2b7894db09e2939e289 --- /dev/null +++ b/samudra/aggregator/validate/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from .main import ValidateAggregator diff --git a/samudra/aggregator/validate/main.py b/samudra/aggregator/validate/main.py new file mode 100644 index 0000000000000000000000000000000000000000..0f920db6518add13132c7edb7d462b45fe6aac09 --- /dev/null +++ b/samudra/aggregator/validate/main.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from samudra.aggregator.train import TrainAggregator +from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator +from samudra.constants import TensorMap +from samudra.utils.data import Normalize, get_aggregator_dicts +from samudra.utils.output import ValBatchOutput +from samudra.utils.wandb import Metrics, MetricsDict + + +class ValidateAggregator(TrainAggregator): + """Aggregates Validation Statistics.""" + + def __init__( + self, + aggregators: dict[str, ValidateSubAggregator], + hist: int, + num_prognostic_channels: int, + *, + tensor_map: TensorMap, + normalize: Normalize, + ): + super().__init__(tensor_map) + self._aggregators = aggregators + self.hist = hist + self.num_prognostic_channels = num_prognostic_channels + self.normalize = normalize + + # TODO(jder): we could remove this by moving from inheritance + # to composition with the TrainAggregator functionality. + def record_batch(self, batch): + raise NotImplementedError( + "Call record_validation_batch instead of record_batch" + ) + + @torch.no_grad() + def record_validation_batch(self, batch: ValBatchOutput): + super().record_batch(batch) # Record losses + + # If there are no log aggregators, omit doing any extra work. + if not self._aggregators: + return + + # Translate the GridContext mask by removing history. + target_data = batch.target_data # [B, C*(hist+1), H, W] + wet = batch.ctx.label_mask # [C*(hist+1), H, W] + assert wet.shape == target_data.shape[1:], ( + "The wetmask must match the target data shape excluding batch." + ) + assert wet.shape[0] % (self.hist + 1) == 0, ( + "The wetmask channel count must be divisible by history size." + ) + first_wetmask_chunk = wet.shape[0] // (self.hist + 1) + wet = wet[:first_wetmask_chunk] # [C, H, W] + + if len(target_data) == 0: + raise ValueError("No data in target_data") + if len(batch.gen_data) == 0: + raise ValueError("No data in gen_data") + + assert target_data.shape[1] == self.num_prognostic_channels + target_data_dict, target_data_unnorm_dict = get_aggregator_dicts( + target_data, + normalize=self.normalize, + tensor_map=self.tensor_map, + wet=wet, + long_rollout=False, + input_type="prognostic", + num_prognostic_channels=self.num_prognostic_channels, + hist=self.hist, + ) + + gen_data_dict, gen_data_unnorm_dict = get_aggregator_dicts( + batch.gen_data, + normalize=self.normalize, + tensor_map=self.tensor_map, + wet=wet, + long_rollout=False, + input_type="prognostic", + num_prognostic_channels=self.num_prognostic_channels, + hist=self.hist, + ) + input_data_dict, input_data_unnorm_dict = get_aggregator_dicts( + batch.input_data, + normalize=self.normalize, + tensor_map=self.tensor_map, + wet=wet, + long_rollout=False, + input_type="input", + num_prognostic_channels=self.num_prognostic_channels, + hist=self.hist, + ) + + for agg in self._aggregators.values(): + agg.record_batch( + loss=batch.loss, + target_data=target_data_unnorm_dict, + gen_data=gen_data_unnorm_dict, + input_data=input_data_unnorm_dict, + target_data_norm=target_data_dict, + gen_data_norm=gen_data_dict, + input_data_norm=input_data_dict, + ) + + @torch.no_grad() + def get_logs(self, label: str = "train") -> Metrics: + logs: MetricsDict = dict(super().get_logs(label)) + for agg_label in self._aggregators: + for k, v in self._aggregators[agg_label].get_logs(label=agg_label).items(): + logs[f"{label}/{k}"] = v + + return logs diff --git a/samudra/aggregator/validate/map.py b/samudra/aggregator/validate/map.py new file mode 100644 index 0000000000000000000000000000000000000000..436913c94dc4ccbe3ab301193f8ca2d0c1f46009 --- /dev/null +++ b/samudra/aggregator/validate/map.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from samudra.aggregator.plotting import plot_paneled_data +from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator +from samudra.utils.distributed import all_reduce_mean, is_main_process + + +class MapAggregator(ValidateSubAggregator): + """ + An aggregator that records the average over batches as function of lat and lon. + """ + + _captions = { + "full-field": ( + "{name} one step mean full field; " + "(top) generated and (bottom) target [{units}]" + ), + "error": ( + "{name} one step mean full field error (generated - target) [{units}]" + ), + } + + def __init__( + self, metadata: dict[str, dict[str, str]] | None = None, hist: int = 0 + ): + """ + Args: + metadata: Mapping of variable names their metadata that will + used in generating logged image captions. + hist: Number of history steps to include in the snapshot. + """ + if metadata is None: + metadata = {} + else: + self._metadata = metadata + self._n_batches = 0 + self._target_data: dict[str, torch.Tensor] = {} + self._gen_data: dict[str, torch.Tensor] = {} + self.hist = hist + + @torch.no_grad() + def record_batch( + self, + loss, + target_data, + gen_data, + input_data, + target_data_norm, + gen_data_norm, + input_data_norm, + ): + self._loss = loss + for name in target_data: + if name in self._target_data: + self._target_data[name] += target_data[name].mean(dim=0) + else: + self._target_data[name] = target_data[name].mean(dim=0) + for name in gen_data: + if name in self._gen_data: + self._gen_data[name] += gen_data[name].mean(dim=0) + else: + self._gen_data[name] = gen_data[name].mean(dim=0) + self._n_batches += 1 + + @torch.no_grad() + def get_logs(self, label: str): + """ + Returns logs as can be reported to WandB. + + Args: + label: Label to prepend to all log keys. + """ + time_dim = 0 + target_time = self.hist # Use latest time step + image_logs = {} + render_images = is_main_process() + sorted_names = sorted(list(self._gen_data.keys())) + for name in sorted_names: + # use first sample in batch + gen = ( + all_reduce_mean( + self._gen_data[name].select(dim=time_dim, index=target_time) + ) + / self._n_batches + ) + target = ( + all_reduce_mean( + self._target_data[name].select(dim=time_dim, index=target_time) + ) + / self._n_batches + ) + if not render_images: + continue + image_logs[f"image-error/{name}"] = plot_paneled_data( + [[(gen - target).cpu().numpy()]], + diverging=True, + caption=self._get_caption("error", name), + ) + image_logs[f"image-full-field/{name}"] = plot_paneled_data( + [ + [gen.cpu().numpy()], + [target.cpu().numpy()], + ], + diverging=False, + caption=self._get_caption("full-field", name), + ) + image_logs = {f"{label}/{key}": image_logs[key] for key in image_logs} + return image_logs + + def _get_caption(self, key: str, name: str) -> str: + if name in self._metadata: + caption_name = self._metadata[name]["long_name"] + units = self._metadata[name]["units"] + else: + caption_name, units = name, "unknown_units" + caption = self._captions[key].format(name=caption_name, units=units) + return caption diff --git a/samudra/aggregator/validate/reduced.py b/samudra/aggregator/validate/reduced.py new file mode 100644 index 0000000000000000000000000000000000000000..70b74783d5e7f67bac5b96323968540d67092475 --- /dev/null +++ b/samudra/aggregator/validate/reduced.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Callable +from functools import partial + +import numpy as np +import torch + +from samudra.aggregator.metrics import ( + area_weighted_gradient_magnitude_percent_diff, + area_weighted_mean_bias, + area_weighted_rmse, +) +from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator +from samudra.utils.device import get_device +from samudra.utils.distributed import all_reduce_mean + + +class AreaWeightedReducedMetric: + """ + A wrapper around an area-weighted metric function. + """ + + def __init__( + self, + device: torch.device, + compute_metric: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + ): + self._compute_metric = compute_metric + self._total: torch.Tensor | None = None + self._device = device + + def record(self, target: torch.Tensor, gen: torch.Tensor): + """Add a batch of data to the metric. + + Args: + target: Target data. Should have shape [batch, time, height, width]. + gen: Generated data. Should have shape [batch, time, height, width]. + """ + new_value = self._compute_metric(target, gen).mean(dim=0) + if self._total is None: + self._total = torch.zeros_like(new_value, device=self._device) + self._total += new_value + + def get(self) -> torch.Tensor: + """Returns the metric.""" + assert self._total is not None + return self._total + + +class MeanAggregator(ValidateSubAggregator): + """ + Aggregator for mean-reduced metrics. + + These are metrics such as means which reduce to a single float for each batch, + and then can be averaged across batches to get a single float for the + entire dataset. This is important because the aggregator uses the mean to combine + metrics across batches and processors. + """ + + def __init__(self, area_weights: torch.Tensor, target_time: int): + self._n_batches = 0 + self._variable_metrics: dict | None = None + self._target_time = target_time + self._area_weights = area_weights + + def _get_variable_metrics(self, gen_data): + if self._variable_metrics is None: + self._variable_metrics = { + "weighted_rmse": {}, + "weighted_bias": {}, + "weighted_grad_mag_percent_diff": {}, + } + device = get_device() + for key in gen_data: + self._variable_metrics["weighted_rmse"][key] = ( + AreaWeightedReducedMetric( + device=device, + compute_metric=partial( + area_weighted_rmse, area_weights=self._area_weights + ), + ) + ) + self._variable_metrics["weighted_bias"][key] = ( + AreaWeightedReducedMetric( + device=device, + compute_metric=partial( + area_weighted_mean_bias, area_weights=self._area_weights + ), + ) + ) + self._variable_metrics["weighted_grad_mag_percent_diff"][key] = ( + AreaWeightedReducedMetric( + device=device, + compute_metric=partial( + area_weighted_gradient_magnitude_percent_diff, + area_weights=self._area_weights, + ), # noqa: E501 + ) + ) + + return self._variable_metrics + + @torch.no_grad() + def record_batch( + self, + *, + loss: torch.Tensor = torch.tensor(np.nan), + target_data, + gen_data, + target_data_norm, + gen_data_norm, + input_data: dict[str, torch.Tensor] | None = None, + input_data_norm: dict[str, torch.Tensor] | None = None, + i_time_start: int = 0, + ): + if input_data is None: + input_data = {} + if input_data_norm is None: + input_data_norm = {} + + variable_metrics = self._get_variable_metrics(gen_data) + time_dim = 1 + time_len = gen_data[list(gen_data.keys())[0]].shape[time_dim] + target_time = self._target_time - i_time_start + if target_time >= 0 and time_len > target_time: + for name in gen_data.keys(): + target = target_data[name].select(dim=time_dim, index=target_time) + gen = gen_data[name].select(dim=time_dim, index=target_time) + for metric in variable_metrics: + variable_metrics[metric][name].record( + target=target, + gen=gen, + ) + # only increment n_batches if we actually recorded a batch + self._n_batches += 1 + + def _get_data(self): + if self._variable_metrics is None or self._n_batches == 0: + raise ValueError("No batches have been recorded.") + data: dict[str, torch.Tensor] = {} + for metric in self._variable_metrics: + for key in self._variable_metrics[metric]: + data[f"{metric}/{key}"] = ( + self._variable_metrics[metric][key].get() / self._n_batches + ) + meaned_data: dict[str, float] = {} + for key in sorted(data.keys()): + meaned_data[key] = float(all_reduce_mean(data[key].detach()).cpu().numpy()) + return meaned_data + + @torch.no_grad() + def get_logs(self, label: str): + """ + Returns logs as can be reported to WandB. + + Args: + label: Label to prepend to all log keys. + """ + return { + f"{label}/{key}": data for key, data in sorted(self._get_data().items()) + } diff --git a/samudra/aggregator/validate/snapshot.py b/samudra/aggregator/validate/snapshot.py new file mode 100644 index 0000000000000000000000000000000000000000..a6e7557e3b6e7f5989cc58794766530ebeac732e --- /dev/null +++ b/samudra/aggregator/validate/snapshot.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from samudra.aggregator.plotting import plot_paneled_data +from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator +from samudra.utils.distributed import is_main_process +from samudra.utils.wandb import Metrics + + +class SnapshotAggregator(ValidateSubAggregator): + """ + An aggregator that records the first sample of the last batch of data. + """ + + _captions = { + "full-field": ( + "{name} one step full field for last sample; " + "(top) generated and (bottom) target [{units}]" + ), + "residual": ( + "{name} one step residual (prediction - previous time) for last sample; " + "(top) generated and (bottom) target [{units}]" + ), + "error": ( + "{name} one step full field error (generated - target) " + "for last sample [{units}]" + ), + } + + def __init__( + self, metadata: dict[str, dict[str, str]] | None = None, hist: int = 0 + ): + """ + Args: + metadata: Mapping of variable names their metadata that will + used in generating logged image captions. + hist: Number of history steps to include in the snapshot. + """ + if metadata is None: + metadata = {} + else: + self._metadata = metadata + self.hist = hist + + @torch.no_grad() + def record_batch( + self, + loss: torch.Tensor, + target_data, + gen_data, + input_data, + target_data_norm, + gen_data_norm, + input_data_norm, + ): + self._loss = loss + self._target_data = target_data + self._gen_data = gen_data + self._target_data_norm = target_data_norm + self._gen_data_norm = gen_data_norm + self._input_data = input_data + self._input_data_norm = input_data_norm + + @torch.no_grad() + def get_logs(self, label: str) -> Metrics: + """ + Returns logs as can be reported to WandB. + + Args: + label: Label to prepend to all log keys. + """ + if not is_main_process(): + return {} + + time_dim = 1 + target_time = 0 # first output time step + input_time = self.hist # last input time step + image_logs = {} + for name in self._gen_data.keys(): + # use first sample in batch + gen = self._gen_data[name].select(dim=time_dim, index=target_time)[0].cpu() + target = ( + self._target_data[name].select(dim=time_dim, index=target_time)[0].cpu() + ) + input = ( + self._input_data[name].select(dim=time_dim, index=input_time)[0].cpu() + ) + images = {} + images["error"] = [[(gen - target).numpy()]] + images["full-field"] = [[gen.numpy()], [target.numpy()]] + images["residual"] = [[(gen - input).numpy()], [(target - input).numpy()]] + for key, data in images.items(): + if key == "error" or key == "residual": + diverging = True + else: + diverging = False + caption = self._get_caption(key, name) + wandb_image = plot_paneled_data(data, diverging, caption=caption) + image_logs[f"image-{key}/{name}"] = wandb_image + image_logs = {f"{label}/{key}": image_logs[key] for key in image_logs} + return image_logs + + def _get_caption(self, key: str, name: str) -> str: + if name in self._metadata: + caption_name = self._metadata[name]["long_name"] + units = self._metadata[name]["units"] + else: + caption_name, units = name, "unknown_units" + caption = self._captions[key].format(name=caption_name, units=units) + return caption diff --git a/samudra/aggregator/validate/sub_aggregator.py b/samudra/aggregator/validate/sub_aggregator.py new file mode 100644 index 0000000000000000000000000000000000000000..f24944b3ba866666a25445e82a717f848033e46d --- /dev/null +++ b/samudra/aggregator/validate/sub_aggregator.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod + +import torch + +from samudra.utils.wandb import Metrics + + +class ValidateSubAggregator(ABC): + @abstractmethod + def get_logs(self, label: str) -> Metrics: ... + + @abstractmethod + def record_batch( + self, + *, + loss: torch.Tensor, + target_data, + gen_data, + input_data, + target_data_norm, + gen_data_norm, + input_data_norm, + ): ... diff --git a/samudra/backend.py b/samudra/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..6517915dd986865534278093121bc366fa5a4d39 --- /dev/null +++ b/samudra/backend.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import logging + +import torch + +logger = logging.getLogger(__name__) + +import os + +from samudra.config import DistributedConfig, EvalBackendConfig, TrainBackendConfig +from samudra.utils.device import set_device +from samudra.utils.distributed import init_distributed_mode + + +def _cuda_diagnostics() -> str: + """Return a short string that helps debug why CUDA is (not) available.""" + visible = os.getenv("CUDA_VISIBLE_DEVICES", None) + return ( + "CUDA diagnostics: " + f"torch={torch.__version__}, " + f"torch.version.cuda={torch.version.cuda}, " + f"torch.cuda.is_available()={torch.cuda.is_available()}, " + f"torch.cuda.device_count()={torch.cuda.device_count()}, " + f"CUDA_VISIBLE_DEVICES={visible!r}" + ) + + +def init_train_backend( + backend: TrainBackendConfig, +) -> tuple[torch.device, DistributedConfig | None]: + """Given backend config, get the device and (if any) distributed configuration.""" + match backend: + case "cpu": + device = torch.device("cpu") + dist_cfg = None + case "cuda": + if not torch.cuda.is_available(): + raise RuntimeError( + "Requested backend='cuda' but CUDA is not available. " + + _cuda_diagnostics() + ) + device = torch.device("cuda") + dist_cfg = None + case "nccl": + if not torch.cuda.is_available(): + raise RuntimeError( + "Requested backend='nccl' but CUDA is not available. " + + _cuda_diagnostics() + ) + device = torch.device("cuda") + dist_cfg = init_distributed_mode() + case "auto" if torch.cuda.is_available(): + logger.info("auto backend detected CUDA") + device = torch.device("cuda") + try: + dist_cfg = init_distributed_mode() + logger.info("succeeded in initializing distributed mode") + except RuntimeError as e: + logger.info( + f"Failed to initialize distributed mode, running on single GPU.", + exc_info=e, + ) + dist_cfg = None + case "auto": + logger.warning( + "auto backend: cuda not found, using CPU. " + _cuda_diagnostics() + ) + device = torch.device("cpu") + dist_cfg = None + case _: + raise ValueError(f"Invalid backend: {backend}") + + # We set this globally so we don't need to hand the device around. + # See https://github.com/m2lines/Samudra/issues/87. + set_device(device) + + return device, dist_cfg + + +def init_eval_backend(backend: EvalBackendConfig) -> torch.device: + """Given evaluation backend config, get the device.""" + match backend: + case "cpu": + device = torch.device("cpu") + case "cuda": + if not torch.cuda.is_available(): + raise RuntimeError( + "Requested backend='cuda' but CUDA is not available. " + + _cuda_diagnostics() + ) + device = torch.device("cuda") + case "auto" if torch.cuda.is_available(): + logger.info("auto backend detected CUDA") + device = torch.device("cuda") + case "auto": + logger.warning( + "auto backend: cuda not found, using CPU. " + _cuda_diagnostics() + ) + device = torch.device("cpu") + case _: + raise ValueError(f"Invalid backend: {backend}") + + set_device(device) + return device diff --git a/samudra/cli.py b/samudra/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..16dd3717d83e1abbe1ddcf73fad0e0e8adc99f20 --- /dev/null +++ b/samudra/cli.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Console-script entry point: ``samudra CONFIG [OVERRIDES...]``. + +Installed as the ``samudra`` command (see ``[project.scripts]`` in +pyproject.toml), so a user who ``pip install samudra`` can run +``samudra train path/to/config.yaml`` without cloning the repo. Each +subcommand forwards to the same ``main`` the module entry points use +(``python -m samudra.train`` etc.), so behaviour is identical. +""" + +import sys + +_COMMANDS = ("train", "eval", "viz") + +_HELP = """\ +samudra — train and evaluate emulators of ocean physics + +Samudra is a PyTorch package for models that auto-regressively predict future +ocean states — temperature, salinity, horizontal currents, sea-surface height, +and surface heat flux — learned from the OM4 ocean model at 1°, 1/2°, and 1/4° +resolution. See https://arxiv.org/abs/2412.03795 for the method. + +Usage: + samudra CONFIG [--KEY VALUE ...] + +Commands: + train Train a model from a config (checkpointing, W&B logging, multi-GPU). + eval Roll a trained model out autoregressively and collect metrics. + viz Render maps, time series, and PDFs from evaluation outputs. + +CONFIG is a path to a YAML file or the name of a bundled preset such as +`samudra_om4/train.yaml`. Any config key can be overridden inline (e.g. +`--epochs 100`); run `samudra --help` for a command's full options. + +Examples: + samudra train samudra_om4/train.yaml --experiment.data_root ./data + samudra eval samudra_om4/eval.yaml --ckpt_path ./checkpoint.pt + +Docs: https://m2lines.github.io/Samudra/docs/ +""" + + +def main() -> None: + """Dispatch ``samudra ...`` to the matching task entry point.""" + argv = sys.argv[1:] + if not argv or argv[0] in ("-h", "--help"): + sys.stdout.write(_HELP) + raise SystemExit(0) + if argv[0] not in _COMMANDS: + sys.stderr.write(f"samudra: unknown command {argv[0]!r}\n\n{_HELP}") + raise SystemExit(2) + + command, rest = argv[0], argv[1:] + # The config loaders parse sys.argv directly, so strip the subcommand token + # and leave them an argv of `CONFIG [OVERRIDES...]`. + sys.argv = [f"samudra {command}", *rest] + + if command == "train": + from samudra.train import main as train_main + + train_main() + elif command == "eval": + from samudra.eval import main as eval_main + + eval_main() + else: # viz + from samudra.viz.config import VizConfig + from samudra.viz.config import main as viz_main + + viz_main(VizConfig.from_yaml_and_cli()) + + +if __name__ == "__main__": + main() diff --git a/samudra/config.py b/samudra/config.py new file mode 100644 index 0000000000000000000000000000000000000000..b1778fb62ea0d2d27d4b5f20629173a12216d12d --- /dev/null +++ b/samudra/config.py @@ -0,0 +1,1339 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import abc +import datetime +from functools import cached_property +from pathlib import Path +from typing import Annotated, Literal, Self, assert_never + +import cftime +import numpy as np +import pandas as pd +import pydantic +import torch +import xarray as xr +from perceiver_pytorch import Perceiver as NaivePerceiver +from pydantic import ( + Field, + PlainSerializer, + PlainValidator, + WithJsonSchema, + model_validator, +) +from torch import nn +from torch.nn import GELU + +from samudra.config_base import BaseConfig, TopLevelConfig +from samudra.constants import ( + DatasetSpec, + GridType, + LoaderVersion, + build_llc_spec, + build_om4_spec, +) +from samudra.models import Samudra, SamudraMini, SamudraMulti +from samudra.models.base import BaseModel +from samudra.models.modules import ( + AvgPool, + BilinearUpsample, + CappedGELU, + ConvBlock, + ConvNeXtBlock, + CoreBlock, + CoreBlockBuilder, + MaxPool, + PerceiverDecoder, + PerceiverEncoder, + ReLU, + TransposedConvUpsample, + UNetBackbone, +) +from samudra.models.modules.augment_input import ( + Concat3dCoordinates, + FourierFeatures2D, + fourier_features_2d_dim, +) +from samudra.models.modules.blocks import ZonallyPeriodicBilinearUpsample +from samudra.models.modules.encoder import patch_from +from samudra.utils.data import DataContainer, DataSource, DataSourceSplits +from samudra.utils.llc import canonicalize_llc_datasets +from samudra.utils.location import ( + LocalLocation, + Location, + ResolvedLocation, + UnresolvedLocation, +) +from samudra.utils.loss import ( + DynamicLoss, + GradientLoss, + LossFnWithContext, + LossMetric, + loss_fn_from_metric, +) +from samudra.utils.profiler import Profiler +from samudra.utils.schedule import SchedulerConfig + + +class WandBConfig(BaseConfig): + mode: Literal["online", "disabled"] = "disabled" + project: str = "default" + entity: str = "ocean_emulators" + group: str | None = None + tags: list[str] | None = None + notes: str | None = None + + +class JulianDate: + """Represents a Julian date as a cftime.datetime at noon on the relevant day. + + This is the format the OM4 data uses, so we match that here. + TODO(jder): probably worth asserting the date format when opening the data. + """ + + datetime: cftime.datetime + + def __init__(self, value: str): + parsed = cftime.datetime.strptime(value, "%Y-%m-%d", calendar="julian") + self.datetime = parsed.replace(hour=12) + + def __str__(self) -> str: + return self.datetime.strftime("%Y-%m-%d") + + +def _julian_date_validator(value: str | JulianDate) -> JulianDate: + return JulianDate(value) if isinstance(value, str) else value + + +JulianDateConfig = Annotated[ + JulianDate, + PlainValidator(_julian_date_validator), + PlainSerializer(JulianDate.__str__), + WithJsonSchema({"type": "string", "format": "date"}), +] + + +# We reuse pydantic's AwareDatetime to match JSONSchema's expected RFC3339 format +_aware_datetime_adapter = pydantic.TypeAdapter(pydantic.AwareDatetime) + + +def _datetime64_validator(value: str | np.datetime64) -> np.datetime64: + if isinstance(value, np.datetime64): + return value.astype("datetime64[ns]") + parsed = _aware_datetime_adapter.validate_python(value) + utc = parsed.astimezone(datetime.UTC).replace(tzinfo=None) + return np.datetime64(utc, "ns") + + +def _serialize_datetime64(value: np.datetime64) -> str: + return str(np.datetime_as_string(value, unit="ns", timezone="UTC")) + + +LlcDatetimeConfig = Annotated[ + np.datetime64, + PlainValidator(_datetime64_validator), + PlainSerializer(_serialize_datetime64), + WithJsonSchema({"type": "string", "format": "date-time"}), +] + + +class Om4TimeConfig(BaseConfig): + model_config = pydantic.ConfigDict(frozen=True) + + start: JulianDateConfig + end: JulianDateConfig + + @property + def time_slice(self) -> slice: + return slice(self.start.datetime, self.end.datetime) + + def overlaps(self, other: Self) -> bool: + return ( + self.start.datetime < other.end.datetime + and self.end.datetime > other.start.datetime + ) + + def __str__(self) -> str: + return f"{self.start} to {self.end}" + + +class LlcTimeConfig(BaseConfig): + model_config = pydantic.ConfigDict(frozen=True) + + start: LlcDatetimeConfig + end: LlcDatetimeConfig + + @property + def time_slice(self) -> slice: + return slice(self.start, self.end) + + def overlaps(self, other: Self) -> bool: + return bool(self.start < other.end and self.end > other.start) + + def __str__(self) -> str: + return f"{self.start} to {self.end}" + + +TimeConfig = Om4TimeConfig | LlcTimeConfig + + +LOCATION_DOCS = ( + "Use a string relative to the `data_root` or use a structured location " + "see location.py for possible types." +) + + +class BaseDataSourceConfig[SourceTimeConfigT: TimeConfig](BaseConfig, abc.ABC): + train_time: SourceTimeConfigT = Field(frozen=True) + val_time: SourceTimeConfigT = Field(frozen=True) + inference_times: tuple[SourceTimeConfigT, ...] = Field(default=(), frozen=True) + data_location: Location = Field( + description="Location of the data; " + LOCATION_DOCS + ) + data_means_location: Location = Field( + description="Location of the data means; " + LOCATION_DOCS + ) + data_stds_location: Location = Field( + description="Location of the data standard deviations; " + LOCATION_DOCS + ) + + @property + @abc.abstractmethod + def dataset_spec(self) -> DatasetSpec: + raise NotImplementedError + + @abc.abstractmethod + def canonicalize_datasets( + self, + data: xr.Dataset, + means: xr.Dataset, + stds: xr.Dataset, + ) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + raise NotImplementedError + + def build( + self, + data_root: ResolvedLocation, + *, + use_dask: bool, + is_primary: bool, + ) -> DataSourceSplits: + source = self._build_source( + data_root, + turn_on_dask=use_dask, + ) + inference_source = None + if is_primary and self.inference_times: + if use_dask: + full_inference_source = source + else: + full_inference_source = self._build_source( + data_root, + turn_on_dask=True, + ) + # TODO: remove multiple inference time ranges altogether (see #813) + assert len(self.inference_times) == 1, ( + "multiple inference time ranges have been deprecated" + ) + inference_source = full_inference_source.slice(self.inference_times[0]) + + return DataSourceSplits( + train=source.slice(self.train_time), + val=source.slice(self.val_time), + inference=inference_source, + ) + + def _build_source( + self, + data_root: ResolvedLocation, + *, + turn_on_dask: bool, + ) -> DataSource: + resolved_data_location = data_root.resolve(self.data_location) + resolved_means_location = data_root.resolve(self.data_means_location) + resolved_stds_location = data_root.resolve(self.data_stds_location) + + chunks: dict[str, int] | None = {} if turn_on_dask else None + data = resolved_data_location.open(chunks) + means = resolved_means_location.open(chunks) + stds = resolved_stds_location.open(chunks) + data, means, stds = self.canonicalize_datasets( + data, + means, + stds, + ) + dataset_spec = self.dataset_spec + + source = DataSource.from_datasets( + data, + means, + stds, + dataset_spec=dataset_spec, + prognostic_var_names=dataset_spec.prognostic_var_names, + boundary_var_names=dataset_spec.boundary_var_names, + name=f"{resolved_data_location}-{turn_on_dask}", + ) + return source + + +class BaseDataLoadingConfig(BaseConfig): + def num_pytorch_workers(self) -> int: + raise NotImplementedError + + def persistent_pytorch_workers(self) -> bool: + raise NotImplementedError + + +class CpuDataLoadingConfig(BaseDataLoadingConfig): + type: Literal["cpu"] = "cpu" + num_workers: int = Field(default=4, ge=0) + persistent_workers: bool = True + + def num_pytorch_workers(self) -> int: + return self.num_workers + + def persistent_pytorch_workers(self) -> bool: + return self.persistent_workers + + +class GpuDataLoadingConfig(BaseDataLoadingConfig): + type: Literal["gpu"] = "gpu" + kvikio_task_size: int = Field(default=64 * 1024 * 1024, gt=0) + kvikio_num_threads: int = Field(default=8, gt=0) + + def num_pytorch_workers(self) -> int: + # When loading data direct to GPU, we don't want worker processes. + # 0 means "load in the main process" + return 0 + + def persistent_pytorch_workers(self) -> bool: + return False + + +DataLoadingConfig = Annotated[ + CpuDataLoadingConfig | GpuDataLoadingConfig, + Field(discriminator="type"), +] + + +class Om4DataSourceConfig(BaseDataSourceConfig[Om4TimeConfig]): + type: Literal["om4"] = "om4" + prognostic_vars_key: str = "thermo_dynamic_all" + boundary_vars_key: str = "tau_hfds" + grid_type: GridType = "gaussian" + + @property + def dataset_spec(self) -> DatasetSpec: + return build_om4_spec( + self.prognostic_vars_key, + self.boundary_vars_key, + grid_type=self.grid_type, + ) + + @pydantic.model_validator(mode="after") + def validate_time_splits(self) -> Self: + if self.train_time.overlaps(self.val_time): + raise ValueError( + f"Training time range {self.train_time} overlaps " + f"with validation time range {self.val_time}" + ) + return self + + def canonicalize_datasets( + self, + data: xr.Dataset, + means: xr.Dataset, + stds: xr.Dataset, + ) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + # TODO: move OM4 canonicalization in here instead of in validation afte. + return data, means, stds + + +class LlcDataSourceConfig(BaseDataSourceConfig[LlcTimeConfig]): + type: Literal["llc"] = "llc" + prognostic_vars_key: str = "single_1" + boundary_vars_key: str = "single_1" + face: int = Field(default=1, ge=0) + i_start: int = Field(default=0, ge=0) + i_end: int = Field(default=720, gt=0) + j_start: int = Field(default=0, ge=0) + j_end: int = Field(default=720, gt=0) + + @property + def dataset_spec(self) -> DatasetSpec: + return build_llc_spec( + self.prognostic_vars_key, + self.boundary_vars_key, + ) + + @pydantic.model_validator(mode="after") + def validate_time_splits(self) -> Self: + if self.train_time.overlaps(self.val_time): + raise ValueError( + f"Training time range {self.train_time} overlaps " + f"with validation time range {self.val_time}" + ) + return self + + @pydantic.model_validator(mode="after") + def validate_crop_bounds(self) -> Self: + if self.i_end <= self.i_start: + raise ValueError("LLC crop bounds must satisfy i_start < i_end") + if self.j_end <= self.j_start: + raise ValueError("LLC crop bounds must satisfy j_start < j_end") + return self + + def canonicalize_datasets( + self, + data: xr.Dataset, + means: xr.Dataset, + stds: xr.Dataset, + ) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + return canonicalize_llc_datasets( + data, + means, + stds, + face=self.face, + i_start=self.i_start, + i_end=self.i_end, + j_start=self.j_start, + j_end=self.j_end, + dataset_spec=self.dataset_spec, + ) + + +DataSourceConfig = Annotated[ + Om4DataSourceConfig | LlcDataSourceConfig, + Field(discriminator="type"), +] + + +class DataConfig(BaseConfig): + sources: list[DataSourceConfig] = Field( + description=( + "Data sources to include, each with explicit data/means/stds " + "locations. These are resolved relative to data_root." + ), + min_length=1, + ) + loading: DataLoadingConfig = Field(default_factory=CpuDataLoadingConfig) + hist: int = 1 + loader_version: str = str(LoaderVersion.OM4_TORCH.value) + normalize_before_mask: bool = True + masked_fill_value: float = 0.0 + concurrent_compute: bool = False + + def build( + self, + data_root: ResolvedLocation, + ) -> DataContainer: + loader_version = LoaderVersion(self.loader_version) + use_dask = loader_version != LoaderVersion.OM4_TORCH + dataset_spec = self.sources[0].dataset_spec + assert all(source.dataset_spec == dataset_spec for source in self.sources), ( + "All data sources must use the same dataset spec" + ) + + source_splits = [ + source_cfg.build( + data_root, + use_dask=use_dask, + is_primary=index == 0, + ) + for index, source_cfg in enumerate(self.sources) + ] + train_sources = [splits.train for splits in source_splits] + val_sources = [splits.val for splits in source_splits] + + return DataContainer( + train_sources=train_sources, + val_sources=val_sources, + inference_source=source_splits[0].inference, + loader_version=loader_version, + dataset_spec=dataset_spec, + ) + + +BlockType = Literal["conv_next_block", "conv_block"] +ActivationType = Literal["relu", "gelu", "capped_gelu"] +NormType = Literal["batch", "instance", "layer"] + + +class BlockConfig(BaseConfig): + block_type: BlockType = "conv_next_block" + kernel_size: int = 3 + activation: ActivationType = "capped_gelu" + upscale_factor: int = 4 + norm: NormType = "batch" + pointwise_linear: bool = False + + def build(self) -> CoreBlockBuilder: + match self.activation: + case "relu": + activation: type[nn.Module] = ReLU + case "capped_gelu": + activation = CappedGELU + case "gelu": + activation = GELU + case _: + assert_never(self.activation) + + def create_block( + in_channels: int, + out_channels: int, + dilation: int, + n_layers: int, + pad: str, + checkpoint_simple: bool, + ) -> CoreBlock: + match self.block_type: + case "conv_block": + return ConvBlock( + in_channels=in_channels, + out_channels=out_channels, + dilation=dilation, + n_layers=n_layers, + pad=pad, + checkpoint_simple=checkpoint_simple, + kernel_size=self.kernel_size, + activation=activation, + ) + case "conv_next_block": + return ConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + dilation=dilation, + n_layers=n_layers, + pad=pad, + checkpoint_simple=checkpoint_simple, + kernel_size=self.kernel_size, + upscale_factor=self.upscale_factor, + norm=self.norm, + activation=activation, + pointwise_linear=self.pointwise_linear, + ) + case _: + assert_never(self.block_type) + + return create_block + + +PerceiverImpl = Literal["auto", "naive", "flash"] + + +class PerceiverConfig(BaseConfig): + """A standard config interface to various perceiver implementations. + + Builds either a regular Perceiver (for the encoder, via ``build``) or a + PerceiverIO (for the decoder, via ``build_io``). Both respect the shared + ``implementation`` setting from ``SamudraMultiConfig.perceiver_implementation``. + """ + + depth: int = 6 + latent_dim: int = Field( + default=128, + description="The small, latent dimension of the Perceiver. This is the `N` dimension for the Perceiver's `O(M*N)` complexity", + ) + num_latents: int = Field( + default=512, + description="The number of latent vectors in the Perceiver. This is the `M` dimension for the Perceiver's `O(M*N)` complexity", + ) + + def build( + self, + in_channels: int, + out_channels: int, + max_patch_size: tuple[int, int], + implementation: PerceiverImpl, + ) -> nn.Module: + """Build a regular Perceiver (used by the encoder).""" + # This is not really a "frequency" but a maximum of the width appears to be reasonable from looking at the code. + max_freq = max(*max_patch_size) + + num_freq_bands = 4 + fourier_dim = fourier_features_2d_dim(num_freq_bands) + # Use the same explicit 2D Fourier features in both implementations so + # intra-patch positions are encoded equivalently. + if _use_flash(implementation): + try: + from flash_perceiver import Perceiver as FlashPerceiver # type: ignore + except ImportError as e: + raise _flash_import_error() from e + from einops.layers.torch import Rearrange + + perceiver: nn.Module = nn.Sequential( + FourierFeatures2D(num_freq_bands=num_freq_bands, max_freq=max_freq), + Rearrange("b ph pw v -> b (ph pw) v"), + FlashPerceiver( + depth=self.depth, + input_dim=in_channels + fourier_dim, + output_dim=out_channels, + output_mode="average", + latent_dim=self.latent_dim, + num_latents=self.num_latents, + use_flash_attn=True, + weight_tie_layers=True, + self_per_cross_attn=2, + ), + ) + elif _use_naive(implementation): + perceiver = nn.Sequential( + FourierFeatures2D(num_freq_bands=num_freq_bands, max_freq=max_freq), + NaivePerceiver( + # Required by perceiver-pytorch even when its internal + # Fourier encoding is disabled below. + num_freq_bands=num_freq_bands, + max_freq=max_freq, + depth=self.depth, + input_axis=2, + input_channels=in_channels + fourier_dim, + num_classes=out_channels, + latent_dim=self.latent_dim, + num_latents=self.num_latents, + weight_tie_layers=True, + fourier_encode_data=False, + self_per_cross_attn=2, + ), + ) + else: + raise ValueError(f"Unknown perceiver implementation: {implementation}.") + + return perceiver + + def build_io( + self, + in_channels: int, + queries_dim: int, + out_channels: int, + implementation: PerceiverImpl, + ) -> nn.Module: + """Build a PerceiverIO (used by the decoder).""" + if _use_flash(implementation): + try: + from flash_perceiver.perceiver import ( # type: ignore + PerceiverIO as FlashPerceiverIO, # type: ignore + ) + except ImportError as e: + raise _flash_import_error() from e + perceiver_io: nn.Module = FlashPerceiverIO( + depth=self.depth, + input_dim=in_channels, + query_dim=queries_dim, + proj_dim=out_channels, + num_latents=self.num_latents, + latent_dim=self.latent_dim, + use_flash_attn=True, + weight_tie_layers=True, + ) + elif _use_naive(implementation): + from perceiver_pytorch.perceiver_io import PerceiverIO as NaivePerceiverIO + + perceiver_io = NaivePerceiverIO( + depth=self.depth, + dim=in_channels, + queries_dim=queries_dim, + logits_dim=out_channels, + num_latents=self.num_latents, + latent_dim=self.latent_dim, + weight_tie_layers=True, + decoder_ff=True, + ) + else: + raise ValueError(f"Unknown perceiver implementation: {implementation}.") + + return perceiver_io + + +def _use_flash(implementation: PerceiverImpl) -> bool: + return ( + implementation == "auto" and torch.cuda.is_available() + ) or implementation == "flash" + + +def _use_naive(implementation: PerceiverImpl) -> bool: + return ( + implementation == "auto" and not torch.cuda.is_available() + ) or implementation == "naive" + + +def _flash_import_error() -> ValueError: + return ValueError( + "`implementation==flash` or flash was automatically chosen for `implementation==auto`, " + "but the flash attention dependencies could not be imported. " + "Please run `uv sync --extra cuda` or specify the `naive` attention implementation." + ) + + +class EncoderConfig(BaseConfig): + perceiver: PerceiverConfig = PerceiverConfig() + + def build( + self, + in_channels: int, + out_channels: int, + patch_extent: tuple[float, float], + max_lat_size: int, + max_lon_size: int, + implementation: PerceiverImpl, + ) -> PerceiverEncoder: + max_patch_size = patch_from(patch_extent, max_lat_size, max_lon_size) + return PerceiverEncoder( + in_channels=in_channels, + out_channels=out_channels, + patch_extent=patch_extent, + perceiver=self.perceiver.build( + in_channels, out_channels, max_patch_size, implementation + ), + ) + + +class DecoderConfig(BaseConfig): + """A PerceiverIO-based decoder configuration. + + Uses PerceiverIO (with an explicit query mechanism) rather than a regular + Perceiver. Output pixel positions are encoded as queries, so the output + size is determined by the query count — not by ``num_latents``. + + When ``window_patches`` is set, the decoder tiles the output grid into + spatial blocks of that many patches per side. Each block's PerceiverIO + call receives only the overlapping latent tokens plus ``context_patches`` + extra rings of neighbors, keeping cost bounded even when the latent grid + is large (i.e. fine ``patch_extent``). + """ + + perceiver: PerceiverConfig = PerceiverConfig() + queries_dim: int = Field( + default=64, + description="Embedding dimension for pixel-position queries in the PerceiverIO decoder head.", + ) + window_patches: int | None = Field( + default=4096, + description="Side length (in patches) of each spatial decode window. " + "None = decode all patches at once (global attention). " + "E.g. window_patches=8 means each PerceiverIO call covers an 8x8 block of patches.", + ) + context_patches: int | None = Field( + default=1, + description="Number of extra patch rings around each window to include as data context. " + "Only used when window_patches is set. None = full context (every window sees all latent tokens).", + ) + + def build( + self, + in_channels: int, + out_channels: int, + patch_extent: tuple[float, float], + implementation: PerceiverImpl, + ) -> PerceiverDecoder: + return PerceiverDecoder( + in_channels=in_channels, + out_channels=out_channels, + patch_extent=patch_extent, + queries_dim=self.queries_dim, + perceiver_io=self.perceiver.build_io( + in_channels, self.queries_dim, out_channels, implementation + ), + window_patches=self.window_patches, + context_patches=self.context_patches, + ) + + +DownSamplingBlocks = Literal["avg_pool", "max_pool"] +UpSamplingBlocks = Literal[ + "bilinear_upsample", "transposed_conv", "zonally_periodic_upsample" +] +Checkpointing = Literal["all", "simple"] + + +class UNetBackboneConfig(BaseConfig): + ch_width: list[int] = [200, 250, 300, 400] + dilation: list[int] = [1, 2, 4, 8] + n_layers: list[int] = [1, 1, 1, 1] + core_block: BlockConfig = BlockConfig() + down_sampling_block: DownSamplingBlocks = "avg_pool" + up_sampling_block: UpSamplingBlocks = "zonally_periodic_upsample" + drop_path_rate: float = Field( + default=0.0, + description="Shortcut dropout rate. The chance we turn off skip connections in the UNet. Reasonable values are 0.1-0.3. Use 0.0 to disable.", + ) + + def build( + self, + in_channels: int, + pad: str, + checkpointing: Checkpointing | None, + ) -> UNetBackbone: + assert len(self.ch_width) == len(self.dilation) == len(self.n_layers), ( + "`ch_width`, `dilation`, and `n_layers` must have the same length." + ) + + def create_upsampling_block(in_channels: int, out_channels: int): + match self.up_sampling_block: + case "bilinear_upsample": + return BilinearUpsample( + in_channels=in_channels, out_channels=out_channels + ) + case "transposed_conv": + return TransposedConvUpsample( + in_channels=in_channels, out_channels=out_channels + ) + case "zonally_periodic_upsample": + return ZonallyPeriodicBilinearUpsample() + case _: + assert_never(self.up_sampling_block) + + match self.down_sampling_block: + case "avg_pool": + downsampling_block: nn.Module = AvgPool() + case "max_pool": + downsampling_block = MaxPool() + case _: + assert_never(self.down_sampling_block) + + return UNetBackbone( + in_channels=in_channels, + ch_width=self.ch_width, + dilation=self.dilation, + n_layers=self.n_layers, + pad=pad, + create_block=self.core_block.build(), + downsampling_block=downsampling_block, + create_upsampling_block=create_upsampling_block, + checkpointing=checkpointing, + drop_path_rate=self.drop_path_rate, + ) + + +class BaseModelConfig(BaseConfig, abc.ABC): + pred_residuals: bool = False + last_kernel_size: int = 3 + pad: str = "circular" + + checkpointing: Checkpointing | None = Field( + default=None, + description="""Strategy for storing activations for the model for use in + the backward pass. If not set, the model will store all activations in memory + (fast but lots of memory). If set to 'all', the model will recompute each + top-level layer (CoreBlocks, scaling layers, etc.) in the backward pass. + If set to 'simple', the model will recompute only cheap layers like scales + and nonlinearities.""", + ) + + gradient_detach_interval: int = Field( + default=0, + description="""Interval for detaching gradients in autoregressive training. `0` means no detaching.""", + ) + + add_3d_coordinates: bool = Field( + default=False, + description="Add 3d coordinates representing position on the Earth (cartesian coordinates on a unit sphere) to the input channels.", + ) + + @abc.abstractmethod + def build( + self, + prog_channels: int, + boundary_channels: int, + out_channels: int, + hist: int, + srcs: list[DataSource], + ) -> BaseModel: + pass + + +class SamudraConfig(BaseModelConfig): + unet: UNetBackboneConfig = UNetBackboneConfig() + pos_channels: int = Field( + default=0, + description="""Number of channels used for a learned positional embedding""", + ) + use_bfloat16: bool = Field( + default=False, + description="Use bfloat16 for most layers rather than float32.", + ) + + def build( + self, + prog_channels: int, + boundary_channels: int, + out_channels: int, + hist: int, + srcs: list[DataSource], + ) -> Samudra: + if len(srcs) != 1: + raise ValueError( + "Samudra only supports training at a single scale! Please configure exactly one data source." + ) + src = srcs[0] + in_channels = prog_channels + boundary_channels + total_in_channels = ( + in_channels + self.pos_channels + (3 if self.add_3d_coordinates else 0) + ) + add_3d_coordinates = Concat3dCoordinates() if self.add_3d_coordinates else None + return Samudra( + in_channels=total_in_channels, + out_channels=out_channels, + pred_residuals=self.pred_residuals, + last_kernel_size=self.last_kernel_size, + pad=self.pad, + unet=self.unet.build( + in_channels=total_in_channels, + pad=self.pad, + checkpointing=self.checkpointing, + ), + pos_channels=self.pos_channels, + add_3d_coordinates=add_3d_coordinates, + hist=hist, + grid_size=src.grid_size, + gradient_detach_interval=self.gradient_detach_interval, + use_bfloat16=self.use_bfloat16, + ) + + +class SamudraMultiConfig(BaseModelConfig): + encoder: EncoderConfig = EncoderConfig() + processor: UNetBackboneConfig = UNetBackboneConfig() + decoder: DecoderConfig = DecoderConfig() + perceiver_implementation: PerceiverImpl = Field( + default="auto", + description="Perceiver attention implementation shared by the encoder and decoder. " + "'auto' selects flash attention when CUDA is available, otherwise naive.", + ) + patch_extent: list[float] = Field( + default=[6.0, 10.0], + description="Target physical extent of each patch in degrees [height_deg, width_deg]. " + "Shared by the encoder and decoder for consistent spatial semantics.", + ) + embedding_dim: int = 128 + use_bfloat16: bool = Field( + default=True, + description="Use bfloat16 for most layers rather than float32. Required for flash attention.", + ) + + def build( + self, + prog_channels: int, + boundary_channels: int, + out_channels: int, + hist: int, + srcs: list[DataSource], + ) -> SamudraMulti: + assert len(self.patch_extent) == 2, "patch_extent must be a pair of floats." + extent = self.patch_extent[0], self.patch_extent[1] + + all_grid_sizes = [s.grid_size for s in srcs] + max_lat_size, max_lon_size = ( + max(g[0] for g in all_grid_sizes), + max(g[1] for g in all_grid_sizes), + ) + + impl = self.perceiver_implementation + if _use_flash(impl) and not self.use_bfloat16: + raise ValueError( + "Perceiver implementation resolves to flash attention. " + "Please set `use_bfloat16=True` or `perceiver_implementation='naive'`." + ) + + in_channels = prog_channels + boundary_channels + total_in_channels = in_channels + (3 if self.add_3d_coordinates else 0) + + encoder = self.encoder.build( + total_in_channels, + self.embedding_dim, + extent, + max_lat_size, + max_lon_size, + impl, + ) + processor = self.processor.build( + self.embedding_dim, + self.pad, + self.checkpointing, + ) + decoder = self.decoder.build( + processor.out_channels, + out_channels, + extent, + impl, + ) + + add_3d_coordinates = Concat3dCoordinates() if self.add_3d_coordinates else None + return SamudraMulti( + in_channels=total_in_channels, + out_channels=out_channels, + pred_residuals=self.pred_residuals, + last_kernel_size=self.last_kernel_size, + pad=self.pad, + encoder=encoder, + processor=processor, + decoder=decoder, + add_3d_coordinates=add_3d_coordinates, + hist=hist, + checkpointing=self.checkpointing, + gradient_detach_interval=self.gradient_detach_interval, + use_bfloat16=self.use_bfloat16, + ) + + +class SamudraMiniConfig(BaseModelConfig): + perceiver: PerceiverConfig = PerceiverConfig() + perceiver_implementation: PerceiverImpl = Field( + default="auto", + description="Perceiver attention implementation for the single PerceiverIO model. " + "'auto' selects flash attention when CUDA is available, otherwise naive.", + ) + embedding_dim: int = Field( + default=128, + description="Dimension of data-token embeddings before PerceiverIO.", + ) + queries_dim: int = Field( + default=128, + description="Dimension of PerceiverIO output queries.", + ) + coordinate_embedding_dim: int = Field( + default=64, + description="Hidden dimension used by learned 3D Cartesian coordinate embeddings.", + ) + query_chunk_size: int | None = Field( + default=None, + description="Optional chunk size for query decoding. If set, PerceiverIO is called " + "over query chunks to reduce memory use.", + ) + use_bfloat16: bool = Field( + default=True, + description="Use bfloat16 for most layers rather than float32. Required for flash attention.", + ) + + def build( + self, + prog_channels: int, + boundary_channels: int, + out_channels: int, + hist: int, + srcs: list[DataSource], + ) -> SamudraMini: + if self.add_3d_coordinates: + raise ValueError( + "SamudraMini always uses learned Cartesian coordinate embeddings. " + "Please set `add_3d_coordinates=False`." + ) + + impl = self.perceiver_implementation + if _use_flash(impl) and not self.use_bfloat16: + raise ValueError( + "Perceiver implementation resolves to flash attention. " + "Please set `use_bfloat16=True` or `perceiver_implementation='naive'`." + ) + + in_channels = prog_channels + boundary_channels + perceiver_io = self.perceiver.build_io( + self.embedding_dim, + self.queries_dim, + out_channels, + impl, + ) + return SamudraMini( + in_channels=in_channels, + out_channels=out_channels, + pred_residuals=self.pred_residuals, + last_kernel_size=self.last_kernel_size, + pad=self.pad, + input_embedding_dim=self.embedding_dim, + coordinate_embedding_dim=self.coordinate_embedding_dim, + queries_dim=self.queries_dim, + query_chunk_size=self.query_chunk_size, + perceiver_io=perceiver_io, + hist=hist, + checkpointing=self.checkpointing, + gradient_detach_interval=self.gradient_detach_interval, + use_bfloat16=self.use_bfloat16, + ) + + +AnyModelConfig = SamudraConfig | SamudraMultiConfig | SamudraMiniConfig + + +class DistributedConfig(BaseConfig): + dist_url: str | None = None + world_size: int | None = None + rank: int | None = None + gpu: int | None = None + dist_backend: str | None = None + + +class ExperimentConfig(BaseConfig): + name: str = "cm4_samudra" + rand_seed: int = 1 + base_output_dir: str = "train" + # we require this to be set by the user but have optional here + # so we can leave it out of config files + data_root: Location | None = None + wandb: WandBConfig + + @cached_property + def output_dir(self) -> Path: + return Path(self.base_output_dir) / f"{self.name}" + + @cached_property + def nets_dir(self) -> Path: + return self.output_dir / "saved_nets" + + @cached_property + def resolved_data_root(self) -> ResolvedLocation: + # Default to the current directory when no data_root is given. Absolute + # data locations (e.g. an s3:// demo source) ignore the root entirely, so + # they run flagless; relative locations resolve against cwd and still + # fail loudly at open() if the data isn't there. + default_root = LocalLocation(path=Path.cwd()) + if self.data_root is None: + return default_root + return default_root.resolve(self.data_root) + + +class ProfilerConfig(BaseConfig): + # How often (in batches processed) to take a snapshot of the CUDA memory + # (None = no snapshots) + cuda_snapshot_frequency: int | None = None + + def build(self, output_dir: Path, device: torch.device) -> Profiler: + if self.cuda_snapshot_frequency is not None and device.type != "cuda": + raise ValueError( + "cuda_snapshot_frequency is only supported on CUDA devices, got " + f"{device.type}" + ) + return Profiler(output_dir, self.cuda_snapshot_frequency) + + +# See backend.py for how these are turned into concrete devices +TrainBackendConfig = Literal["cpu", "cuda", "nccl", "auto"] + + +class DynamicLossConfig(pydantic.BaseModel): + type: Literal["dynamic"] = "dynamic" + metric: LossMetric = "mse" + limit: float | None = Field( + description="The ratio of the largest weight to the smallest weight across all channels which we'll allow. Default of None means no limit.", + default=None, + ge=1.0, + ) + + +class GradientLossConfig(pydantic.BaseModel): + type: Literal["gradient"] = "gradient" + # at the moment this metric is only used for the non-gradient loss + # (and would take a bit of refactoring to make it work for the gradient loss too) + # so we fix it to MAE for now until it's clear we what flexibility is needed here. + # TODO(#497): support other metrics for the gradient loss + metric: Literal["mae"] = "mae" + alpha: float = Field( + description="Scaling factor for the gradient penalty term (alpha in the gradient-weighted loss).", + default=0.1, + ge=0.0, + ) + + +Loss = LossMetric | DynamicLossConfig | GradientLossConfig + + +def build_loss_fn( + loss_cfg: Loss, + device: torch.device, + num_channels: int, + pad_mode: str, +) -> LossFnWithContext: + match loss_cfg: + case str(): + return loss_fn_from_metric(loss_cfg) + case DynamicLossConfig(metric=metric, limit=limit): + loss_fn = loss_fn_from_metric(metric) + return DynamicLoss( + loss_fn=loss_fn, + limit=limit, + device=device, + num_channels=num_channels, + ) + case GradientLossConfig(metric=metric, alpha=alpha): + loss_fn = loss_fn_from_metric(metric) + return GradientLoss( + loss_fn=loss_fn, + gradient_weight=alpha, + pad_mode=pad_mode, + ) + case _: + assert_never(loss_cfg) + + +class TrainConfig(TopLevelConfig): + # Training parameters + disk_mode: bool = True + pin_mem: bool = True + save_freq: int = 5 + validation_image_log_freq: int = Field( + default=10, + ge=1, + description=( + "How often to log expensive validation images. Epochs are 1-based, so " + "a value of 10 logs on epochs 1, 11, 21, ..." + ), + ) + epochs: int = 120 + preemptible: bool = True + batch_size: int = 2 + learning_rate: float = 2e-4 + gradient_accumulation_steps: int = 1 + scheduler: SchedulerConfig | None = None + loss: Loss = "mse" + finetune: bool = False + resume_ckpt_path: str | None = None + debug: bool = False + test_using_ema: bool = True + ema_decay: float = 0.999 + faster_decay_at_start: bool = True + delayed_loss_estimate: bool = False + backend: TrainBackendConfig = "auto" + + # Profiling parameters + profiler: ProfilerConfig = ProfilerConfig() + + # Data parameters at root level + data_stride: list[int] = [1] + steps: list[int] = [4] + step_transition: list[int] = [] + inference_epochs: list[int] = [-1] + + # Config components + experiment: ExperimentConfig + data: DataConfig + model: AnyModelConfig + + def prepare_output_dirs(self) -> None: + self.experiment.nets_dir.mkdir(parents=True, exist_ok=True) + self.experiment.output_dir.mkdir(parents=True, exist_ok=True) + + +class ObsMetricsConfig(BaseConfig): + """Where the observation products live, and over what period to score. + + Supplying this block is what turns observation metrics on: an eval job then + compares its rollout against DUACS, OISST and ARGO-IAP once the rollout is + on disk. Omit it (or pass `--observations=null`) to skip that phase. + """ + + duacs_location: Location = Field( + default=UnresolvedLocation(path="obs/duacs.zarr"), + description="DUACS surface geostrophic velocity, on its native grid.", + ) + oisst_location: Location = Field( + default=UnresolvedLocation(path="obs/oisst.zarr"), + description="OISST sea-surface temperature, on its native grid.", + ) + argo_iap_location: Location = Field( + default=UnresolvedLocation(path="obs/argo-iap.zarr"), + description="ARGO-IAP gridded temperature, on its native grid.", + ) + rmse_start: str = Field( + default="2015-01-01", + description=( + "Start of the primary scoring window. Must begin a complete calendar " + "year: the score and its bootstrap both use equal-year blocks." + ), + ) + rmse_end: str = Field( + default="2022-12-31", + description="End of the primary scoring window; must end a complete calendar year.", + ) + bootstrap_samples: int = Field( + default=10_000, + ge=0, + description="Calendar-year block-bootstrap draws for the 95% CI; 0 disables.", + ) + velocity_kind: Literal["absolute", "anomaly"] = Field( + default="absolute", + description=( + "Which DUACS geostrophic velocity to compare against. 'absolute' " + "matches the model velocity derived from zos." + ), + ) + baselines: list[str] = Field( + default_factory=lambda: ["om4"], + description=( + "Reference rollouts scored alongside the model. 'om4' scores the " + "ground-truth OM4 data already staged for the eval, which is what " + "makes the model's own numbers interpretable." + ), + ) + + @model_validator(mode="after") + def _check(self) -> "ObsMetricsConfig": + unknown = set(self.baselines) - {"om4"} + if unknown: + raise ValueError(f"Unknown baselines {sorted(unknown)}; supported: ['om4']") + if self.window[0] > self.window[1]: + raise ValueError( + f"rmse_start ({self.rmse_start}) must not be after rmse_end " + f"({self.rmse_end})" + ) + return self + + @property + def window(self) -> tuple[pd.Timestamp, pd.Timestamp]: + """The scoring window, with a date-only end bound taken to end-of-day. + + OM4 samples are stamped at 12:00, so a bare `2022-12-31` end bound would + exclude that day's sample. Currently masked because the observation + products end earlier and the window gets trimmed to the shared span, but + it becomes a silent one-sample loss the moment coverage extends past + `rmse_end` -- and the 7.5-day tolerance in the complete-calendar-year + check is far too loose to notice. + """ + end = pd.Timestamp(self.rmse_end) + if end == end.normalize(): + end = end + pd.Timedelta(days=1) - pd.Timedelta(1, "ns") + return pd.Timestamp(self.rmse_start), end + + +# See backend.py for how these are turned into concrete devices +EvalBackendConfig = Literal["cpu", "cuda", "auto"] + + +class EvalConfig(TopLevelConfig): + # Basic parameters + debug: bool = False + save_zarr: bool = False + disk_mode: bool = True + # we require this to be set by the user but have optional here + # so we can leave it out of config files + ckpt_path: str | None = None + num_model_steps_forward: int = 200 + backend: EvalBackendConfig = "auto" + + # Config components + experiment: ExperimentConfig + data: DataConfig + model: AnyModelConfig + observations: ObsMetricsConfig | None = Field( + default=None, + description=( + "Observation-based metrics computed after the rollout finishes. " + "Present means enabled; omit to skip the phase." + ), + ) + + @pydantic.model_validator(mode="after") + def _observations_need_a_saved_rollout(self) -> Self: + # Caught here rather than after the rollout: observation metrics score a + # rollout read back from disk, and discovering the misconfiguration at + # the end would waste the whole job. + if self.observations is not None and not self.save_zarr: + raise ValueError( + "observations requires save_zarr=true: the metrics are computed " + "from the predictions.zarr the rollout writes." + ) + return self + + def prepare_output_dirs(self) -> None: + self.experiment.output_dir.mkdir(parents=True, exist_ok=True) + + +AnyTopLevelConfig = TrainConfig | EvalConfig diff --git a/samudra/config_base.py b/samudra/config_base.py new file mode 100644 index 0000000000000000000000000000000000000000..7d7e36edd2fdc842f2e2f97a694f7466b4d4e439 --- /dev/null +++ b/samudra/config_base.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import importlib.resources +import os +import textwrap +from pathlib import Path +from typing import Any, Self + +import yaml +from pydantic import BaseModel, ConfigDict +from pydantic.fields import FieldInfo +from pydantic_settings import ( + BaseSettings, + CliSettingsSource, + PydanticBaseSettingsSource, + SettingsConfigDict, + YamlConfigSettingsSource, +) + + +def register_include_constructor(): + """Set up yaml.safe_load to include other yaml files via !include.""" + + def include_constructor(loader: yaml.Loader, node: yaml.Node) -> Any: + if hasattr(loader.stream, "name"): + name = loader.stream.name # type: ignore + else: + raise ValueError( + "To support includes, you must load a file object, not a string" + ) + filename = os.path.normpath(os.path.join(os.path.dirname(name), node.value)) + with open(filename) as f: + return yaml.safe_load(f) + + # This is arguably unsafe, but we don't parse untrusted YAML + yaml.loader.SafeLoader.add_constructor("!include", include_constructor) + + +register_include_constructor() + + +def resolve_config_path(config: str) -> Path: + """Resolve a config argument to a file on disk. + + A real filesystem path wins, so a checkout behaves exactly as before. If it + doesn't exist, fall back to a preset bundled in the installed package (e.g. + ``samudra_om4/train.yaml``), so ``samudra train samudra_om4/train.yaml`` + works with no checkout. The bundled configs install as real co-located + files, so their relative ``!include``s resolve unchanged. + """ + path = Path(config).expanduser() + if not path.exists(): + bundled = importlib.resources.files("samudra") / "configs" / config + if bundled.is_file(): + path = Path(str(bundled)) + return path.resolve() + + +class BaseConfig(BaseModel): + """ + Base class for all configs. + """ + + model_config = ConfigDict(extra="forbid") + + +class TopLevelConfig(BaseSettings): + """ + Base class for top-level configs (ie tasks like train or eval). + """ + + model_config = SettingsConfigDict(nested_model_default_partial_update=True) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @classmethod + def settings_customise_sources( + cls, + settings_cls: type[BaseSettings], + init_settings: PydanticBaseSettingsSource, + env_settings: PydanticBaseSettingsSource, + dotenv_settings: PydanticBaseSettingsSource, + file_secret_settings: PydanticBaseSettingsSource, + ) -> tuple[PydanticBaseSettingsSource, ...]: + # We don't need env/dotenv/secrets, yaml & CLI are injected in from_yaml_and_cli + return (init_settings,) + + @classmethod + def from_yaml_and_cli( + cls, + args_to_parse: list[str] | None = None, + ) -> Self: + """Load config from YAML & CLI with validation.""" + parser = argparse.ArgumentParser( + description=cls.__doc__, + epilog=textwrap.dedent( + """ + CONFIG may be a path to a YAML file or, when samudra is installed, + the name of a bundled preset such as `samudra_om4/train.yaml`. + YAML files can include other YAML files using the !include tag, + as in `data: !include ../data/om4.yaml` (resolved relative to the + including file). You can also replace any JSON argument listed above + with a YAML file by specifying it with an @ symbol, + eg `--some_param=@../data/om4.yaml`. + """ + ), + ) + parser.add_argument( + "config", + type=str, + help="Path to a config YAML file, or a bundled preset name", + ) + + cli_source = IncludeYamlCliSettingsSource( + cls, + root_parser=parser, + # If args_to_parse is None, we parse argv, which is what `True` does + cli_parse_args=args_to_parse if args_to_parse is not None else True, + ) + + # We do this after creating CliSettingsSource (which populates the parser) + # so the help is complete on error. + args = parser.parse_args(args_to_parse) + + # Then we read the YAML file specified in the CLI. A bare path resolves + # against the filesystem; if it's missing we try a preset bundled in the + # installed package before giving up. + # Note that by default, YamlConfigSettingsSource will ignore missing files + config_path = resolve_config_path(args.config) + if not config_path.exists(): + raise FileNotFoundError( + f"Config file `{args.config}` (full path: {config_path}) not found" + ) + yaml_values = YamlConfigSettingsSource(cls, yaml_file=config_path)() + + return cls( + _cli_settings_source=cli_source, + **yaml_values, + ) + + def save_yaml(self, save_path: Path) -> None: + """Save config to YAML file.""" + with open(save_path, "w") as f: + yaml.dump(self.model_dump(), f) + + +class IncludeYamlCliSettingsSource(CliSettingsSource): + """CliSettingsSource which permits @filename.yaml for JSON arguments.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def decode_complex_value( + self, field_name: str, field: FieldInfo, value: Any + ) -> Any: + if isinstance(value, str) and value.startswith("@"): + # Resolve like the positional config: a real path wins, otherwise + # fall back to a bundled preset (so `--data @data/om4_demo.yaml` + # works without a checkout). + with open(resolve_config_path(value[1:])) as f: + return yaml.safe_load(f) + return super().decode_complex_value(field_name, field, value) diff --git a/samudra/config_schema.py b/samudra/config_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..692fbf531cdc581c46b944d3b7fd2052383efdda --- /dev/null +++ b/samudra/config_schema.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import json +import re +import sys +from pathlib import Path +from typing import get_args, get_origin + +import pydantic +import yaml + +from samudra.config import EvalConfig, TrainConfig +from samudra.viz import VizConfig + + +def get_pydantic_models( + model: type[pydantic.BaseModel], + models: dict[str, type[pydantic.BaseModel]] | None = None, + seen: set[int] | None = None, +) -> dict[str, type[pydantic.BaseModel]]: + """Recursively find all Pydantic models in a model's fields. + + Args: + model: The Pydantic model to start from + models: Collected models keyed by model name + seen: Visited annotation objects keyed by identity + + Returns: + Dictionary of model names to model classes + """ + if models is None: + models = {} + if seen is None: + seen = set() + + _collect_pydantic_models(model, models, seen) + return models + + +def _collect_pydantic_models( + annotation: object, + models: dict[str, type[pydantic.BaseModel]], + seen: set[int], +) -> None: + annotation_id = id(annotation) + if annotation_id in seen: + return + seen.add(annotation_id) + + if isinstance(annotation, type) and issubclass(annotation, pydantic.BaseModel): + models[annotation.__name__] = annotation + for field in annotation.model_fields.values(): + _collect_pydantic_models(field.annotation, models, seen) + return + + # We check if this is a parameterized type and recurse if so (eg Union[int, str]) + # (Union is the origin in that case.) + origin = get_origin(annotation) + if origin is None: + return + + for arg in get_args(annotation): + _collect_pydantic_models(arg, models, seen) + + +def generate_schemas( + output_dir: Path, + models: dict[str, type[pydantic.BaseModel]], +) -> None: + """Generate JSON schemas for all Pydantic models and save them to output_dir. + + Args: + output_dir: Directory to save the generated schemas + models: Dictionary of model names to model classes + """ + output_dir.mkdir(parents=True, exist_ok=True) + + for model_name, model in models.items(): + schema = model.model_json_schema() + output_path = output_dir / f"{model_name}.json" + + # Check if file exists and content is different + if output_path.exists(): + with open(output_path) as f: + existing_schema = json.load(f) + if existing_schema == schema: + continue + + # Write only if content differs or file doesn't exist + with open(output_path, "w") as f: + json.dump(schema, f, indent=2) + print(f"🆕 Updated schema for {model_name} at {output_path}") + + +def validate_config_files( + config_dir: Path, + models: dict[str, type[pydantic.BaseModel]], +) -> bool: + """Validate YAML configuration files against their Pydantic models. + + Args: + config_dir: Directory containing YAML configuration files to validate + models: Dictionary of model names to model classes + """ + valid = True + + # Find all YAML files in the config directory + yaml_files = list(config_dir.rglob("*.yaml")) + + for yaml_file in yaml_files: + # Read the YAML file + with open(yaml_file) as f: + yaml_content = f.read() + + # Extract the schema path from the YAML file's first line + schema_match = re.search( + r"\s*#\s+yaml-language-server:\s+\$schema=(.*\.json)", yaml_content + ) + if not schema_match: + print(f"⚠️ No schema specified in {yaml_file}") + continue + + schema_name = Path(schema_match.group(1)).stem + if schema_name not in models: + print(f"✗ Model {schema_name} not found for {yaml_file}") + valid = False + continue + + # Load and validate the YAML content + try: + with open(yaml_file) as f: + config = yaml.safe_load(f) + # Validate using Pydantic model + models[schema_name].model_validate(config) + print(f"✅ {yaml_file} is a valid {schema_name}") + except Exception as e: + print(f"❌ {yaml_file} is an invalid {schema_name}: {str(e)}") + valid = False + + return valid + + +def main(): + parser = argparse.ArgumentParser( + description="Generate JSON schemas and validate config files" + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("src/samudra/configs/schemas"), + help="Directory to save the generated schemas", + ) + parser.add_argument( + "--validate-dir", + type=Path, + default=Path("src/samudra/configs"), + help="Directory of config files to validate", + ) + + args = parser.parse_args() + + # Get all available models + models = get_pydantic_models(TrainConfig) + models.update(get_pydantic_models(EvalConfig)) + models.update(get_pydantic_models(VizConfig)) + + generate_schemas(args.output_dir, models) + + if not validate_config_files(args.validate_dir, models): + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/samudra/configs/data/llc.yaml b/samudra/configs/data/llc.yaml new file mode 100644 index 0000000000000000000000000000000000000000..62bb5ad0d91ddb7b916bda0fbdd6d5cef1c0e151 --- /dev/null +++ b/samudra/configs/data/llc.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: 2026 Ocean Emulator Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/DataConfig.json + +sources: + - type: llc + face: 1 + i_start: 0 + i_end: 720 + j_start: 0 + j_end: 720 + train_time: + start: "2011-09-10T12:00:00Z" + end: "2012-09-01T12:00:00Z" + val_time: + start: "2012-09-01T12:00:00Z" + end: "2012-11-15T12:00:00Z" + data_location: LLC.zarr + data_means_location: LLC_means.zarr + data_stds_location: LLC_stds.zarr +concurrent_compute: false +loading: + type: cpu + num_workers: 4 diff --git a/samudra/configs/data/obs.yaml b/samudra/configs/data/obs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2e9c782b50d35f881dad50d3e271c51bb4e04f5c --- /dev/null +++ b/samudra/configs/data/obs.yaml @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# Observation-based evaluation metrics, scored once the rollout is on disk. +# Including this file is what enables them; omit it to skip the phase. +# +# The three products are published anonymously readable, so an eval run needs +# neither credentials nor a staging step to produce these metrics. They were +# built from the providers by `python -m ocean_preprocessing.obs_preprocessing` +# (see data/README.md), and each store records its own source, attribution and +# provenance in its zarr attrs. +# +# To read a locally staged copy instead -- worth it on repeat runs, since DUACS +# alone is 46 GB -- set OSN_OBS_PATH in the skypilot eval job and override the +# locations to the relative paths `obs/duacs.zarr` and so on, which resolve +# against experiment.data_root. + +duacs_location: &obs_s3 + type: s3 + bucket: m2lines-pubs + path: Samudra/v2026-07/obs/duacs.zarr + endpoint_url: https://nyu1.osn.mghpcc.org + anon: true +oisst_location: + <<: *obs_s3 + path: Samudra/v2026-07/obs/oisst.zarr +argo_iap_location: + <<: *obs_s3 + path: Samudra/v2026-07/obs/argo-iap.zarr + +# Eight complete calendar years, so the score and its block bootstrap share +# equal-weight blocks. +rmse_start: "2015-01-01" +rmse_end: "2022-12-31" diff --git a/samudra/configs/data/om4.yaml b/samudra/configs/data/om4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..70e2d8040a4b3493119c506bf1a1a0952ba41e81 --- /dev/null +++ b/samudra/configs/data/om4.yaml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/DataConfig.json + +sources: + - type: om4 + prognostic_vars_key: thermo_dynamic_all + boundary_vars_key: tau_hfds_hfds_anom + train_time: + start: "1975-01-03" + end: "2013-10-05" + val_time: + start: "2013-10-05" + end: "2014-10-05" + inference_times: + - start: "2014-10-10" + end: "2022-12-24" + data_location: OM4.zarr + data_means_location: OM4_means.zarr + data_stds_location: OM4_stds.zarr +concurrent_compute: true +loading: + type: cpu diff --git a/samudra/configs/data/om4_demo.yaml b/samudra/configs/data/om4_demo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5864cfdc95788ca6ff5afee186f150cc8f6d2d97 --- /dev/null +++ b/samudra/configs/data/om4_demo.yaml @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/DataConfig.json + +# Two-degree (2°) OM4 demo data source that streams straight from the public +# m2lines-pubs bucket on OSN over S3 — no local download or credentials needed +# (anon: true reads the public bucket anonymously). The locations below are +# absolute (type: s3), so they are read as-is and `experiment.data_root` is +# unused — you can omit it entirely. + +sources: + - type: om4 + prognostic_vars_key: thermo_dynamic_all + # tau_hfds (tauuo, tauvo, hfds) rather than the samudra_om4 presets' + # tau_hfds_hfds_anom: the public 2deg data has no `hfds_anomalies` variable. + # This is a self-contained train -> eval -> viz demo (train a fresh model on + # this data); its boundary channels differ from the released checkpoints, so + # don't eval a preset-trained checkpoint against this source. + boundary_vars_key: tau_hfds + train_time: + start: "1975-01-03" + end: "2013-10-05" + val_time: + start: "2013-10-05" + end: "2014-10-05" + inference_times: + - start: "2014-10-10" + end: "2022-12-24" + data_location: + type: s3 + endpoint_url: "https://nyu1.osn.mghpcc.org" + anon: true + bucket: m2lines-pubs + path: Samudra/v2026-07/om4_twodeg/OM4.zarr + data_means_location: + type: s3 + endpoint_url: "https://nyu1.osn.mghpcc.org" + anon: true + bucket: m2lines-pubs + path: Samudra/v2026-07/om4_twodeg/OM4_means.zarr + data_stds_location: + type: s3 + endpoint_url: "https://nyu1.osn.mghpcc.org" + anon: true + bucket: m2lines-pubs + path: Samudra/v2026-07/om4_twodeg/OM4_stds.zarr +concurrent_compute: true +loading: + type: cpu diff --git a/samudra/configs/samudra_mini_om4/README.md b/samudra/configs/samudra_mini_om4/README.md new file mode 100644 index 0000000000000000000000000000000000000000..afe9a085f85b5438e385e9c95cbdb774be509dfd --- /dev/null +++ b/samudra/configs/samudra_mini_om4/README.md @@ -0,0 +1,11 @@ + + +# SamudraMini OM4 + +These configurations are for SamudraMini, a single PerceiverIO model that treats +each `(lat, lon)` point as a token and uses learned embeddings of 3D Cartesian +coordinates for both data tokens and output queries. diff --git a/samudra/configs/samudra_mini_om4/model.yaml b/samudra/configs/samudra_mini_om4/model.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1dd37da2a78505cb11bf71f76c1d3e1725b1d31f --- /dev/null +++ b/samudra/configs/samudra_mini_om4/model.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/SamudraMiniConfig.json + +pred_residuals: false +last_kernel_size: 3 +pad: "circular" +checkpointing: "all" +perceiver_implementation: "auto" +embedding_dim: 128 +queries_dim: 128 +coordinate_embedding_dim: 64 +query_chunk_size: 8192 + +perceiver: + depth: 4 + latent_dim: 96 + num_latents: 256 diff --git a/samudra/configs/samudra_mini_om4/train.yaml b/samudra/configs/samudra_mini_om4/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c2543f7e0c4a6bceaff9053e56da8c0d1d2f292c --- /dev/null +++ b/samudra/configs/samudra_mini_om4/train.yaml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/TrainConfig.json + +debug: false +disk_mode: true +pin_mem: true +save_freq: 5 +epochs: 70 +batch_size: 2 +learning_rate: 0.0006 +scheduler: { type: cosine } +loss: + type: dynamic + metric: mse +finetune: false +resume_ckpt_path: null +inference_epochs: [] +data_stride: [1] +steps: [4] +step_transition: [] +preemptible: false + +backend: auto + +experiment: + name: samudra_mini_om4 + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_multi_om4/README.md b/samudra/configs/samudra_multi_om4/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ebbf6541b4d06228398bc4544c745289796e9e2a --- /dev/null +++ b/samudra/configs/samudra_multi_om4/README.md @@ -0,0 +1,11 @@ + + +# samudra-multi OM4 + +These configurations are for the samudra-multi (encoder-processor-decoder) model, +trained on OM4 data. The intent is for this to support heterogeneous +resolutions and modalities. diff --git a/samudra/configs/samudra_multi_om4/eval.yaml b/samudra/configs/samudra_multi_om4/eval.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3f9d8cfb260fa713d0a0bd61c3ef2c7a1f520097 --- /dev/null +++ b/samudra/configs/samudra_multi_om4/eval.yaml @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/EvalConfig.json +# +# Eval config for the multi-scale samudra-multi model. +# +# The eval code runs inference on the *first* data source only. +# To evaluate all three scales, run three separate eval jobs and +# reorder the sources list so the target resolution is first. +# All three sources must always be present so the model builds +# with the correct max grid size for the encoder patch grid. +# The max grid size affects the rotary positional embedding +# dimension for the encoder's Perceiver. If it is set wrong, +# it would likely cause hard-to-debug model level errors. +# Setting multiple data sources lets us use checkpoints for +# samudra-multi models that use multiple scales. + +debug: false +save_zarr: true +disk_mode: true +num_model_steps_forward: 25 + +experiment: + name: samudra_multi_om4_eval + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default + +data: + concurrent_compute: true + sources: + # First source is used for inference. Reorder to eval a different scale. + - type: om4 + prognostic_vars_key: thermo_dynamic_all + boundary_vars_key: tau_hfds + train_time: &train_time + start: "1975-01-03" + end: "2013-10-05" + val_time: &val_time + start: "2013-10-05" + end: "2014-10-05" + inference_times: &inference_times + - start: "2014-10-05" + end: "2015-10-05" + data_location: om4_halfdeg_v4/OM4.zarr + data_means_location: om4_halfdeg_v4/OM4_means.zarr + data_stds_location: om4_halfdeg_v4/OM4_stds.zarr + - type: om4 + prognostic_vars_key: thermo_dynamic_all + boundary_vars_key: tau_hfds + train_time: *train_time + val_time: *val_time + inference_times: *inference_times + data_location: om4_onedeg_v3/OM4.zarr + data_means_location: om4_onedeg_v3/OM4_means.zarr + data_stds_location: om4_onedeg_v3/OM4_stds.zarr + - type: om4 + prognostic_vars_key: thermo_dynamic_all + boundary_vars_key: tau_hfds + train_time: *train_time + val_time: *val_time + inference_times: *inference_times + data_location: om4_quarterdeg_v2/OM4.zarr + data_means_location: om4_quarterdeg_v2/OM4_means.zarr + data_stds_location: om4_quarterdeg_v2/OM4_stds.zarr + +model: !include model.yaml diff --git a/samudra/configs/samudra_multi_om4/model.yaml b/samudra/configs/samudra_multi_om4/model.yaml new file mode 100644 index 0000000000000000000000000000000000000000..612999d2c0058fcf477cf953dcc9338ce1080016 --- /dev/null +++ b/samudra/configs/samudra_multi_om4/model.yaml @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/SamudraMultiConfig.json + +pred_residuals: false +last_kernel_size: 3 +pad: "circular" +checkpointing: "all" +embedding_dim: 128 # Reduced from 280 +perceiver_implementation: "auto" +# TODO(alxmrs): How fine of a resolution should each patch be? This produces a 60x72 latent grid. +patch_extent: [3.0, 5.0] + +encoder: + perceiver: + depth: 2 + latent_dim: 64 + num_latents: 256 + +decoder: + perceiver: + depth: 2 + latent_dim: 64 + num_latents: 256 + queries_dim: 64 # on the small end, but typical for PercieverIO models (usually in the ~100s-1000s). + window_patches: 12 # Must evenly divide latent grid (60x72). Enables real spatial windowing. + context_patches: 1 # Look at this many window_patches around the target patch. + +processor: + ch_width: [380, 480, 520] + dilation: [1, 1, 1] + n_layers: [1, 1, 1] + + core_block: + block_type: "conv_next_block" + kernel_size: 3 + activation: "capped_gelu" + upscale_factor: 2 + norm: "instance" + pointwise_linear: true + + down_sampling_block: "avg_pool" + up_sampling_block: "zonally_periodic_upsample" diff --git a/samudra/configs/samudra_multi_om4/train.yaml b/samudra/configs/samudra_multi_om4/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c7482a3ac4f5c2e7d3ed5fc8c1f3cc641d4d2cdd --- /dev/null +++ b/samudra/configs/samudra_multi_om4/train.yaml @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/TrainConfig.json + +debug: false +disk_mode: true +pin_mem: true +save_freq: 5 +epochs: 70 +batch_size: 2 +learning_rate: 0.0006 +scheduler: { type: cosine } +loss: + type: dynamic + metric: mse + limit: 20 +finetune: false +resume_ckpt_path: null +inference_epochs: [] +data_stride: [1] +steps: [4] +step_transition: [] +preemptible: false + +backend: auto + +experiment: + name: samudra_multi_om4 + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: + concurrent_compute: true + loading: + type: cpu + num_workers: 4 + sources: + - type: om4 + prognostic_vars_key: thermo_dynamic_all + boundary_vars_key: tau_hfds + train_time: &train_time + start: "1975-01-03" + end: "2013-10-05" + val_time: &val_time + start: "2013-10-05" + end: "2014-10-05" + inference_times: &inference_times + - start: "2014-10-05" + end: "2015-10-05" + data_location: om4_quarterdeg_v2/OM4.zarr + data_means_location: om4_quarterdeg_v2/OM4_means.zarr + data_stds_location: om4_quarterdeg_v2/OM4_stds.zarr + - type: om4 + prognostic_vars_key: thermo_dynamic_all + boundary_vars_key: tau_hfds + train_time: *train_time + val_time: *val_time + inference_times: *inference_times + data_location: om4_halfdeg_v4/OM4.zarr + data_means_location: om4_halfdeg_v4/OM4_means.zarr + data_stds_location: om4_halfdeg_v4/OM4_stds.zarr + - type: om4 + prognostic_vars_key: thermo_dynamic_all + boundary_vars_key: tau_hfds + train_time: *train_time + val_time: *val_time + inference_times: *inference_times + data_location: om4_onedeg_v3/OM4.zarr + data_means_location: om4_onedeg_v3/OM4_means.zarr + data_stds_location: om4_onedeg_v3/OM4_stds.zarr +model: !include model.yaml diff --git a/samudra/configs/samudra_om4/README.md b/samudra/configs/samudra_om4/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cb8c0cba344de453d0581376d3f6d204259a5fd1 --- /dev/null +++ b/samudra/configs/samudra_om4/README.md @@ -0,0 +1,10 @@ + + +# Samudra OM4 + +These configurations are our current WIP Samudra-style (UNet) model, +trained on OM4 data. diff --git a/samudra/configs/samudra_om4/eval.yaml b/samudra/configs/samudra_om4/eval.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ca0ad5938578069cb9925d172845e19b48d8b2e7 --- /dev/null +++ b/samudra/configs/samudra_om4/eval.yaml @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/EvalConfig.json + +debug: false +save_zarr: true +disk_mode: true +num_model_steps_forward: 25 + +experiment: + name: samudra_om4_eval + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml +observations: !include ../data/obs.yaml diff --git a/samudra/configs/samudra_om4/model.yaml b/samudra/configs/samudra_om4/model.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0bcb2c0f5c79e4e08585d195e8a0fadc265cedce --- /dev/null +++ b/samudra/configs/samudra_om4/model.yaml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/SamudraConfig.json + +checkpointing: all +pred_residuals: false +last_kernel_size: 3 +pad: "circular" + +unet: + ch_width: [280, 380, 480, 520] + dilation: [1, 2, 4, 8] + n_layers: [1, 1, 1, 1] + + core_block: + block_type: "conv_next_block" + kernel_size: 3 + activation: "capped_gelu" + upscale_factor: 2 + norm: "batch" + + down_sampling_block: "avg_pool" + up_sampling_block: "zonally_periodic_upsample" diff --git a/samudra/configs/samudra_om4/train.yaml b/samudra/configs/samudra_om4/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..683b3544bb44d6c416990edf55dd5334d75b6b82 --- /dev/null +++ b/samudra/configs/samudra_om4/train.yaml @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/TrainConfig.json + +debug: false +disk_mode: true +pin_mem: true +save_freq: 5 +epochs: 70 +batch_size: 4 +learning_rate: 0.0006 +gradient_accumulation_steps: 1 +scheduler: { type: cosine } +loss: + type: dynamic + metric: mse +finetune: false +resume_ckpt_path: null +inference_epochs: [] +data_stride: [1] +steps: [4] +step_transition: [] +preemptible: false + +backend: auto + +experiment: + name: samudra_om4 + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_om4/viz.yaml b/samudra/configs/samudra_om4/viz.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f2837474f74b7fb3c7171dfbe3d6b151d522febd --- /dev/null +++ b/samudra/configs/samudra_om4/viz.yaml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/VizConfig.json + +base_output_dir: ".LOCAL" +name: "samudra_om4_viz" +dataset_name: "OM4" +# Ground truth comes from this data source (same shape as train/eval), so +# `--data @data/om4_demo.yaml` points viz at the demo data just like the others. +data: !include ../data/om4.yaml +basins_location: + type: s3 + anon: true + bucket: "m2lines-pubs" + path: "Samudra/basins/basin_masks_original.zarr" + endpoint_url: "https://nyu1.osn.mghpcc.org" +groundtruth_time_range: + start: "2014-10-20" + end: "2022-12-24" +runs: + - name: "samudra_om4_eval" + location: ".LOCAL/samudra_om4_eval/predictions.zarr" +observations: !include ../data/obs.yaml diff --git a/samudra/configs/samudra_om4_v1/README.md b/samudra/configs/samudra_om4_v1/README.md new file mode 100644 index 0000000000000000000000000000000000000000..bb7af68f7a76954c14697a3b5cca048855aa8c0e --- /dev/null +++ b/samudra/configs/samudra_om4_v1/README.md @@ -0,0 +1,10 @@ + + +# Samudra OM4 V1 + +These configurations are our replication of the [original Samudra model](https://arxiv.org/abs/2412.03795) +trained on OM4 data. diff --git a/samudra/configs/samudra_om4_v1/eval.yaml b/samudra/configs/samudra_om4_v1/eval.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b158e2d467e80755c21e8edf89a57b4b785b813b --- /dev/null +++ b/samudra/configs/samudra_om4_v1/eval.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/EvalConfig.json + +debug: false +save_zarr: true +disk_mode: true +num_model_steps_forward: 25 + +experiment: + name: samudra_om4_v1_eval + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_om4_v1/model.yaml b/samudra/configs/samudra_om4_v1/model.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eeb4a6d54f7d5c4549256424e4770b2a722edca8 --- /dev/null +++ b/samudra/configs/samudra_om4_v1/model.yaml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/SamudraConfig.json + +pred_residuals: false +last_kernel_size: 3 +pad: "circular" +checkpointing: "all" + +unet: + ch_width: [200, 250, 300, 400] + dilation: [1, 2, 4, 8] + n_layers: [1, 1, 1, 1] + + core_block: + block_type: "conv_next_block" + kernel_size: 3 + activation: "capped_gelu" + upscale_factor: 4 + norm: "batch" + + down_sampling_block: "avg_pool" + up_sampling_block: "bilinear_upsample" diff --git a/samudra/configs/samudra_om4_v1/train.yaml b/samudra/configs/samudra_om4_v1/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..489ab3699b420f0722c5a312b104e322f7c7245d --- /dev/null +++ b/samudra/configs/samudra_om4_v1/train.yaml @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/TrainConfig.json + +debug: false +disk_mode: true +pin_mem: true +save_freq: 5 +epochs: 70 +batch_size: 4 +learning_rate: 0.0006 +scheduler: { type: cosine } +loss: mse +finetune: false +resume_ckpt_path: null +inference_epochs: [] +data_stride: [1] +steps: [4] +step_transition: [] +preemptible: false + +backend: auto + +experiment: + name: samudra_om4_v1 + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_om4_v1/viz.yaml b/samudra/configs/samudra_om4_v1/viz.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5d94a1a4a9efbcdd2ea28d5eb52ef891ce8aba92 --- /dev/null +++ b/samudra/configs/samudra_om4_v1/viz.yaml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/VizConfig.json + +base_output_dir: ".LOCAL" +name: "samudra_om4_v1_viz" +dataset_name: "OM4" +groundtruth_location: "OM4.zarr" +basins_location: + type: s3 + anon: true + bucket: "m2lines-pubs" + path: "Samudra/basins/basin_masks_original.zarr" + endpoint_url: "https://nyu1.osn.mghpcc.org" +groundtruth_time_range: + start: "2014-10-20" + end: "2022-12-24" +runs: + - name: "samudra_om4_v1_eval" + location: ".LOCAL/samudra_om4_v1_eval/predictions.zarr" diff --git a/samudra/configs/samudra_om4_v2/README.md b/samudra/configs/samudra_om4_v2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..42228a500fe99af6bbfe8f2956f1ac944eca0fb3 --- /dev/null +++ b/samudra/configs/samudra_om4_v2/README.md @@ -0,0 +1,11 @@ + + +# Samudra OM4 V2 + +These configurations are the v2 version of Samudra-style (UNet) model, +trained on OM4 data. Major changes from the v1 model are that it is +wider and uses a dynamically-weighted loss. diff --git a/samudra/configs/samudra_om4_v2/eval.yaml b/samudra/configs/samudra_om4_v2/eval.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9e0fd823dcaf026b3b7ed3614c2549fbdccab583 --- /dev/null +++ b/samudra/configs/samudra_om4_v2/eval.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/EvalConfig.json + +debug: false +save_zarr: true +disk_mode: true +num_model_steps_forward: 25 + +experiment: + name: samudra_om4_v2_eval + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_om4_v2/model.yaml b/samudra/configs/samudra_om4_v2/model.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0bcb2c0f5c79e4e08585d195e8a0fadc265cedce --- /dev/null +++ b/samudra/configs/samudra_om4_v2/model.yaml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/SamudraConfig.json + +checkpointing: all +pred_residuals: false +last_kernel_size: 3 +pad: "circular" + +unet: + ch_width: [280, 380, 480, 520] + dilation: [1, 2, 4, 8] + n_layers: [1, 1, 1, 1] + + core_block: + block_type: "conv_next_block" + kernel_size: 3 + activation: "capped_gelu" + upscale_factor: 2 + norm: "batch" + + down_sampling_block: "avg_pool" + up_sampling_block: "zonally_periodic_upsample" diff --git a/samudra/configs/samudra_om4_v2/train.yaml b/samudra/configs/samudra_om4_v2/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2d0e5dd6cf0853ac86cd4ce615cb318c45402ebd --- /dev/null +++ b/samudra/configs/samudra_om4_v2/train.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/TrainConfig.json + +debug: false +disk_mode: true +pin_mem: true +save_freq: 5 +epochs: 70 +batch_size: 4 +learning_rate: 0.0006 +gradient_accumulation_steps: 1 +scheduler: { type: cosine } +loss: + type: dynamic + metric: mse + limit: 20 +finetune: false +resume_ckpt_path: null +inference_epochs: [] +data_stride: [1] +steps: [4] +step_transition: [] +preemptible: false + +backend: auto + +experiment: + name: samudra_om4_v2 + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_om4_v2/viz.yaml b/samudra/configs/samudra_om4_v2/viz.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f06e80e162d2919a96b84b13bf1cdca3f2865456 --- /dev/null +++ b/samudra/configs/samudra_om4_v2/viz.yaml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/VizConfig.json + +base_output_dir: ".LOCAL" +name: "samudra_om4_v2_viz" +dataset_name: "OM4" +groundtruth_location: "OM4.zarr" +basins_location: + type: s3 + anon: true + bucket: "m2lines-pubs" + path: "Samudra/basins/basin_masks_original.zarr" + endpoint_url: "https://nyu1.osn.mghpcc.org" +groundtruth_time_range: + start: "2014-10-20" + end: "2022-12-24" +runs: + - name: "samudra_om4_v2_eval" + location: ".LOCAL/samudra_om4_v2_eval/predictions.zarr" diff --git a/samudra/configs/samudra_om4_v2_highres/README.md b/samudra/configs/samudra_om4_v2_highres/README.md new file mode 100644 index 0000000000000000000000000000000000000000..18e6b7846b15270be210ae77cc940730d8c0f119 --- /dev/null +++ b/samudra/configs/samudra_om4_v2_highres/README.md @@ -0,0 +1,11 @@ + + +# Samudra OM4 V2 High-Res Configration + +A variant of the V2 configuration which uses a reduced +batch size, instance norm, and bfloat16, for handling +higher spatial resolutions. diff --git a/samudra/configs/samudra_om4_v2_highres/eval.yaml b/samudra/configs/samudra_om4_v2_highres/eval.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c19936383e2806ae2a77282bc5f75623701b0c56 --- /dev/null +++ b/samudra/configs/samudra_om4_v2_highres/eval.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/EvalConfig.json + +debug: false +save_zarr: true +disk_mode: true +num_model_steps_forward: 25 + +experiment: + name: samudra_om4_v2_highres_eval + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_om4_v2_highres/model.yaml b/samudra/configs/samudra_om4_v2_highres/model.yaml new file mode 100644 index 0000000000000000000000000000000000000000..559cb1f855d3626faeb21834a63598fcf2b313bf --- /dev/null +++ b/samudra/configs/samudra_om4_v2_highres/model.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/SamudraConfig.json + +checkpointing: all +pred_residuals: false +last_kernel_size: 3 +pad: "circular" +use_bfloat16: true + +unet: + ch_width: [280, 380, 480, 520] + dilation: [1, 2, 4, 8] + n_layers: [1, 1, 1, 1] + + core_block: + block_type: "conv_next_block" + kernel_size: 3 + activation: "capped_gelu" + upscale_factor: 2 + norm: "instance" + + down_sampling_block: "avg_pool" + up_sampling_block: "zonally_periodic_upsample" diff --git a/samudra/configs/samudra_om4_v2_highres/train.yaml b/samudra/configs/samudra_om4_v2_highres/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..01788e3891ee57c19d0ae1059399d447ad913b00 --- /dev/null +++ b/samudra/configs/samudra_om4_v2_highres/train.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/TrainConfig.json + +debug: false +disk_mode: true +pin_mem: true +save_freq: 5 +epochs: 70 +batch_size: 1 +learning_rate: 0.0006 +gradient_accumulation_steps: 4 +scheduler: { type: cosine } +loss: + type: dynamic + metric: mse + limit: 20 +finetune: false +resume_ckpt_path: null +inference_epochs: [] +data_stride: [1] +steps: [4] +step_transition: [] +preemptible: false + +backend: auto + +experiment: + name: samudra_om4_v2_highres + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/om4.yaml +model: !include model.yaml diff --git a/samudra/configs/samudra_om4_v2_highres/viz.yaml b/samudra/configs/samudra_om4_v2_highres/viz.yaml new file mode 100644 index 0000000000000000000000000000000000000000..354ce830f5738b40366841d420a3a46d02797083 --- /dev/null +++ b/samudra/configs/samudra_om4_v2_highres/viz.yaml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/VizConfig.json + +base_output_dir: ".LOCAL" +name: "samudra_om4_v2_highres_viz" +dataset_name: "OM4" +groundtruth_location: "OM4.zarr" +basins_location: + type: s3 + anon: true + bucket: "m2lines-pubs" + path: "Samudra/basins/basin_masks_original.zarr" + endpoint_url: "https://nyu1.osn.mghpcc.org" +groundtruth_time_range: + start: "2014-10-20" + end: "2022-12-24" +runs: + - name: "samudra_om4_v2_highres_eval" + location: ".LOCAL/samudra_om4_v2_highres_eval/predictions.zarr" diff --git a/samudra/constants.py b/samudra/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..a3e33eab8ae135195c3a13ceb08dd2f19f8b99a9 --- /dev/null +++ b/samudra/constants.py @@ -0,0 +1,570 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import enum +import logging +from typing import Literal, Self + +logger = logging.getLogger(__name__) + +import torch +import xarray as xr +from jaxtyping import Bool, Float +from torch import Tensor + +# Common Type Aliases +# See "Existing jaxtyping annotations" section of +# https://docs.kidger.site/jaxtyping/api/array/#array + +Lat = Float[Tensor, "lat"] +Lon = Float[Tensor, "lon"] +Grid = Float[Tensor, "lat lon"] +GridSize = tuple[int, int] +Prognostic = Float[ + Grid, "*batch prognostic_vars" +] # equivalent to "*batch prognostic_vars lat lon" +Boundary = Float[Grid, "*batch boundary_vars"] +# A note from jaxtyping (why we can't do "prognostic_vars+boundary_vars"): +# In practice you should usually only use symbolic axes in annotations +# for return types, referring only to axes annotated for arguments. +# So, we'll leave this default and use symbolic axes locally. +type Input = Float[Grid, "*batch total_vars"] + +Example = tuple[ + Prognostic, Boundary, Prognostic +] # (prognostic_input, boundary_input, label) + +GridMask = Bool[Tensor, "lat lon"] +PrognosticMask = Bool[GridMask, "prognostic_vars"] + +SingleChannelVar = Float[Tensor, "batch time lat lon"] +DictSingleChannelVar = dict[str, SingleChannelVar] +SinglePrognosticTimeSeries = Float[Grid, "*batch time"] + +SingleTimeSeriesOutput = Float[Tensor, "batch=1 time prognostic_vars lat lon"] +BatchTimeSeriesOutput = Float[Tensor, "batch time=(hist+1) prognostic_vars lat lon"] +HistBatched = Float[Tensor, "batch_hist prognostic_vars lat lon"] +HistChanneled = Float[Tensor, "batch hist_prognostic_vars lat lon"] + + +MAX_TRAIN_MODEL_STEPS_FORWARD = 200 + +DatasetType = Literal["om4", "llc"] + +# Horizontal grid geometry. "gaussian" is a regular/rectilinear lat-lon grid whose +# 2D lat/lon are the outer product of the 1D axes, so they can be reconstructed by +# broadcasting. "tripolar" is curvilinear: lat/lon vary along both horizontal dims +# and cannot be rebuilt by broadcasting, so the real 2D coordinates must be carried +# with the data. Downstream code that reconstructs geometry must branch on this. +GridType = Literal["gaussian", "tripolar"] + +PrognosticVarNames = list[str] +BoundaryVarNames = list[str] + + +@dataclasses.dataclass(frozen=True) +class DatasetSpec: + type: DatasetType + depth_levels: tuple[float, ...] + depth_thickness: tuple[float, ...] + mask_vars: tuple[str, ...] + mask_all_levels_var: str + seconds_per_time_step: int + prognostic_var_names: PrognosticVarNames + boundary_var_names: BoundaryVarNames + default_metadata: dict[str, dict[str, str]] + ocean_heat_temperature_var: str + surface_heat_flux_var: str + grid_type: GridType = "gaussian" + + def __post_init__(self) -> None: + if len(self.depth_levels) != len(self.depth_thickness): + raise ValueError( + "depth_levels and depth_thickness must have the same length" + ) + if len(self.depth_levels) != len(self.mask_vars): + raise ValueError("depth_levels and mask_vars must have the same length") + + @property + def depth_i_levels(self) -> tuple[str, ...]: + return tuple(str(i) for i in range(len(self.depth_levels))) + + @property + def num_prognostic_depth_levels(self) -> int: + depth_indices = [] + for var_name in self.prognostic_var_names: + _, sep, suffix = var_name.rpartition("_") + if sep and suffix.isdigit(): + depth_indices.append(int(suffix)) + return max(depth_indices) + 1 if depth_indices else 1 + + +# Experiment prognostic and boundary variables +# Assumption that all 3D variables are appended with depth_i_levels +# and all 2D variables do not have any digits / underscores in their names + +RHO_0 = 1035.0 # DENSITY_OF_WATER_CM4 kg/m^3 +CP_SW = 3992.0 # SPECIFIC_HEAT_OF_WATER_CM4 J/kg/K + + +def _select_var_names( + vars_by_key: dict[str, list[str]], + key: str, + dataset_type: DatasetType, + var_kind: str, +) -> list[str]: + try: + return list(vars_by_key[key]) + except KeyError as exc: + choices = ", ".join(sorted(vars_by_key)) + raise ValueError( + f"Unsupported {dataset_type} {var_kind} variable key {key!r}. " + f"Expected one of: {choices}" + ) from exc + + +def build_om4_spec( + prognostic_vars_key: str = "thermo_dynamic_all", + boundary_vars_key: str = "tau_hfds", + grid_type: GridType = "gaussian", +) -> DatasetSpec: + return DatasetSpec( + type="om4", + depth_levels=( + 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, + ), + depth_thickness=( + 5.0, + 10.0, + 15.0, + 20.0, + 30.0, + 50.0, + 70.0, + 100.0, + 150.0, + 200.0, + 250.0, + 300.0, + 400.0, + 500.0, + 600.0, + 800.0, + 1000.0, + 1000.0, + 1000.0, + ), + mask_vars=tuple(f"mask_{i}" for i in range(19)), + mask_all_levels_var="wetmask", + seconds_per_time_step=5 * 24 * 60 * 60, + prognostic_var_names=_select_var_names( + { + "thetao_1": ["thetao_0"], + "thermo_dynamic_5": [ + k + str(j) + for k in ["uo_", "vo_", "thetao_", "so_"] + for j in range(5) + ] + + ["zos"], + "thermo_dynamic_all": [ + k + str(j) + for k in ["uo_", "vo_", "thetao_", "so_"] + for j in range(19) + ] + + ["zos"], + "thermo_all": [ + k + str(j) for k in ["thetao_", "so_"] for j in range(19) + ] + + ["zos"], + }, + prognostic_vars_key, + "om4", + "prognostic", + ), + boundary_var_names=_select_var_names( + { + "hfds": ["hfds"], + "tau_hfds": ["tauuo", "tauvo", "hfds"], + "tau_hfds_hfds_anom": ["tauuo", "tauvo", "hfds", "hfds_anomalies"], + }, + boundary_vars_key, + "om4", + "boundary", + ), + default_metadata={ + "thetao": { + "long_name": "Sea Water Potential Temperature", + "units": r"\degree C", + }, + "so": { + "long_name": "Sea Water Salinity", + "units": "psu", + }, + "uo": { + "long_name": "Sea Water X Velocity", + "units": "m/s", + }, + "vo": { + "long_name": "Sea Water Y Velocity", + "units": "m/s", + }, + "zos": { + "long_name": "Sea surface height above geoid", + "units": "m", + }, + "tos": { + "long_name": "Sea surface temperature", + "units": r"\degree C", + }, + "tauuo": { + "long_name": "Surface Downward X Stress", + "units": "N/m^2", + }, + "tauvo": { + "long_name": "Surface Downward Y Stress", + "units": "N/m^2", + }, + "hfds": { + "long_name": "Surface ocean heat flux from " + "SW+LW+latent+sensible+masstransfer+frazil+seaice_melt_heat", + "units": "W/m^2", + }, + "hfds_anomalies": { + "long_name": "hfds anomalies", + "units": "W/m^2", + }, + }, + ocean_heat_temperature_var="thetao", + surface_heat_flux_var="hfds", + grid_type=grid_type, + ) + + +def build_llc_spec( + prognostic_vars_key: str = "single_1", + boundary_vars_key: str = "single_1", +) -> DatasetSpec: + return DatasetSpec( + type="llc", + depth_levels=( + 0.5, + 1.57, + 2.79, + 4.185, + 5.78, + 7.595, + 9.66, + 12.01, + 14.68, + 17.705, + 21.125, + 24.99, + 29.345, + 34.24, + 39.725, + 45.855, + 52.69, + 60.28, + 68.685, + 77.965, + 88.175, + 99.37, + 111.6, + 124.915, + 139.365, + 154.99, + 171.825, + 189.9, + 209.235, + 229.855, + 251.77, + 274.985, + 299.505, + 325.32, + 352.42, + 380.79, + 410.41, + 441.255, + 473.305, + 506.54, + 540.935, + 576.465, + 613.11, + 650.855, + 689.685, + 729.595, + 770.585, + 812.66, + 855.835, + 900.135, + 945.595, + ), + depth_thickness=( + 1.07, + 1.22, + 1.395, + 1.595, + 1.815, + 2.065, + 2.35, + 2.67, + 3.025, + 3.42, + 3.865, + 4.355, + 4.895, + 5.485, + 6.13, + 6.835, + 7.59, + 8.405, + 9.28, + 10.21, + 11.195, + 12.23, + 13.315, + 14.45, + 15.625, + 16.835, + 18.075, + 19.335, + 20.62, + 21.915, + 23.215, + 24.52, + 25.815, + 27.1, + 28.37, + 29.62, + 30.845, + 32.05, + 33.235, + 34.395, + 35.53, + 36.645, + 37.745, + 38.83, + 39.91, + 40.99, + 42.075, + 43.175, + 44.3, + 45.46, + 54.405, + ), + mask_vars=tuple(f"wetmask_{i}" for i in range(51)), + mask_all_levels_var="wetmask", + seconds_per_time_step=60, + prognostic_var_names=_select_var_names( + { + "single_1": ["Theta_0"], + "all": [ + k + str(j) + for k in ["U_", "V_", "Theta_", "Salt_"] + for j in range(51) + ] + + ["Eta"], + }, + prognostic_vars_key, + "llc", + "prognostic", + ), + boundary_var_names=_select_var_names( + { + "single_1": ["oceQnet"], + "all": ["oceTAUX", "oceTAUY", "oceQnet", "Eta"], + }, + boundary_vars_key, + "llc", + "boundary", + ), + default_metadata={ + "Theta": { + "long_name": "Sea Water Potential Temperature", + "units": r"\degree C", + }, + "Salt": { + "long_name": "Sea Water Salinity", + "units": "psu", + }, + "U": { + "long_name": "Sea Water X Velocity", + "units": "m/s", + }, + "V": { + "long_name": "Sea Water Y Velocity", + "units": "m/s", + }, + "Eta": { + "long_name": "Sea surface height above geoid", + "units": "m", + }, + "oceTAUX": { + "long_name": "Surface Downward X Stress", + "units": "N/m^2", + }, + "oceTAUY": { + "long_name": "Surface Downward Y Stress", + "units": "N/m^2", + }, + "oceQnet": { + "long_name": "Surface ocean heat flux from " + "SW+LW+latent+sensible+masstransfer+frazil+seaice_melt_heat", + "units": "W/m^2", + }, + }, + ocean_heat_temperature_var="Theta", + surface_heat_flux_var="oceQnet", + # LLC (lat-lon-cap) is curvilinear, so its 2D geometry can't be broadcast + # from 1D axes -- same broadcast-unsafe class as the tripolar grid. + grid_type="tripolar", + ) + + +def construct_metadata( + data: xr.Dataset, + dataset_spec: DatasetSpec, +) -> dict[str, dict[str, str]]: + metadata = {} + for var in data.variables: + try: + metadata[str(var)] = { + "long_name": data[var].long_name, + "units": data[var].units, + } + except AttributeError: + if var in dataset_spec.default_metadata.keys(): + metadata[str(var)] = dataset_spec.default_metadata[str(var)] + elif ( + key := str(var).split("_")[0] + ) in dataset_spec.default_metadata.keys(): + metadata[str(var)] = dataset_spec.default_metadata[key] + else: + logger.info(f"{var} does not have any default metadata") + metadata[str(var)] = { + "long_name": "Unknown", + "units": "Unknown", + } + + return metadata + + +class LoaderVersion(enum.Enum): + OM4_TORCH = "om4-torch" + + +class TensorMap: + def __init__( + self, + dataset_spec: DatasetSpec, + ): + """ + Maps input variables / depth levels to their indices in the input tensor. + + VAR_3D_IDX maps the input variables to their indices in the input tensor + DP_3D_IDX maps the depth levels to their indices in the input tensor + """ + self.dataset_spec = dataset_spec + self.VAR_3D_IDX: dict[str, torch.Tensor] = {} + self.DP_3D_IDX: dict[str, torch.Tensor] = {} + + self.INPT_BOUNDARY_IDX: dict[str, torch.Tensor] = {} + self.VAR_SET_2D = [] + self.VAR_SET_3D = [] + self.prognostic_var_names = dataset_spec.prognostic_var_names + self.boundary_var_names = dataset_spec.boundary_var_names + for out in self.prognostic_var_names: + var_split = out.split("_") + if len(var_split) == 1: + self.VAR_SET_2D.append(var_split[0]) + else: + self.VAR_SET_3D.append(var_split[0]) + + # Consistent order of variables + self.VAR_SET = list( + dict.fromkeys([out.split("_")[0] for out in self.prognostic_var_names]) + ) + + levels = dataset_spec.num_prognostic_depth_levels + + self.DEPTH_SET = list(dataset_spec.depth_i_levels[:levels]) + self.dz = torch.tensor(dataset_spec.depth_thickness[:levels]) + + self._populate_var_3d_idx() + self._populate_dp_3d_idx() + self._populate_boundary_idx() + + def _populate_var_3d_idx(self): + for kt in self.VAR_SET: + self.VAR_3D_IDX[kt] = torch.tensor([]) + for i, k in enumerate(self.prognostic_var_names): + if kt in k: + self.VAR_3D_IDX[kt] = torch.cat( + [self.VAR_3D_IDX[kt], torch.tensor([i])] + ) + self.VAR_3D_IDX[kt] = self.VAR_3D_IDX[kt].to(torch.int32) + + def _populate_dp_3d_idx(self): + for d in self.DEPTH_SET: + self.DP_3D_IDX[d] = torch.tensor([]) + for i, k in enumerate(self.prognostic_var_names): + k_split = k.split("_") + if len(k_split) == 1: + continue + elif d == k_split[-1]: + self.DP_3D_IDX[d] = torch.cat( + [self.DP_3D_IDX[d], torch.tensor([i])] + ) + self.DP_3D_IDX[d] = self.DP_3D_IDX[d].to(torch.int32) + + self.DP_3D_IDX[self.DEPTH_SET[0]] = torch.cat( + [ + self.DP_3D_IDX[self.DEPTH_SET[0]], + torch.tensor([self.VAR_3D_IDX[var_2D] for var_2D in self.VAR_SET_2D]), + ] + ).to(torch.int32) + + def _populate_boundary_idx(self): + """ + Populates the indices of the boundary variables in the input tensor. + + We assume the indices INPT_BOUNDARY_IDX will be used after the boundary + condition is extracted from the input tensor + """ + for i, k in enumerate(self.boundary_var_names): + self.INPT_BOUNDARY_IDX[k] = torch.tensor([i]) + + def to(self, device: torch.device) -> Self: + """Move all index tensors to the given device. + + Call this once after model initialization so that indexing a GPU tensor + with these indices stays on-device and avoids implicit CUDA syncs. + """ + for k in self.VAR_3D_IDX: + self.VAR_3D_IDX[k] = self.VAR_3D_IDX[k].to(device) + for k in self.DP_3D_IDX: + self.DP_3D_IDX[k] = self.DP_3D_IDX[k].to(device) + for k in self.INPT_BOUNDARY_IDX: + self.INPT_BOUNDARY_IDX[k] = self.INPT_BOUNDARY_IDX[k].to(device) + self.dz = self.dz.to(device) + return self diff --git a/samudra/datasets.py b/samudra/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..5955ca41efd8ee4667463ff7900ceffecbdaa93a --- /dev/null +++ b/samudra/datasets.py @@ -0,0 +1,726 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import logging +import time +from concurrent.futures import wait +from concurrent.futures.thread import ThreadPoolExecutor +from typing import ClassVar, final + +import numpy as np +import torch +import xarray as xr +from einops import rearrange +from jaxtyping import Float +from torch.utils.data import Dataset +from xarray_einstats.einops import rearrange as xr_rearrange # noqa: F401 + +from samudra.constants import ( + Boundary, + BoundaryVarNames, + Example, + GridMask, + Input, + LoaderVersion, + Prognostic, + PrognosticMask, + PrognosticVarNames, +) +from samudra.utils.ctx import GridContext +from samudra.utils.data import DataSource, LoadStats, OceanData, conditional_rearrange +from samudra.utils.device import using_gpu +from samudra.utils.logging import elapsed + +logger = logging.getLogger(__name__) + + +class InferenceDataset(Dataset): + """This class is used for inference rollouts. + + It creates rolling indices to keep track of histories/past states. + For example, + Hist=0 ; 0->[0, 1]; 1->[1, 2]; 2->[2, 3]; 3->[3, 4] + Hist=1 ; 0->[[0, 1], [2, 3]]; 1->[[2, 3], [4, 5]]; + 2->[[4, 5], [6, 7]]; 3->[[6, 7], [8, 9]] + Hist=2 ; 0->[[0, 1, 2], [3, 4, 5]]; + 1->[[3, 4, 5], [6, 7, 8]]; + 2->[[6, 7, 8], [9, 10, 11]]; + 3->[[9, 10, 11], [12, 13, 14]] + """ + + @elapsed + def __init__( + self, + src: DataSource, + prognostic_var_names, + boundary_var_names, + hist, + normalize_before_mask, + masked_fill_value, + long_rollout, + ): + super().__init__() + # NOTE: Keep tensors on CPU during initialization. This allows the dataset + # to be passed between DataLoader worker processes. Call to(device) before + # using the dataset for inference. + + self.hist = hist + + self.num_prognostic_channels = (hist + 1) * len(prognostic_var_names) + data = src.data + self.input_res = src.resolution + self._prognostic_src = src.filter(prognostic_var_names, prefix="prognostic") + self._boundary_src = src.filter(boundary_var_names, prefix="boundary") + self._times = data.time + self.normalize_before_mask = normalize_before_mask + self.masked_fill_value = masked_fill_value + + time_indices = np.arange(data.time.size) + indices = xr.DataArray( + time_indices, + dims=["time"], + coords={"time": time_indices}, + ) + total_steps = 2 * self.hist + 1 + rolling_indices = indices.rolling( + time=len(time_indices) - total_steps, center=False + ).construct("window_dim") + rolling_indices = rolling_indices.transpose("window_dim", "time").isel( + time=slice(len(time_indices) - total_steps - 1, None) + ) # Remove first few null indices + self.rolling_indices = rolling_indices.isel( + window_dim=slice(0, None, self.hist + 1) + ) # Skip indices based on history + self.rolling_indices = self.rolling_indices.astype(int) + + if long_rollout: + logger.info( + f"Long rollout will use input at time {data.time.values[0]} and produce" + f" output at {data.time.values[self.hist + 1]}" + ) + + self.wet: PrognosticMask = src.masks.prognostic + self.wet_surface: GridMask = src.masks.boundary + self.wet_label = src.masks.prognostic_with_hist(self.hist) + self.size = len(self.rolling_indices) + + if using_gpu(): + self.wet = self.wet.pin_memory() + self.wet_surface = self.wet_surface.pin_memory() + self.wet_label = self.wet_label.pin_memory() + + # Inference only currently supports the same output resolution as the input + # resolution. + self.ctx = GridContext(self.wet_label, self.input_res, self.input_res) + + def __len__(self): + return self.size + + def to(self, device: torch.device) -> "InferenceDataset": + """Move the dataset's context tensors to the specified device. + + Call this before using the dataset for inference to ensure tensors + are on the correct device (GPU). + """ + self.ctx = self.ctx.to(device) + self.wet_label = self.wet_label.to(device, non_blocking=True) + return self + + @property + def initial_prognostic(self): + x_index = self._get_x_index(0) + data_in = self._get_prognostic(x_index) + return data_in + + def inference_target(self, step: int | slice): + x_index = self._get_x_index(step) + label = self._get_label(x_index) + return label + + def get_initial_input(self) -> tuple[Prognostic, Boundary]: + prog, boundary, _ = self.__getitem__(0) + return prog, boundary + + def get_target_time(self, start_step: int, num_steps: int): + x_index = self._get_x_index(start_step) + batch_index = x_index.values[0] + steps_predicted = len(batch_index) // 2 + start_target_index = batch_index[steps_predicted] + + return self._times.isel( + time=slice( + start_target_index, start_target_index + num_steps * steps_predicted + ) + ) + + def get_boundary(self, step: int) -> Boundary: + """Return boundary at the requested step.""" + x_index = self._get_x_index(step) + boundary = self._get_boundary(x_index) + return boundary + + @elapsed(level=logging.DEBUG) + def __getitem__(self, idx): + x_index = self._get_x_index(idx) + data_in_prog = self._get_prognostic(x_index) + data_in_boundary = self._get_boundary(x_index) + label = self._get_label(x_index) + return (data_in_prog, data_in_boundary, label) + + def _get_x_index(self, idx): + if isinstance(idx, slice): + if ( + (idx.start is not None and idx.start < 0) + or (idx.stop is not None and idx.stop < 0) + or (idx.step is not None and idx.step < 0) + ): + raise IndexError("Sorry, negative indexing is not supported!") + if idx.step is None: + idx = slice(idx.start, idx.stop, 1) + if idx.start is None and idx.stop is None: + idx = slice(0, self.size, idx.step) + elif idx.start is None: + idx = slice(0, idx.stop, idx.step) + elif idx.stop is None: + idx = slice(idx.start, self.size, idx.step) + elif isinstance(idx, int): + if idx < 0: + raise IndexError("Sorry, negative indexing is not supported!") + elif idx >= self.size: + raise IndexError(f"Index {idx} out of range with size {self.size}") + idx = slice(idx, idx + 1, 1) + + rolling_idx = self.rolling_indices.isel(window_dim=idx) + x_index = xr.Variable(["window_dim", "time"], rolling_idx) + + return x_index + + def _get_prognostic(self, x_index): + data_in_src = self._prognostic_src.map_data( + lambda ds: ds.isel(time=x_index).isel(time=slice(None, self.hist + 1)) + ) + if self.normalize_before_mask: + data_in_ds = data_in_src.normalize() + else: + data_in_ds = data_in_src.data + + if "lev" in data_in_ds.dims: + data_in_np: np.ndarray = ( + conditional_rearrange( + data_in_ds, + "window_dim time (variable lev)=var lat lon", + concat_dim="var", + ) + .rename({"var": "variable"}) + .to_numpy() + ) + else: + data_in_np = ( + data_in_ds.to_array() + .transpose("window_dim", "time", "variable", "lat", "lon") + .to_numpy() + ) + data_in: torch.Tensor = torch.from_numpy(data_in_np).float() + data_in = torch.where(self.wet, data_in, self.masked_fill_value) + if not self.normalize_before_mask: + data_in = self._prognostic_src.normalize_with(data_in, variable_axis=2) + data_in = rearrange( + data_in, + "window_dim time variable lat lon -> window_dim (time variable) lat lon", + ) + return data_in + + def _get_boundary(self, x_index): + """ + This function returns the boundary condition for the current time step. + + With hist > 0, the boundary condition considered is always the last step of + the input. + """ + data_in_boundary_src = self._boundary_src.map_data( + lambda ds: ds.isel(time=x_index).isel(time=slice(None, self.hist + 1)) + ) + if self.normalize_before_mask: + data_in_boundary_ds = data_in_boundary_src.normalize() + else: + data_in_boundary_ds = data_in_boundary_src.data + data_in_boundary_np: np.ndarray = ( + data_in_boundary_ds.to_array() + .transpose("window_dim", "time", "variable", "lat", "lon") + .to_numpy() + ) + data_in_boundary: torch.Tensor = torch.from_numpy(data_in_boundary_np).float() + data_in_boundary = torch.where( + self.wet_surface, data_in_boundary, self.masked_fill_value + ) + if not self.normalize_before_mask: + data_in_boundary = self._boundary_src.normalize_with( + data_in_boundary, variable_axis=2 + ) + data_in_boundary = rearrange( + data_in_boundary, + "window_dim time variable lat lon -> window_dim (time variable) lat lon", + ) + return data_in_boundary + + def _get_label(self, x_index): + label_src = self._prognostic_src.map_data( + lambda ds: ds.isel(time=x_index).isel(time=slice(self.hist + 1, None)) + ) + if self.normalize_before_mask: + label_ds = label_src.normalize() + else: + label_ds = label_src.data + if "lev" in label_ds.dims: + label_np: np.ndarray = ( + conditional_rearrange( + label_ds, + "window_dim time (variable lev)=var lat lon", + concat_dim="var", + ) + .rename({"var": "variable"}) + .to_numpy() + ) + else: + label_np = ( + label_ds.to_array() + .transpose("window_dim", "time", "variable", "lat", "lon") + .to_numpy() + ) + label: torch.Tensor = torch.from_numpy(label_np).float() + label = torch.where(self.wet, label, self.masked_fill_value) + if not self.normalize_before_mask: + label = self._prognostic_src.normalize_with(label, variable_axis=2) + label = rearrange( + label, + "window_dim time variable lat lon -> window_dim (time variable) lat lon", + ) + return label + + def get_coords_dict(self): + return { + co: self._prognostic_src.data[co] for co in self._prognostic_src.data.coords + } + + +class InferenceDatasets(Dataset): + def __init__(self, datasets: list[InferenceDataset], lengths: list[int]): + self.datasets = datasets + self.lengths = lengths + + def __len__(self): + return len(self.datasets) + + def __getitem__(self, idx): + return (self.datasets[idx], self.lengths[idx]) + + +class RawTrainData: + def __init__(self, dataset_id: "TorchTrainDataset.Id"): + self.dataset_id: TorchTrainDataset.Id = dataset_id + self.raw_data: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + self.load_stats: LoadStats | None = None + + def insert( + self, + input_: torch.Tensor, + boundary: torch.Tensor, + label: torch.Tensor, + ): + """Add a prognostic input, boundary, and prognostic label as the last step.""" + self.raw_data.append((input_, boundary, label)) + + def to(self, device: torch.device): + self.raw_data = [ + ( + input_.to(device, non_blocking=True), + boundary.to(device, non_blocking=True), + label.to(device, non_blocking=True), + ) + for input_, boundary, label in self.raw_data + ] + + def pin_memory(self): + self.raw_data = [ + ( + input_.pin_memory(), + boundary.pin_memory(), + label.pin_memory(), + ) + for input_, boundary, label in self.raw_data + ] + return self + + +class TrainData: + """A single batch of training data. + + A single batch contains multiple steps worth of ``Example`` entries, each + of which is a ``(prognostic_input, boundary_input, label)`` triple. The + prognostic and boundary tensors are carried separately because the + samudra-multi model encodes them separately (Samudra just concatenates them later). + """ + + def __init__( + self, num_prognostic_channels: int, num_boundary_channels: int, ctx: GridContext + ): + self.num_prognostic_channels = num_prognostic_channels + self.num_boundary_channels = num_boundary_channels + self.ctx = ctx + self.example_by_step: list[Example] = [] + self.load_stats: LoadStats | None = None + + def append( + self, prognostic_input: Prognostic, boundary_input: Boundary, label: Prognostic + ) -> None: + """Add another Example as a new step.""" + self.example_by_step.append((prognostic_input, boundary_input, label)) + + def get_initial_input(self) -> tuple[Prognostic, Boundary]: + return self.get_input(0) + + def get_input(self, step: int) -> tuple[Prognostic, Boundary]: + prog, boundary, _ = self.example_by_step[step] + return prog, boundary + + def get_label(self, step: int) -> Prognostic: + return self.example_by_step[step][2] + + def __getitem__(self, step: int) -> Example: + """Converts index (step) into (prognostic, boundary, label) triple.""" + return self.example_by_step[step] + + def __len__(self) -> int: + return len(self.example_by_step) + + def __iter__(self): + return iter(range(len(self))) + + def to(self, device: torch.device) -> None: + for step in self: + prog, boundary, label = self.example_by_step[step] + self.example_by_step[step] = ( + prog.to(device, non_blocking=True), + boundary.to(device, non_blocking=True), + label.to(device, non_blocking=True), + ) + + def pin_memory(self): + for step in self: + prog, boundary, label = self.example_by_step[step] + self.example_by_step[step] = ( + prog.pin_memory(), + boundary.pin_memory(), + label.pin_memory(), + ) + return self + + +@final +class TorchTrainDataset(Dataset[RawTrainData]): + """ + This class is used for training and validation. + + It creates rolling indices to keep track of histories/past states. But different + from InferenceDataset, as it creates rolling indices based on stride. By default, + the sliding window / stride is 1. + + We make use of TrainData class to store a single sample. + + For example, + Hist=0 ; TD: step=0->[0, 1]; step=1->[1, 2]; step=2->[2, 3]; step=3->[3, 4] + Hist=1 ; TD: step=0->[[0, 1], [2, 3]]; step=1->[[2, 3], [4, 5]]; + step=2->[[4, 5], [6, 7]]; step=3->[[6, 7], [8, 9]] + Hist=2 ; TD: step=0->[[0, 1, 2], [3, 4, 5]]; + step=1->[[3, 4, 5], [6, 7, 8]]; + step=2->[[6, 7, 8], [9, 10, 11]]; + step=3->[[9, 10, 11], [12, 13, 14]] + """ + + type Id = str + + FLAG = LoaderVersion.OM4_TORCH + + # Shared across all instances within a process. Created lazily on first + # __getitem__ call so that each DataLoader worker gets its own clean + # executor. + _shared_executor: ClassVar[ThreadPoolExecutor | None] = None + + @classmethod + def _get_executor(cls) -> ThreadPoolExecutor: + if cls._shared_executor is None: + cls._shared_executor = ThreadPoolExecutor( + max_workers=None, thread_name_prefix="concurrent_compute" + ) + return cls._shared_executor + + @elapsed + def __init__( + self, + src: DataSource, + prognostic_var_names: PrognosticVarNames, + boundary_var_names: BoundaryVarNames, + hist: int, + steps: int, + normalize_before_mask: bool, + masked_fill_value: float, + stride: int = 1, + concurrent_compute_: bool = False, + ): + super().__init__() + self.id = f"{self.__class__.__name__}_{str(id(self))}" + + self.hist: int = hist + self.steps: int = steps + self.stride: int = stride + self.normalize_before_mask: bool = normalize_before_mask + self.masked_fill_value: float = masked_fill_value + self._concurrent_compute = concurrent_compute_ + + self.num_prognostic_channels: int = (hist + 1) * len(prognostic_var_names) + self.num_boundary_channels: int = (hist + 1) * len(boundary_var_names) + time_ = src.data.time + self.prognostic_src = src.filter(prognostic_var_names, prefix="prog") + self.boundary_src = src.filter(boundary_var_names, prefix="boundary") + + # This class will be used only for training and validation + total_steps: int = 2 * self.hist + 2 + + # Calculate the number of windows + num_windows = time_.size - (total_steps - 1) * self.stride + + # Create base indices + indices = np.arange(num_windows) + indices_da = xr.DataArray(indices, dims=["window"]) + + # Create window dimension + window_dim = xr.DataArray(np.arange(total_steps), dims=["time"]) + + # Construct rolling indices + self.rolling_indices: Float[xr.DataArray, "window time"] = ( + indices_da + stride * window_dim + ) + + # NB(alxmrs): Keep masks on CPU - will be moved to GPU in to_train_data() + self.wet_prognostic: PrognosticMask = src.masks.prognostic + self.wet_surface: GridMask = src.masks.boundary + + self.ctx = GridContext( + label_mask=self.prognostic_src.masks.prognostic_with_hist(self.hist), + input_resolution_cpu=self.prognostic_src.resolution, + output_resolution_cpu=self.prognostic_src.resolution, + ) + + self.size: int = ( + time_.size + - self.steps * (self.hist + 1) * self.stride + - self.hist * self.stride + ) + + def __len__(self) -> int: + return self.size + + @elapsed(level=logging.DEBUG) + def __getitem__(self, idx: int): + start_time = time.perf_counter() + TD = RawTrainData(self.id) + + for step in range(self.steps): + x_index = self._get_x_index(idx, step) + current_x_index = x_index.isel(time=slice(0, self.hist + 1)) + forecast_x_index = x_index.isel(time=slice(self.hist + 1, None)) + + # Only materialize the time ranges we actually use to reduce memory. + input_selected = self.prognostic_src.data.isel(time=current_x_index) + boundary_selected = self.boundary_src.data.isel(time=current_x_index) + label_selected = self.prognostic_src.data.isel( + time=forecast_x_index + ) # forecasted data + prognostic_selected = [input_selected, label_selected] + + if self._concurrent_compute: + datasets = prognostic_selected + [boundary_selected] + concurrent_compute( + *datasets, + executor=self._get_executor(), + ) + + if "lev" in prognostic_selected[0].dims: + prognostics = [ + torch.from_numpy( + conditional_rearrange( + selected, + "time (variable lev)=var lat lon", + concat_dim="var", + ) + .rename({"var": "variable"}) + .to_numpy() + .astype(np.float32, copy=False) + ) + for selected in prognostic_selected + ] + else: + prognostics = [ + torch.from_numpy( + selected.to_array() + .transpose("time", "variable", "lat", "lon") + .to_numpy() + .astype(np.float32, copy=False) + ) + for selected in prognostic_selected + ] + boundary = torch.from_numpy( + boundary_selected.to_array() + .transpose("time", "variable", "lat", "lon") + .to_numpy() + .astype(np.float32, copy=False) + ) + input_, label = prognostics[0], prognostics[-1] + TD.insert(input_, boundary, label) + TD.load_stats = LoadStats(time.perf_counter() - start_time) + + return TD + + def to_train_data( + self, raw_train_data: RawTrainData, device: torch.device + ) -> TrainData: + """Convert RawTrainData to TrainData, moving tensors to the specified device. + + Args: + raw_train_data: CPU data from worker process + device: Target device (typically GPU) to move tensors to + + Returns: + TrainData with tensors on the target device + """ + train_data = TrainData( + self.num_prognostic_channels, + self.num_boundary_channels, + self.ctx.to(device), + ) + for input_, boundary, label in raw_train_data.raw_data: + prog_input, boundary_input, label_tensor = self._to_example( + OceanData.from_data_source( + input_, + self.wet_prognostic, + self.prognostic_src, + ).to(device=device, non_blocking=True), + OceanData.from_data_source( + boundary, + self.wet_surface, + self.boundary_src, + ).to(device=device, non_blocking=True), + OceanData.from_data_source( + label, self.wet_prognostic, self.prognostic_src + ).to(device=device, non_blocking=True), + ) + train_data.append(prog_input, boundary_input, label_tensor) + train_data.load_stats = raw_train_data.load_stats + return train_data + + def _to_example( + self, input_: OceanData, boundary: OceanData, label: OceanData + ) -> Example: + # Input/boundary only include current steps; label only includes forecasted steps. + prog_input = self._prep_tensor_steps(input_) + boundary_input = self._prep_tensor_steps(boundary) + label_tensor = self._prep_tensor_steps(label) + return prog_input, boundary_input, label_tensor + + def _prep_tensor_steps(self, ocean_data: OceanData) -> Input: + """Normalize, mask, and flatten (time, variable) dims into a channel dim.""" + steps = ocean_data.normalize_and_mask( + self.normalize_before_mask, self.masked_fill_value + ) + return rearrange( + steps, "batch time variable lat lon -> batch (time variable) lat lon" + ) + + def _get_x_index(self, idx: int, step: int) -> xr.DataArray: + assert isinstance(idx, int) + if idx < 0: + raise IndexError("Sorry, negative indexing is not supported!") + if idx >= len(self): + raise IndexError("Index out of range") + + window_index = idx + step * (self.hist + 1) * self.stride + return self.rolling_indices.isel(window=window_index, drop=True) + + +def concurrent_compute( + *datasets: xr.Dataset, + executor: ThreadPoolExecutor, +) -> None: + def load_variable_data(var: xr.Variable) -> None: + var.load() + + futures = [] + for ds in datasets: + for var in ds.data_vars.variables.values(): + futures.append(executor.submit(load_variable_data, var)) + + wait(futures) + + +@final +class TrainDataLoader: + """Wrapper around a torch DataLoader that handles GPU post-processing. + + This class wraps a DataLoader[RawTrainData] and converts the raw data + to TrainData by applying GPU-based normalization and masking. This allows + the data loading process to handle I/O while the main process handles + GPU operations. + + Since the data samples flow from one process to the other, we want to tie + them back to the dataset they came from which knows how to do that second + half once they're in the main process which has GPU access set up. To do that, + each data sample (which could come from a different dataset) has a dataset ID + -- `datasets` maps from those IDs to the original datasets. + """ + + def __init__( + self, + dataloader: torch.utils.data.DataLoader[RawTrainData], + datasets: list[TorchTrainDataset], + device: torch.device, + ): + self._dataloader = dataloader + self._datasets = {dataset.id: dataset for dataset in datasets} + self._device = device + + def __iter__(self): + """Iterate over the dataloader, converting RawTrainData to TrainData.""" + for raw_train_data in self._dataloader: + dataset = self._datasets[raw_train_data.dataset_id] + train_data = dataset.to_train_data(raw_train_data, self._device) + yield train_data + + def __len__(self) -> int: + return len(self._dataloader) + + def __getitem__(self, index: int) -> TrainData: + """Access a single item by index, converting RawTrainData to TrainData. + + Note: This bypasses the DataLoader's sampling/batching and directly accesses + the underlying dataset for test purposes. + """ + # Access the underlying dataset directly + raw_train_data = self._dataloader.dataset[index] + # Apply the collate function to add batch dimension (expects a list) + collate_fn = self._dataloader.collate_fn + if collate_fn is not None: + raw_train_data = collate_fn([raw_train_data]) + # Get the dataset that created this raw data + dataset = self._datasets[raw_train_data.dataset_id] + # Convert to TrainData + train_data = dataset.to_train_data(raw_train_data, self._device) + return train_data + + @property + def dataset(self): + return self._dataloader.dataset + + @property + def sampler(self): + return self._dataloader.sampler diff --git a/samudra/derived_variables.py b/samudra/derived_variables.py new file mode 100644 index 0000000000000000000000000000000000000000..ae5d6f8f795300eb7934a48dcadded67b7e43e12 --- /dev/null +++ b/samudra/derived_variables.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Callable + +import torch +from jaxtyping import Float + +from samudra.constants import CP_SW, RHO_0, Grid, TensorMap + + +def compute_ocean_heat_content( + T: Float[Grid, "*batch depth"] | Float[Grid, "*batch hist depth"], dz: torch.Tensor +) -> Float[Grid, "*batch"] | Float[Grid, "*batch hist"]: + """Compute the heat content of the ocean. + + Args: + T: Temperature tensor + dz: Depth tensor of shape (depth,) + + Returns: + Heat content tensor + """ + mask = torch.isnan(T).any() + + # Compute heat content per layer + if T.ndim == 4: + depth_dim = 1 + surface_wet_mask = ~T[:, 0].isnan() + elif T.ndim == 5: + depth_dim = 2 + surface_wet_mask = ~T[:, :, 0].isnan() + else: + raise ValueError(f"Invalid number of dimensions: {T.ndim}") + + view_shape = [1] * T.ndim + view_shape[depth_dim] = -1 + + # Compute heat content per layer + HC_t = RHO_0 * CP_SW * T * dz.view(view_shape) + + # Column integrated heat content + total_HC_t = torch.nansum(HC_t, dim=depth_dim) + + if mask: + total_HC_t = torch.where(surface_wet_mask, total_HC_t, torch.nan) + + return total_HC_t + + +def compute_global_ocean_heat_content( + T: Float[Grid, "*batch depth"] | Float[Grid, "*batch hist depth"], + dz: torch.Tensor, + area_weighted_func: Callable, +) -> torch.Tensor: + """Compute the global heat content of the ocean. + + Args: + T: Temperature tensor of shape + dz: Depth tensor of shape (depth,) + area_weighted_func: Area weighted function + Returns: + Global heat content tensor of shape (batch_size,) + """ + total_HC_t = compute_ocean_heat_content(T, dz) + + # Area weighted sum + global_HC_t = area_weighted_func(total_HC_t) # (batch,) [J] + + return global_HC_t + + +def add_derived_variables( + tensor_out: torch.Tensor, + *, + tensor_map: TensorMap, +) -> dict[str, torch.Tensor]: + """ + Add derived variables to the output. + """ + # Ocean heat content + derived_vars = {} + dz = tensor_map.dz.to(tensor_out.device) + thetao = tensor_out[ + :, :, tensor_map.VAR_3D_IDX[tensor_map.dataset_spec.ocean_heat_temperature_var] + ] + ohct = compute_ocean_heat_content(thetao, dz) + derived_vars["ocean_heat_content"] = ohct + + return derived_vars diff --git a/samudra/eval.py b/samudra/eval.py new file mode 100644 index 0000000000000000000000000000000000000000..c1d590eee6700abde787995af1694d2ab0b99405 --- /dev/null +++ b/samudra/eval.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import datetime +import logging +import time +from collections import OrderedDict + +import torch + +from samudra.aggregator import Aggregator +from samudra.backend import init_eval_backend +from samudra.config import EvalConfig, ObsMetricsConfig +from samudra.constants import BoundaryVarNames, Grid, PrognosticVarNames, TensorMap +from samudra.datasets import InferenceDataset +from samudra.metrics.run import open_predictions, run_observation_metrics +from samudra.stepper import run_rollout +from samudra.utils.data import Normalize, get_inference_steps, spherical_area_weights +from samudra.utils.device import using_gpu +from samudra.utils.distributed import is_main_process, set_seed +from samudra.utils.logging import get_model_summary, handle_logging, handle_warnings +from samudra.utils.wandb import WandBLogger + +logger = logging.getLogger(__name__) + + +class Eval: + """Evaluation pipeline for ocean emulator models. + + Runs long autoregressive rollouts and computes metrics against ground-truth + ocean states. + """ + + def __init__(self, cfg: EvalConfig) -> None: + cfg.prepare_output_dirs() + + self.device = init_eval_backend(cfg.backend) + + # Adjust workers and memory pinning based on device + data_num_workers = cfg.data.loading.num_pytorch_workers() + if not using_gpu(): + data_num_workers = 0 # Disable multi-processing on CPU + elif cfg.disk_mode: + data_num_workers = torch.cuda.device_count() * data_num_workers + + # Set seeds + set_seed(cfg.experiment.rand_seed) + + logger.info("Loading data") + self.data_container = cfg.data.build( + cfg.experiment.resolved_data_root, + ) + + # Getting prognostic and boundary variables + self.dataset_spec = self.data_container.dataset_spec + self.prognostic_var_names: PrognosticVarNames = ( + self.dataset_spec.prognostic_var_names + ) + self.boundary_var_names: BoundaryVarNames = self.dataset_spec.boundary_var_names + self.levels = self.dataset_spec.num_prognostic_depth_levels + + str_prognostics = ", ".join([i for i in self.prognostic_var_names]) + str_boundaries = ", ".join([i for i in self.boundary_var_names]) + + logger.info(f"Prognostic variables: {str_prognostics}") + logger.info(f"Boundary variables: {str_boundaries}") + logger.info(f"Levels: {self.levels}") + + self.N_bound = len(self.boundary_var_names) + self.N_prog = len(self.prognostic_var_names) + + self.num_prog_in = int((cfg.data.hist + 1) * self.N_prog) + self.num_boundary_in = int((cfg.data.hist + 1) * self.N_bound) + self.num_in = self.num_prog_in + self.num_boundary_in + self.num_out = self.num_prog_in + + self.tensor_map = TensorMap(dataset_spec=self.dataset_spec).to(self.device) + + logger.info(f"Number of inputs (prognostic + boundary): {self.num_in}") + logger.info(f"Number of outputs (prognostic): {self.num_out}") + + # Dataloaders + if self.data_container.inference_source is None: + raise ValueError( + "Inference time is not configured for the first data source" + ) + self.src = self.data_container.inference_source + self.data = self.src.data + self.metadata = self.src.metadata + self.wet = self.src.masks.prognostic_with_hist(cfg.data.hist) + self.area_weights: Grid = spherical_area_weights(self.data) + self.area_weights = self.area_weights.to(self.device) + + self.normalize = Normalize( + self.src, + prognostic_var_names=self.prognostic_var_names, + boundary_var_names=self.boundary_var_names, + ) + + # Model + self.model = cfg.model.build( + prog_channels=self.num_prog_in, + boundary_channels=self.num_boundary_in, + out_channels=self.num_out, + hist=cfg.data.hist, + srcs=self.data_container.train_sources, + ).to(self.device) + + get_model_summary(self.model, None, cfg.debug) + + if cfg.ckpt_path is None: + raise ValueError( + "ckpt_path must be set; try --ckpt_path=path/to/checkpoint" + ) + self.load_checkpoint(cfg.ckpt_path) + + self.network = self.model.__class__.__name__ + + # Initialize WandB + self.wandb_logger = WandBLogger.init_instance() + self.wandb_logger.configure( + cfg.experiment.wandb.mode == "online", is_main_process() + ) + + # Set up wandb run + self.wandb_id, self.wandb_name = self.wandb_logger.setup_run( + None, cfg, data_container=self.data_container, finetune=False + ) + + # Eval + self.hist = cfg.data.hist + self.output_dir = cfg.experiment.output_dir + self.debug = cfg.debug + self.num_workers = data_num_workers + self.num_model_steps_forward = cfg.num_model_steps_forward + self.save_zarr = cfg.save_zarr + self.model_path = cfg.ckpt_path + self.normalize_before_mask = cfg.data.normalize_before_mask + self.masked_fill_value = cfg.data.masked_fill_value + self.observations = cfg.observations + self.resolved_data_root = cfg.experiment.resolved_data_root + self.experiment_name = cfg.experiment.name + self.init_inference_store() + + def load_checkpoint(self, ckpt_path: str): + checkpoint = torch.load(ckpt_path, map_location=torch.device(self.device)) + model_state_dict = checkpoint["model"] + new_state_dict = OrderedDict() + for k, v in model_state_dict.items(): + name = k.removeprefix("module.") + new_state_dict[name] = v + self.model.load_state_dict(new_state_dict) + + def init_inference_store(self): + self.num_time_steps = get_inference_steps( + self.src, + hist=self.hist, + ) + self.inference_dataset = InferenceDataset( + src=self.src, + prognostic_var_names=self.prognostic_var_names, + boundary_var_names=self.boundary_var_names, + hist=self.hist, + normalize_before_mask=self.normalize_before_mask, + masked_fill_value=self.masked_fill_value, + long_rollout=True, + ) + + def run(self) -> None: + start_time = time.perf_counter() + inf_stats = self.standalone_inference() + time_elapsed = time.perf_counter() - start_time + + log_stats = { + **inf_stats, + "eval_total_seconds": time_elapsed, + } + + if is_main_process(): + self.wandb_logger.log(log_stats, step=None) + + # Logged after the rollout's own metrics, deliberately: the observation + # phase reads the rollout back off disk and can fail on data problems + # that have nothing to do with the model. When it does, it should raise + # loudly without also costing us the rollout results already computed. + try: + if self.observations is not None and is_main_process(): + self.report_observation_metrics(self.observations) + finally: + # The observation phase reads ~65 GB back off remote storage, so it + # can fail for reasons unrelated to the model. It should still + # raise -- but not by skipping `finish()`, which would leave the run + # marked crashed in W&B with the rollout's own metrics stranded + # inside it. + total_time = time.perf_counter() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + logger.info(f"Eval time (Including wandb logging) {total_time_str}") + self.finish() + + def report_observation_metrics(self, obs_cfg: ObsMetricsConfig) -> None: + """Score the finished rollout against observation products.""" + logger.info("Computing observation metrics") + baselines = {} + if "om4" in obs_cfg.baselines: + # The ground-truth data the eval already staged. The metric table + # is only interpretable next to it, so scoring it costs one extra + # pass over data that is local anyway. + baselines["om4"] = self.src.data + + frame, scalars = run_observation_metrics( + obs_cfg, + predictions=open_predictions(self.output_dir), + dataset_spec=self.dataset_spec, + data_root=self.resolved_data_root, + model_label=self.experiment_name, + baselines=baselines, + output_dir=self.output_dir, + ) + self.wandb_logger.log( + {**scalars, "obs/metrics_table": self.wandb_logger.Table(dataframe=frame)}, + step=None, + ) + + @torch.no_grad() + def standalone_inference(self): + self.model.eval() + inf_aggregator = Aggregator.get_standalone_inference_aggregator( + self.num_time_steps, + self.metadata, + self.hist, + self.area_weights, + self.src.masks.prognostic.to(self.device), + self.num_out, + self.tensor_map, + self.normalize, + self.prognostic_var_names, + ) + + run_rollout( + model=self.model, + dataset=self.inference_dataset, + inf_aggregator=inf_aggregator, + epoch=0, + output_dir=self.output_dir, + model_path=self.model_path, + num_model_steps_forward=self.num_model_steps_forward, + save_zarr=self.save_zarr, + tensor_map=self.tensor_map, + normalize=self.normalize, + ) + logs = inf_aggregator.get_summary_logs() + return {f"inference/{k}": v for k, v in logs.items()} + + def finish(self): + self.wandb_logger.finish() + + +def main(): + cfg = EvalConfig.from_yaml_and_cli() + cfg.prepare_output_dirs() # we do this first so logging can use them + + handle_logging(cfg.debug, cfg.experiment.output_dir) + handle_warnings() + + Evaluator = Eval(cfg) + + try: + Evaluator.run() + except Exception as e: + # Log the exception with traceback + logger.exception("Evaluation failed with an exception") + raise e + + +if __name__ == "__main__": + main() diff --git a/samudra/metrics/__init__.py b/samudra/metrics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aac5532dbef74cb3e7362beef7953b6762d89f5e --- /dev/null +++ b/samudra/metrics/__init__.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Observation-based evaluation metrics for finished rollouts. + +Unlike `samudra.aggregator`, which accumulates statistics incrementally as a +rollout streams past, these metrics operate on a *finished* rollout that has +been written to disk. Several of them (calendar-year block bootstrap, detrended +and deseasonalized residual variance) have no incremental form at all: they need +the whole multi-year time series in hand. + +The pieces: + +- `kernels`: pure functions over `xr.DataArray`. No I/O, no plotting, no state. +- `spectra`: the same, for the transforms that answer *at which scales* a model + is wrong rather than by how much. +- `observations`: loaders for the DUACS / OISST / ARGO-IAP products. +- `comparisons`: model and observation fields aligned and put on one grid, which + is everything scoring and drawing have in common. +- `report`: the driver, producing one tidy `pandas.DataFrame` of metric rows. +- `run`: one observation pass, returning the products, the prepared rollouts, + the frame and the W&B scalars. + +`ObsMetricsConfig` lives in `samudra.config` with the other task configs. + +`samudra.eval` takes the scalars from `run` and logs them; `samudra.viz` takes +the same pass and draws it. Both reduce through `kernels`, so a number on a +figure cannot drift from the scalar logged beside it: they are one calculation. + +One rule is easy to get wrong and worth stating up front: **reduce before you +regrid, whenever the reduction is nonlinear.** Horizontal interpolation is a +weighted average, so it commutes with differences but not with variance, energy, +or anything else quadratic. Interpolating first blends neighbouring cells and +destroys the fine-scale structure those quantities are largely made of, biasing +them low. Residual variance is therefore computed on the model's native grid and +only the resulting variance map is regridded. + +The bias is invisible when the observation grid is as fine as the model's, which +is what makes it dangerous: it hid in SST (quarter degree to quarter degree, a +0.18% error) while costing 6% in OHC (quarter to half degree), and it scales +with the resolution ratio -- so it distorts comparisons between models of +different resolution rather than applying a constant offset. +""" diff --git a/samudra/metrics/__main__.py b/samudra/metrics/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..430509c02af3ce62cac20a0c90a1d40d038884ab --- /dev/null +++ b/samudra/metrics/__main__.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +r"""Score an existing rollout against observations, without rerunning it. + + python -m samudra.metrics configs/samudra_om4_v2/eval.yaml \\ + --experiment.name=my_past_eval + +Takes the same `EvalConfig` as `samudra.eval` and reads the `predictions.zarr` +that job already wrote, so iterating on a metric costs minutes rather than the +hours a fresh rollout would. No GPU and no checkpoint required. +""" + +import logging + +from samudra.config import EvalConfig +from samudra.metrics.run import open_predictions, run_observation_metrics +from samudra.utils.data import DataContainer +from samudra.utils.logging import handle_logging, handle_warnings + +logger = logging.getLogger(__name__) + + +def main() -> None: + cfg = EvalConfig.from_yaml_and_cli() + cfg.prepare_output_dirs() + handle_logging(cfg.debug, cfg.experiment.output_dir) + handle_warnings() + + if cfg.observations is None: + raise ValueError( + "This config has no `observations` block, so there is nothing to " + "score. Add one (see configs/data/obs.yaml) or use a config that " + "includes it." + ) + + data_container: DataContainer = cfg.data.build(cfg.experiment.resolved_data_root) + src = data_container.inference_source + if src is None: + raise ValueError("Inference time is not configured for the first data source") + + baselines = {"om4": src.data} if "om4" in cfg.observations.baselines else {} + + frame, scalars = run_observation_metrics( + cfg.observations, + predictions=open_predictions(cfg.experiment.output_dir), + dataset_spec=data_container.dataset_spec, + data_root=cfg.experiment.resolved_data_root, + model_label=cfg.experiment.name, + baselines=baselines, + output_dir=cfg.experiment.output_dir, + ) + + logger.info("Observation metrics:") + for key, value in sorted(scalars.items()): + logger.info(" %s = %s", key, value) + logger.info("Wrote %d rows of metric detail alongside the rollout", len(frame)) + + +if __name__ == "__main__": + main() diff --git a/samudra/metrics/comparisons.py b/samudra/metrics/comparisons.py new file mode 100644 index 0000000000000000000000000000000000000000..61b61ece8460348b1ed6883006ad5eb40eb9f52a --- /dev/null +++ b/samudra/metrics/comparisons.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Model and observation fields, aligned and ready to reduce. + +Scoring a rollout and drawing it are the same pipeline up to the final +reduction: find the observed variable, derive the matching model field, restrict +both to the span and timestamps they share, and put the model on the +observation grid without breaking the quantities that are not linear. Only then +do the two jobs diverge -- one takes an area-weighted RMSE, the other draws a +map or a spectrum. + +That shared prefix lives here so `samudra.metrics.report` and +`samudra.viz.observations` cannot drift: a difference between a logged number +and a drawn one would have to come from the reduction, which is the part the +reader can see. + +Each builder returns a `Comparison`, which holds the model field on its own grid +as well as on the observation grid. Which one to reduce is not a detail: linear +differences may be interpolated first, but variance and eddy energy are +quadratic, and interpolating before reducing damps them by an amount set by the +resolution ratio rather than by the model. +""" + +from __future__ import annotations + +import dataclasses +import functools +import logging +from dataclasses import dataclass +from typing import Self, cast + +import numpy as np +import pandas as pd +import xarray as xr + +from samudra.constants import DatasetSpec +from samudra.metrics import kernels, observations + +logger = logging.getLogger(__name__) + +SPATIAL_DIMS = ("lat", "lon") + + +def ever_finite(field: xr.DataArray) -> xr.DataArray: + """A cell mask: true where the field carries data at any time.""" + finite = cast(xr.DataArray, np.isfinite(field)) + return finite.any("time") if "time" in field.dims else finite + + +def _without_time_mean(component: Comparison) -> Comparison: + """The model side of a comparison, as an anomaly about its own time mean.""" + return dataclasses.replace( + component, native=component.native - component.native.mean("time") + ) + + +@dataclass(frozen=True) +class Pairing: + """How much of the observation ocean survived pairing with a model. + + Linear interpolation makes an observation cell NaN whenever its model-grid + neighbours include land, so the comparison set is eroded by a band one model + cell wide -- all of it coastal, which is where model error is largest. The + eroded fraction therefore scales with model resolution: a coarse run is + scored on a smaller and easier subset of the ocean than a fine one. + Recording it makes that visible instead of leaving cross-resolution + comparison merely "not comparable" in prose. + """ + + n_paired_cells: int + n_observed_cells: int + + @property + def paired_ocean_fraction(self) -> float: + """The share of the product's ocean the model reaches.""" + if not self.n_observed_cells: + return float("nan") + return self.n_paired_cells / self.n_observed_cells + + def as_row(self) -> dict[str, float]: + """The metric-frame columns this contributes.""" + return { + "n_paired_cells": float(self.n_paired_cells), + "paired_ocean_fraction": self.paired_ocean_fraction, + } + + +@dataclass +class Comparison: + """One model field and one observation field, aligned and comparable. + + Both cover the same span and the same timestamps. `native` is on the + model's own grid and `on_obs_grid` is the interpolation of it, computed once + and reused. + """ + + label: str + native: xr.DataArray + obs: xr.DataArray + area: xr.DataArray + + @functools.cached_property + def on_obs_grid(self) -> xr.DataArray: + """The model field interpolated onto the observation grid.""" + return kernels.model_field_on_obs_grid(self.native, self.obs) + + @functools.cached_property + def error_squared(self) -> xr.DataArray: + """Squared model-minus-observation error, on the observation grid.""" + return (self.on_obs_grid - self.obs) ** 2 + + @property + def time(self) -> xr.DataArray: + return self.obs["time"] + + def slice(self, time_slice: slice) -> Self: + """Restrict the comparison to a time slice.""" + return dataclasses.replace( + self, + native=self.native.sel(time=time_slice), + obs=self.obs.sel(time=time_slice), + ) + + @property + def pairing(self) -> Pairing: + """How much of the observation ocean this model reaches.""" + obs_mask = ever_finite(self.obs) + return Pairing( + n_paired_cells=int((obs_mask & ever_finite(self.on_obs_grid)).sum()), + n_observed_cells=int(obs_mask.sum()), + ) + + +@dataclass +class LayerComparison(Comparison): + """An ocean-heat-content layer, with the caveat that qualifies it. + + `ohc_per_area_layer_maps` integrates whatever levels a cell has. Where two + products agree a column is shallow that is exact; where their bathymetry + *disagrees* the deeper one integrates water the other lacks, and the + difference reads as a heat deficit of roughly 2e9 J m^-2 per 50 m -- the + size of the whole 0-700 m score. + + `bathymetry_disagreement_m` measures that directly, as the area-weighted + metres of water the two sides differ by. Counting each side's incomplete + columns separately would only bound it: two products can be shallow in the + same places, or in different ones, and the fractions cannot tell those + apart. `model_partial_columns` is kept beside it to say which side is the + shallow one. + """ + + bathymetry_disagreement_m: float = float("nan") + model_partial_columns: float = float("nan") + + +@dataclass +class VelocityComparison: + """The two geostrophic velocity components, aligned together. + + Kept as a pair because the vector RMSE combines them before reducing over + time, and because the derived energies need both on the native grid. + """ + + eastward: Comparison + northward: Comparison + kind: str = "absolute" + + @property + def area(self) -> xr.DataArray: + return self.eastward.area + + def _rebased(self) -> Self: + """Put the model on the same footing as the observed velocity. + + DUACS reports either an absolute geostrophic velocity or an anomaly + about its own reference period, and the model's velocity derived from + `zos` is always absolute. Against the anomaly product the model's time + mean has to come off, or the two are not the same quantity. + + It cannot be left to a detrend further downstream: kinetic energy is + quadratic, so a retained mean flow contributes a constant *and* a cross + term with the eddies, and no linear detrend of the energy series + removes either. + """ + if self.kind != "anomaly": + return self + return dataclasses.replace( + self, + eastward=_without_time_mean(self.eastward), + northward=_without_time_mean(self.northward), + ) + + @property + def time(self) -> xr.DataArray: + return self.eastward.obs["time"] + + def slice(self, time_slice: slice) -> Self: + """Restrict both components to a time slice. + + Slicing has to happen here rather than on a derived quantity: the + energies reduce over time, so a mean taken before the slice is a mean of + the wrong record and cannot be corrected afterwards. + """ + # Rebased after slicing, not before: for the anomaly kind the mean has + # to be the one over the span actually being compared. + return dataclasses.replace( + self, + eastward=self.eastward.slice(time_slice), + northward=self.northward.slice(time_slice), + )._rebased() + + @functools.cached_property + def vector_error_squared(self) -> xr.DataArray: + """Squared vector error: a cell is penalised for direction, not just speed.""" + return self.eastward.error_squared + self.northward.error_squared + + def eddy_kinetic_energy(self) -> Comparison: + """Instantaneous EKE, reduced on the native grid before regridding.""" + return Comparison( + label=f"{self.eastward.label} EKE", + native=kernels.instantaneous_surface_eke( + self.eastward.native, self.northward.native + ), + obs=kernels.instantaneous_surface_eke( + self.eastward.obs, self.northward.obs + ), + area=self.area, + ) + + def mean_eddy_kinetic_energy(self) -> Comparison: + """Time-mean EKE map, for the figures rather than for scoring.""" + return Comparison( + label=f"{self.eastward.label} mean EKE", + native=kernels.mean_surface_eke( + self.eastward.native, self.northward.native + ), + obs=kernels.mean_surface_eke(self.eastward.obs, self.northward.obs), + area=self.area, + ) + + def kinetic_energy(self) -> Comparison: + """Surface KE, likewise squared on the native grid before regridding.""" + return Comparison( + label=f"{self.eastward.label} KE", + native=0.5 * (self.eastward.native**2 + self.northward.native**2), + obs=0.5 * (self.eastward.obs**2 + self.northward.obs**2), + area=self.area, + ) + + +def align( + model_field: xr.DataArray, + obs_field: xr.DataArray, + window: tuple[pd.Timestamp, pd.Timestamp] | None, + context: str, +) -> tuple[xr.DataArray, xr.DataArray]: + """Restrict both fields to their common span and require timestamps to match. + + The requested window is intersected with the span the two products actually + share, then exact agreement is required *within* that overlap. Products + legitimately end at different dates -- DUACS's last centered 5-day mean + predates OM4's final timestamp, for instance -- and demanding identical + spans would make the caller responsible for knowing every product's exact + coverage. Trimming the edges is safe; a disagreement anywhere inside the + overlap is still an error, which is the property that actually matters. + + A `None` window means the full shared span, which is what the + residual-variance and spectral diagnostics use: they characterise the whole + rollout rather than the primary scoring years. + """ + shared = ( + max( + pd.Timestamp(model_field["time"].values.min()), + pd.Timestamp(obs_field["time"].values.min()), + ), + min( + pd.Timestamp(model_field["time"].values.max()), + pd.Timestamp(obs_field["time"].values.max()), + ), + ) + if window is not None: + requested = window + window = (max(shared[0], window[0]), min(shared[1], window[1])) + if window != requested: + logger.info( + " %s: window trimmed to shared coverage %s..%s (requested %s..%s)", + context, + f"{window[0]:%Y-%m-%d}", + f"{window[1]:%Y-%m-%d}", + f"{requested[0]:%Y-%m-%d}", + f"{requested[1]:%Y-%m-%d}", + ) + else: + window = shared + + if window[0] > window[1]: + raise ValueError( + f"{context}: model and observation coverage do not overlap " + f"inside the requested window" + ) + + selection = slice(window[0], window[1]) + model_field = model_field.sel(time=selection) + obs_field = obs_field.sel(time=selection) + # Two empty axes compare equal, so the exact-match check below would pass + # and the emptiness would only surface much later. + if model_field.sizes.get("time", 0) == 0 or obs_field.sizes.get("time", 0) == 0: + raise ValueError( + f"{context}: no samples inside the requested window " + f"{window[0]:%Y-%m-%d} to {window[1]:%Y-%m-%d}" + ) + kernels.require_exact_time_match(model_field, obs_field, context) + return model_field, obs_field + + +def whole_years( + model_field: xr.DataArray, obs_field: xr.DataArray, context: str +) -> tuple[xr.DataArray, xr.DataArray]: + """Restrict a monthly pair to the calendar years both cover in full. + + Equal-year blocks are what make the score and its bootstrap comparable, so + a year missing a month has to go rather than be silently down-weighted. + Dropping it here -- where the shortfall is visible and attributable -- + beats failing later in the kernel, which can only report that some year is + ragged without knowing why. + """ + months = pd.DatetimeIndex(model_field["time"].values) + counts = pd.Series(1, index=months).groupby(months.year).sum() + whole = sorted(int(str(year)) for year, n in counts.items() if n == 12) + if not whole: + raise ValueError( + f"{context}: no calendar year is covered by twelve whole months " + f"(months available: {months.min():%Y-%m} to {months.max():%Y-%m})" + ) + dropped = sorted({int(str(year)) for year in counts.index} - set(whole)) + if dropped: + logger.info( + " %s: scoring %d-%d; dropped partly covered year(s) %s", + context, + whole[0], + whole[-1], + dropped, + ) + # By year membership, not a slice: a slice would only trim the ends and + # leave a ragged interior year in place. + keep = model_field["time"].dt.year.isin(whole) + return model_field.sel(time=keep), obs_field.sel(time=keep) + + +def surface_velocity( + rollout: xr.Dataset, + duacs: xr.Dataset, + window: tuple[pd.Timestamp, pd.Timestamp] | None, + context: str, + velocity_kind: str = "absolute", +) -> VelocityComparison: + """Surface geostrophic velocity, derived from the rollout's sea level. + + The model has velocities of its own, but DUACS reports geostrophic velocity + derived from altimetric sea level, so the comparable model field is derived + the same way. + """ + obs_u, obs_v = observations.duacs_velocity(duacs, velocity_kind) + sim_u, sim_v = kernels.geostrophic_velocity_from_zos( + rollout["zos"], lat_dim="lat", lon_dim="lon" + ) + sim_u, obs_u = align(sim_u, obs_u, window, f"{context} eastward velocity") + sim_v, obs_v = align(sim_v, obs_v, window, f"{context} northward velocity") + area = duacs["area"] + return VelocityComparison( + eastward=Comparison(f"{context} eastward velocity", sim_u, obs_u, area), + northward=Comparison(f"{context} northward velocity", sim_v, obs_v, area), + kind=velocity_kind, + )._rebased() + + +def sea_surface_temperature( + rollout: xr.Dataset, + oisst: xr.Dataset, + window: tuple[pd.Timestamp, pd.Timestamp] | None, + context: str, +) -> Comparison: + """The rollout's top model level against OISST.""" + obs = oisst[observations.find_var_name(oisst, observations.SST_ALIASES)] + sim, obs = align(rollout["thetao"].isel(lev=0), obs, window, f"{context} SST") + return Comparison(f"{context} SST", sim, obs, oisst["area"]) + + +def ohc_layer( + rollout: xr.Dataset, + argo: xr.Dataset, + layer: kernels.DepthLayer, + window: tuple[pd.Timestamp, pd.Timestamp] | None, + context: str, + model_dz: xr.DataArray, + complete_years_only: bool = True, +) -> LayerComparison: + """Per-area ocean heat content over a depth layer, against ARGO-IAP. + + ARGO-IAP is monthly, so both sides reduce to months and partially covered + months are dropped rather than compared against full ones. `complete_years_ + only` additionally trims to whole calendar years, which the equal-year block + reductions require; the diagnostics that are not scored that way pass False. + """ + obs_temp = argo[ + observations.find_var_name(argo, observations.ARGO_TEMPERATURE_ALIASES) + ] + obs_depth = observations.find_coord_name(obs_temp, observations.DEPTH_ALIASES) + assert obs_depth is not None # find_coord_name raises when required + + obs = kernels.monthly_mean_of_complete_months( + kernels.ohc_per_area_layer_maps(obs_temp, depth_name=obs_depth)[layer.label] + ) + sim = kernels.monthly_mean_of_complete_months( + kernels.ohc_per_area_layer_maps( + rollout["thetao"], native_dz=model_dz, depth_name="lev" + )[layer.label] + ) + label = f"{context} OHC {layer.label}" + sim, obs = align(sim, obs, window, label) + if complete_years_only: + # A rollout ending 24 December contributes no December once partial + # months are dropped, so scoring its final year would weigh 11 months + # against 12. + sim, obs = whole_years(sim, obs, label) + # Bathymetry does not change with time, so one step answers these and + # spares a pass over the whole record. + model_column = rollout["thetao"].isel(time=0) + observed_column = obs_temp.isel(time=0) + model_metres = kernels.integrated_layer_thickness( + model_column, layer, native_dz=model_dz, depth_name="lev" + ) + observed_metres = kernels.integrated_layer_thickness( + observed_column, layer, depth_name=obs_depth + ) + return LayerComparison( + label, + sim, + obs, + argo["area"], + bathymetry_disagreement_m=kernels.area_weighted_mean( + abs( + kernels.model_field_on_obs_grid(model_metres, observed_metres) + - observed_metres + ), + argo["area"], + ), + model_partial_columns=kernels.partial_column_fraction( + model_column, layer, native_dz=model_dz, depth_name="lev" + ), + ) + + +def model_depth_thickness(rollout: xr.Dataset, spec: DatasetSpec) -> xr.DataArray: + """Re-exported so callers need only this module to build a comparison.""" + return observations.model_depth_thickness(rollout, spec) diff --git a/samudra/metrics/kernels.py b/samudra/metrics/kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..86c5080d061233ea41d189a75f0cc6407ffc92a9 --- /dev/null +++ b/samudra/metrics/kernels.py @@ -0,0 +1,1076 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Pure metric kernels for observation-based evaluation. + +Implements the observation-evaluation metrics as specified in the metrics +document [1], at the revision incorporating its scientific review: instantaneous +(not time-mean) EKE, exact OHC layer-overlap weights, the exact-timestamp guard, +and calendar-year block-bootstrap uncertainty. + +[1]: https://docs.google.com/document/d/1cWzIrcajfW-gvoMzV3PU0JDAouVt5xerRmEiBIUbV1s + +Every function here is pure: it takes arrays and returns arrays or floats. No +file access, no plotting, no module-level mutable state, no `argparse`. That is +what lets `samudra.eval` and `samudra.viz` share them and get identical numbers. + +Two conventions run through the module: + +- `e_i(t)` is always *model minus observation* at cell `i` and time `t`. +- Whenever a spatial field collapses to one scalar it is area-weighted over the + cells where model, observation, and area are all finite. +""" + +from __future__ import annotations + +import dataclasses +import logging +from typing import cast + +import numpy as np +import pandas as pd +import xarray as xr + +logger = logging.getLogger(__name__) + +# Seawater properties used for ocean heat content. Held here rather than read +# from data so that OHC values are reproducible against externally published +# numbers, which depend on the choice. +RHO = 1035.0 # kg m^-3 +CP = 3850.0 # J kg^-1 K^-1 +OHC_REFERENCE = "temperature_integral_relative_to_0_degC" + +EARTH_RADIUS_M = 6.371e6 +GRAVITY = 9.81 +EARTH_ROTATION_RATE = 7.2921e-5 + +# Geostrophy is ill-conditioned near the equator, where the Coriolis parameter +# vanishes; DUACS comparisons drop this band entirely. +EQUATORIAL_MASK_DEG = 5.0 + +DEFAULT_BOOTSTRAP_SAMPLES = 10_000 + +# Least area-weighted fraction of a year's paired cells that must carry data +# before the year may contribute an equal-weight block. Set well below the +# coverage a healthy product gives, so it catches a genuinely starved year +# rather than ordinary gaps. +MIN_YEAR_COVERAGE = 0.5 + +# Fewest valid samples a cell needs before its residual variance means +# anything. Deseasonalizing removes a per-bin climatology, so a cell with about +# one sample per bin residualizes to exactly zero however variable it truly is. +# 24 admits two years of monthly data; the pentad path has 73 bins a year, so +# any cell that survives at all clears this comfortably. +MIN_RESIDUAL_SAMPLES = 24 + +# Fewest calendar-year blocks before a percentile bootstrap interval is worth +# reporting. With two blocks it degenerates to the min and max of two numbers. +MIN_BOOTSTRAP_BLOCKS = 5 + + +@dataclasses.dataclass(frozen=True) +class DepthLayer: + """A half-open depth interval `[min_depth, max_depth)` in metres.""" + + min_depth: float + max_depth: float + label: str + + +OHC_LAYERS: tuple[DepthLayer, ...] = ( + DepthLayer(0.0, 700.0, "Surface (0-700m)"), + DepthLayer(700.0, 2000.0, "Intermediate (700-2000m)"), +) + + +# -------------------------------------------------------------------------- +# Coordinate and time coercion +# -------------------------------------------------------------------------- + + +def coerce_datetime_values(values: np.ndarray) -> pd.DatetimeIndex: + """Convert a time coordinate to a `DatetimeIndex`, including cftime objects.""" + flat = np.asarray(values) + if np.issubdtype(flat.dtype, np.datetime64): + return pd.to_datetime(flat) + if flat.size == 0: + return pd.DatetimeIndex([]) + + first = flat.flat[0] + if hasattr(first, "year") and hasattr(first, "month") and hasattr(first, "day"): + return pd.to_datetime( + [ + f"{val.year:04d}-{val.month:02d}-{val.day:02d} " + f"{getattr(val, 'hour', 0):02d}:" + f"{getattr(val, 'minute', 0):02d}:" + f"{getattr(val, 'second', 0):02d}" + for val in flat + ] + ) + return pd.to_datetime(flat.astype(str)) + + +def normalize_lon[T: (xr.Dataset, xr.DataArray)](obj: T, lon_name: str = "lon") -> T: + """Put longitudes on `[0, 360)` and sort, if they are not already.""" + lon = obj[lon_name] + if np.nanmin(lon.values) < 0: + obj = obj.assign_coords({lon_name: (lon % 360)}).sortby(lon_name) + return obj + + +def require_exact_time_match( + model: xr.Dataset | xr.DataArray, + observation: xr.Dataset | xr.DataArray, + context: str, +) -> None: + """Fail rather than silently interpolate when model and observation times differ. + + Temporal interpolation between a model and an observation product smooths + model variability, which deflates variance metrics and inflates RMSE. So it + is disabled outright: timestamps must line up exactly before any horizontal + interpolation, and a mismatch is an error the caller has to resolve. + """ + if "time" not in model.dims or "time" not in observation.dims: + raise ValueError( + f"{context}: both model and observation fields must have a time dimension" + ) + + model_time = coerce_datetime_values(model["time"].values) + obs_time = coerce_datetime_values(observation["time"].values) + if model_time.equals(obs_time): + return + + common = model_time.intersection(obs_time) + model_only = model_time.difference(obs_time) + obs_only = obs_time.difference(model_time) + details = [ + f"model samples={len(model_time)}", + f"observation samples={len(obs_time)}", + f"common timestamps={len(common)}", + ] + if len(model_only): + details.append(f"first model-only timestamp={model_only[0]}") + if len(obs_only): + details.append(f"first observation-only timestamp={obs_only[0]}") + if not len(model_only) and not len(obs_only): + details.append("the timestamp order differs") + raise ValueError( + f"{context}: model and observation timestamps must match exactly before " + f"spatial interpolation ({'; '.join(details)}). Temporal interpolation is " + "disabled to avoid silently smoothing model variability." + ) + + +def median_time_step_days(index: pd.DatetimeIndex) -> float: + """Median spacing of a time index, in days; NaN when undefined.""" + if len(index) < 2: + return float("nan") + delta_days = np.diff(index.values).astype("timedelta64[s]").astype(float) / 86400.0 + delta_days = delta_days[np.isfinite(delta_days) & (delta_days > 0)] + return float(np.median(delta_days)) if delta_days.size else float("nan") + + +def complete_calendar_years( + time: xr.DataArray, max_gap_steps: float = 2.5, min_sample_ratio: float = 0.75 +) -> list[int]: + """Years that span the full calendar with no unexplained hole in the middle. + + Completeness is judged by the largest gap between consecutive samples, not + by a sample count. A count tolerance has to be loose enough to absorb the + cadence variation a real product has -- OM4's 5-day calendar carries + occasional 6-day steps -- and that same looseness then admits a fortnight + of genuinely missing data. A gap threshold separates the two: it accepts a + 6-day step and rejects a 15-day hole, which is what "complete" is meant to + mean. + + Args: + time: The time coordinate to inspect. + max_gap_steps: Largest interior gap to tolerate, in multiples of the + median step. + min_sample_ratio: Fewest samples a year may hold, as a fraction of the + busiest year in the record. + """ + index = pd.DatetimeIndex(time.values) + if index.empty: + return [] + step_days = median_time_step_days(index) + if not np.isfinite(step_days) or step_days <= 0: + return [] + # The year boundary will not land on a sample, so allow the record to start + # and end up to one-and-a-half steps inside the calendar year. + boundary_tolerance = pd.Timedelta(days=1.5 * step_days) + max_gap = pd.Timedelta(days=max_gap_steps * step_days) + + # A gap threshold alone cannot see thinning: dropping every second sample + # leaves gaps of exactly two median steps. Check the count as well. + counts = pd.Series(1, index=index).groupby(index.year).sum() + busiest = int(counts.max()) + + years = [] + for year in sorted({int(y) for y in index.year}): + year_index = index[index.year == year].sort_values() + if year_index.size < 2: + continue + if year_index.size < min_sample_ratio * busiest: + continue + covers_calendar = ( + year_index.min() <= pd.Timestamp(f"{year}-01-01") + boundary_tolerance + and year_index.max() + >= pd.Timestamp(f"{year}-12-31 23:59:59.999999999") - boundary_tolerance + ) + largest_gap = pd.Timedelta(np.diff(year_index.values).max()) + if covers_calendar and largest_gap <= max_gap: + years.append(year) + return years + + +def require_complete_calendar_years(time: xr.DataArray, context: str) -> list[int]: + """Return the covered years, failing if the period contains a partial one. + + Equal-year blocks are what make the primary score and its bootstrap + comparable, so a ragged final year has to be an error rather than a silently + down-weighted block. + """ + index = pd.DatetimeIndex(time.values) + present_years = sorted({int(year) for year in index.year}) + complete = complete_calendar_years(time) + if complete != present_years: + incomplete = sorted(set(present_years) - set(complete)) + # This fires only after the rollout, so the message has to carry the fix + # rather than leave the reader to work out which bound is wrong. The + # usual cause is an observation product ending days short of the year: + # the shipped window clears the tolerance by about half a day, so one + # extra 6-day step in a rebuilt archive is enough to trip it. + usable = [year for year in complete if year not in incomplete] + suggestion = ( + f" The last fully covered year is {max(usable)}, so " + f'rmse_end="{max(usable)}-12-31" would score {len(usable)} years.' + if usable + else "" + ) + raise ValueError( + f"{context}: uncertainty-enabled primary RMSE requires complete calendar " + f"years, but {incomplete} {'is' if len(incomplete) == 1 else 'are'} only " + f"partly covered (data ends {index.max():%Y-%m-%d}).{suggestion}" + ) + return complete + + +# -------------------------------------------------------------------------- +# Depth geometry +# -------------------------------------------------------------------------- + + +def depth_edges_from_centers(depth_vals: np.ndarray) -> np.ndarray: + """Reconstruct cell edges from depth centers, for products that ship only centers.""" + depth = np.asarray(depth_vals, dtype=float) + if depth.ndim != 1 or depth.size < 2: + raise ValueError("Expected at least two 1D depth centers") + if not np.all(np.isfinite(depth)) or np.any(np.diff(depth) <= 0): + raise ValueError("Depth centers must be finite and strictly increasing") + edges = np.empty(depth.size + 1, dtype=float) + edges[1:-1] = 0.5 * (depth[:-1] + depth[1:]) + edges[0] = max(0.0, depth[0] - 0.5 * (depth[1] - depth[0])) + edges[-1] = depth[-1] + 0.5 * (depth[-1] - depth[-2]) + return edges + + +def layer_overlap_thickness( + depth: xr.DataArray, + min_depth: float, + max_depth: float, + native_dz: xr.DataArray | None = None, +) -> xr.DataArray: + """Thickness of each native cell inside the half-open interval `[min_depth, max_depth)`. + + Boundary cells are prorated by their true overlap rather than included or + excluded on the basis of their center depth. Selecting by center is what + made a nominal "0-700 m" integral actually cover 0-650 m on the OM4 vertical + grid, where a cell centered at 775 m spans 650-900 m. + + Observation products supply only depth centers, so their edges are + reconstructed from adjacent centers. Model products additionally supply + exact native thicknesses; pass those as `native_dz` to recover exact bounds. + """ + min_depth = float(min_depth) + max_depth = float(max_depth) + if ( + not np.isfinite(min_depth) + or not np.isfinite(max_depth) + or max_depth <= min_depth + ): + raise ValueError(f"Invalid target depth interval [{min_depth}, {max_depth})") + + if native_dz is None: + edges = depth_edges_from_centers(depth.values) + lower = xr.DataArray(edges[:-1], dims=depth.dims, coords=depth.coords) + upper = xr.DataArray(edges[1:], dims=depth.dims, coords=depth.coords) + else: + if depth.ndim != 1: + raise ValueError( + f"Expected a 1D native depth coordinate, got dimensions {depth.dims}" + ) + # Accumulate from the surface rather than assuming centers are cell + # midpoints. A uniform edge shift would preserve the summed thickness + # while getting every boundary weight wrong. + thickness = np.asarray(native_dz.values, dtype=float) + edges = np.concatenate([[0.0], np.cumsum(thickness)]) + lower = xr.DataArray(edges[:-1], dims=depth.dims, coords=depth.coords) + upper = xr.DataArray(edges[1:], dims=depth.dims, coords=depth.coords) + + clipped_lower = xr.where(lower > min_depth, lower, min_depth) + clipped_upper = xr.where(upper < max_depth, upper, max_depth) + overlap = (clipped_upper - clipped_lower).clip(min=0.0) + overlap.name = "layer_overlap_dz" + return overlap + + +# -------------------------------------------------------------------------- +# Area-weighted reductions +# -------------------------------------------------------------------------- + + +def area_weighted_mean(field: xr.DataArray, area: xr.DataArray) -> float: + """Area-weighted mean over the finite cells of `field`.""" + mask = np.isfinite(field) + weighted_sum = (field.where(mask) * area).sum() + weight_sum = area.where(mask).sum() + if float(weight_sum.values) == 0.0: + return float("nan") + return float((weighted_sum / weight_sum).values) + + +def area_weighted_map_rmse(bias: xr.DataArray, area: xr.DataArray) -> float: + """Area-weighted RMS of a bias or per-cell RMSE map.""" + mean_square = area_weighted_mean(bias**2, area) + if not np.isfinite(mean_square): + return float("nan") + return float(np.sqrt(mean_square)) + + +def area_weighted_timeseries( + field: xr.DataArray, + area: xr.DataArray, + spatial_dims: tuple[str, str], +) -> xr.DataArray: + """Collapse a field's spatial dims to an area-weighted time series.""" + mask = np.isfinite(field) + weighted_sum = (field.where(mask) * area).sum(spatial_dims, skipna=True) + weight_sum = area.where(mask).sum(spatial_dims, skipna=True) + return weighted_sum / weight_sum + + +def area_weighted_pattern_corr( + model: xr.DataArray, + observation: xr.DataArray, + area: xr.DataArray, +) -> float: + """Centered pattern correlation on the common finite, positive-area mask.""" + model, observation, area = xr.align(model, observation, area, join="inner") + valid = ( + np.isfinite(model) & np.isfinite(observation) & np.isfinite(area) & (area > 0) + ) + weights = area.where(valid) + weight_sum = weights.sum(skipna=True) + weight_sum_value = float(weight_sum.values) + if not np.isfinite(weight_sum_value) or weight_sum_value <= 0.0: + return float("nan") + + model_mean = (model.where(valid) * weights).sum(skipna=True) / weight_sum + obs_mean = (observation.where(valid) * weights).sum(skipna=True) / weight_sum + model_anomaly = model - model_mean + obs_anomaly = observation - obs_mean + covariance = (weights * model_anomaly * obs_anomaly).sum(skipna=True) + model_variance = (weights * model_anomaly**2).sum(skipna=True) + obs_variance = (weights * obs_anomaly**2).sum(skipna=True) + denominator = cast(xr.DataArray, np.sqrt(model_variance * obs_variance)) + denominator_value = float(denominator.values) + if not np.isfinite(denominator_value) or denominator_value <= 0.0: + return float("nan") + return float((covariance / denominator).values) + + +def monthly_mean_of_complete_months( + field: xr.DataArray, time_dim: str = "time" +) -> xr.DataArray: + """Monthly mean, keeping only months the source actually covers end to end. + + The month at either edge of a record is usually partial. A rollout that + stops on the 24th contributes 24 days to December while a monthly + observation product contributes all 31, yet both reduce to a single + December label -- so the comparison silently pits different quantities + against each other, and every downstream check counts it as one good month. + + A month is kept when it holds at least as many samples as its length allows + at the source's own cadence, which admits every month of a monthly product + and rejects only genuinely short ones in a finer series. + """ + index = pd.DatetimeIndex(field[time_dim].values) + monthly = field.resample({time_dim: "MS"}).mean() + step = median_time_step_days(index) + if not np.isfinite(step) or step <= 0: + return monthly + + counts = pd.Series(1, index=index).resample("MS").sum() + keep = [] + for label in pd.DatetimeIndex(monthly[time_dim].values): + days = label.days_in_month + expected = max(1, int(days // step)) + keep.append(int(counts.get(label, 0)) >= expected) + + kept = np.array(keep, dtype=bool) + if not kept.all(): + dropped = [ + f"{stamp:%Y-%m}" + for stamp, ok in zip(pd.DatetimeIndex(monthly[time_dim].values), kept) + if not ok + ] + logger.info(" dropping partially covered month(s): %s", ", ".join(dropped)) + return monthly.isel({time_dim: kept}) + + +# -------------------------------------------------------------------------- +# Velocity, EKE +# -------------------------------------------------------------------------- + + +def apply_equatorial_mask( + da: xr.DataArray, + lat_dim: str = "lat", + equatorial_mask_deg: float = EQUATORIAL_MASK_DEG, +) -> xr.DataArray: + """Blank the equatorial band, where geostrophic balance does not hold.""" + lat_vals = np.asarray(da[lat_dim].values, dtype=float) + eq_mask = np.abs(lat_vals) >= equatorial_mask_deg + eq_mask_da = xr.DataArray(eq_mask, dims=[lat_dim], coords={lat_dim: da[lat_dim]}) + return da.where(eq_mask_da) + + +def _spans_globe(lon_values: np.ndarray) -> bool: + """Whether a longitude axis wraps the globe, so its ends are neighbours.""" + values = np.asarray(lon_values, dtype=float) + if values.size < 2: + return False + spacing = float(np.median(np.diff(values))) + if not np.isfinite(spacing) or spacing <= 0: + return False + # Tight tolerance: one whole cell is the width at which "global" and "one + # column short" stop being distinguishable. + return bool( + np.isclose(values[-1] - values[0] + spacing, 360.0, atol=1e-6 * abs(spacing)) + ) + + +def _differentiate_lon(field: xr.DataArray, lon_dim: str) -> xr.DataArray: + """Longitude derivative, wrapping across the 0/360 seam on a global grid. + + `xarray.differentiate` falls back to a one-sided difference at the first and + last index, which on a global grid means the two meridians beside the seam + are differenced against the wrong neighbour rather than across it. That + error propagates into the geostrophic velocity and from there into both the + velocity RMSE and EKE. Padding one column from the opposite edge makes the + seam an interior point. + + Regional grids are differentiated as-is: their edges really are edges. + """ + values = np.asarray(field[lon_dim].values, dtype=float) + if not _spans_globe(values): + return field.differentiate(lon_dim) + + padded = _wrap_lon(field, lon_dim) + # Concatenating single-column pads leaves 1-wide dask chunks at each edge, + # and `np.gradient` needs at least two points per chunk. Rechunking the + # padded axis into one block is safe here: it is a single row of longitudes, + # and the alternative -- a partially rechunked axis -- would reintroduce the + # same failure at whatever new boundary it created. + if padded.chunks is not None: + padded = padded.chunk({lon_dim: -1}) + return padded.differentiate(lon_dim).isel({lon_dim: slice(1, -1)}) + + +def geostrophic_velocity_from_zos( + zos: xr.DataArray, + lat_dim: str = "y", + lon_dim: str = "x", + equatorial_mask_deg: float = EQUATORIAL_MASK_DEG, +) -> tuple[xr.DataArray, xr.DataArray]: + """Surface geostrophic velocity from sea-surface height. + + This is what makes the model comparable to DUACS, which reports geostrophic + velocity derived from altimetric sea level rather than a model velocity. + """ + deg_per_m = 180.0 / (np.pi * EARTH_RADIUS_M) + + lat_da = zos[lat_dim] + f = 2.0 * EARTH_ROTATION_RATE * np.sin(np.deg2rad(lat_da)) + coslat = np.cos(np.deg2rad(lat_da)) + + deta_dy = zos.differentiate(lat_dim) * deg_per_m + deta_dx = _differentiate_lon(zos, lon_dim) * deg_per_m / coslat + + # `f` vanishes at the equator and `coslat` at the poles, so both quotients + # blow up there; the equatorial mask below removes the band that matters and + # the polar rows fall outside every observation product's coverage. + with np.errstate(divide="ignore", invalid="ignore"): + u_g = -GRAVITY / f * deta_dy + v_g = GRAVITY / f * deta_dx + + u_g = apply_equatorial_mask(u_g, lat_dim, equatorial_mask_deg) + v_g = apply_equatorial_mask(v_g, lat_dim, equatorial_mask_deg) + return u_g, v_g + + +def instantaneous_surface_eke(u: xr.DataArray, v: xr.DataArray) -> xr.DataArray: + """Eddy kinetic energy at each timestamp, from velocity anomalies about the time mean. + + Instantaneous, not the EKE of the time-mean field: scoring a time-mean EKE + map against observations measures a bias, not an RMSE, and is not comparable + with the other per-cell time-RMSE metrics. + """ + u_anom = u - u.mean("time") + v_anom = v - v.mean("time") + return 0.5 * (u_anom**2 + v_anom**2) + + +def mean_surface_eke(u: xr.DataArray, v: xr.DataArray) -> xr.DataArray: + """Time-mean EKE map, for the map figures rather than for scoring.""" + return instantaneous_surface_eke(u, v).mean("time", skipna=True) + + +# -------------------------------------------------------------------------- +# Ocean heat content +# -------------------------------------------------------------------------- + + +def integrated_layer_thickness( + temp_field: xr.DataArray, + layer: DepthLayer, + native_dz: xr.DataArray | None = None, + depth_name: str = "depth", +) -> xr.DataArray: + """Metres of water each cell actually contributes to a layer. + + `ohc_per_area_layer_maps` integrates whatever levels a cell has, so this is + the weight behind that integral: the layer's nominal thickness where the + column is complete, less wherever it stops early. Differencing it between + two products is what turns "their bathymetry disagrees" into a number, and + metres of water convert to the heat the deeper product counted and the + other did not. + + Thickness is linear, so unlike the heat content itself this may be + interpolated onto another grid before differencing. + """ + overlap_dz = layer_overlap_thickness( + temp_field[depth_name], layer.min_depth, layer.max_depth, native_dz=native_dz + ) + finite = cast(xr.DataArray, np.isfinite(temp_field)) + if "time" in finite.dims: + finite = finite.any("time") + return (overlap_dz * finite).sum(depth_name) + + +def partial_column_fraction( + temp_field: xr.DataArray, + layer: DepthLayer, + native_dz: xr.DataArray | None = None, + depth_name: str = "depth", +) -> float: + """Fraction of wet cells whose column stops before the layer bottom. + + `ohc_per_area_layer_maps` integrates whatever levels a cell has, so where + two products disagree about bathymetry the missing water reads as a + deficit. This measures how much of the comparison is exposed to that, + rather than leaving the risk as an assertion. + """ + overlap_dz = layer_overlap_thickness( + temp_field[depth_name], layer.min_depth, layer.max_depth, native_dz=native_dz + ) + wanted = overlap_dz > 0 + finite = cast(xr.DataArray, np.isfinite(temp_field.where(wanted))) + present = finite.sum(depth_name) + if "time" in present.dims: + present = present.max("time") + wet = present > 0 + if not bool(wet.any()): + return float("nan") + complete = present >= int(wanted.sum()) + return float(1.0 - (complete & wet).sum() / wet.sum()) + + +def ohc_per_area_layer_maps( + temp_field: xr.DataArray, + layers: tuple[DepthLayer, ...] = OHC_LAYERS, + native_dz: xr.DataArray | None = None, + depth_name: str = "depth", +) -> dict[str, xr.DataArray]: + """Ocean heat content per unit area, in J m^-2, for each depth layer. + + Temperature is integrated relative to 0 degrees C (see `OHC_REFERENCE`), so + these are not IAP anomalies; comparing against published OHC numbers means + matching both the reference and the volumetric heat capacity. + + A cell contributes if it has water anywhere in the layer, and levels it + lacks are skipped. Where two products agree a column is shallow this is + exact, but where their bathymetry *disagrees* the deeper product integrates + water the other lacks and the difference reads as a heat deficit -- roughly + 2e9 J m^-2 per 50 m. `partial_column_fraction` measures how much of a + comparison is exposed to this, and `report` emits it beside every layer + score. Removing the effect, rather than measuring it, would require masking + each cell to the depth range both products share. + """ + out = {} + for layer in layers: + overlap_dz = layer_overlap_thickness( + temp_field[depth_name], + layer.min_depth, + layer.max_depth, + native_dz=native_dz, + ) + # An axis stopping short of the layer would integrate less water than + # the other product, giving a difference dominated by missing water + # rather than model error. `thermo_dynamic_5` reaches only 65 m. + covered = float(overlap_dz.sum()) + nominal = layer.max_depth - layer.min_depth + if not np.isclose(covered, nominal, rtol=1e-6): + raise ValueError( + f"Depth axis {depth_name!r} covers {covered:.1f} m of the " + f"{nominal:.1f} m layer [{layer.min_depth:g}, {layer.max_depth:g}) " + f"('{layer.label}'). OHC over a partial column is not comparable " + "against a product that spans the full layer; use a variable set " + "that reaches the layer bottom, or disable the OHC metrics." + ) + temp_layer = temp_field.where(overlap_dz > 0) + # np.isfinite on a DataArray returns a DataArray, but the numpy stubs + # widen it to ndarray, which has no `dim=` keyword. + finite = cast(xr.DataArray, np.isfinite(temp_layer)) + # A cell contributes whatever water it has; `skipna` treats levels it + # lacks as zero heat. See the docstring for when that matters. + valid = finite.any(dim=depth_name) + ohc = ( + (temp_layer * overlap_dz * RHO * CP) + .sum(depth_name, skipna=True) + .where(valid) + ) + out[layer.label] = ohc + return out + + +# -------------------------------------------------------------------------- +# Anomalies: seasonal cycle, trend, residual variance +# -------------------------------------------------------------------------- + + +def pentad_index(index: pd.DatetimeIndex) -> np.ndarray: + """Map timestamps to 5-day climatological bins on a 365-day calendar. + + Using a no-leap calendar keeps post-February dates in the same bin across + leap and non-leap years; February 29 joins the February 28 bin. + """ + dayofyear = index.dayofyear.to_numpy().copy() + is_leap = np.asarray(index.is_leap_year) + month = index.month.to_numpy() + day = index.day.to_numpy() + dayofyear[is_leap & (month > 2)] -= 1 + dayofyear[is_leap & (month == 2) & (day == 29)] = 59 + return ((dayofyear - 1) // 5 + 1).astype(int) + + +def has_repeated_calendar_months(index: pd.DatetimeIndex) -> bool: + """True when every calendar month appears at least twice.""" + counts = pd.Series(index.month).value_counts() + return bool(len(counts) == 12 and (counts >= 2).all()) + + +def _use_pentad_climatology(index: pd.DatetimeIndex) -> bool: + """Whether the sampling is fine enough, and repeats enough, for pentad bins.""" + if ( + not np.isfinite(median_time_step_days(index)) + or median_time_step_days(index) > 10 + ): + return False + counts = pd.Series(pentad_index(index)).value_counts() + return bool((counts >= 2).sum() >= 24) + + +def detrend_linear_dataarray( + field: xr.DataArray, time_dim: str = "time" +) -> xr.DataArray: + """Remove a least-squares linear trend independently at each grid cell.""" + if field.sizes.get(time_dim, 0) < 2: + return field - field.mean(time_dim, skipna=True) + + # Elapsed days, not sample index: the axis is not uniform once + # `monthly_mean_of_complete_months` drops a month, and an index regression + # would treat the gap as though no time had passed. + stamps = pd.DatetimeIndex(field[time_dim].values) + elapsed = (stamps - stamps[0]).days.to_numpy(dtype=float) + x = xr.DataArray(elapsed, dims=[time_dim], coords={time_dim: field[time_dim]}) + finite = np.isfinite(field) + x_mean = x.where(finite).mean(time_dim, skipna=True) + y_mean = field.mean(time_dim, skipna=True) + cov = ((x - x_mean) * (field - y_mean)).where(finite).sum(time_dim, skipna=True) + var = ((x - x_mean) ** 2).where(finite).sum(time_dim, skipna=True) + slope = xr.where(var > 0, cov / var, 0.0) + intercept = y_mean - slope * x_mean + return field - (slope * x + intercept) + + +def without_seasonal_cycle(field: xr.DataArray, time_dim: str = "time") -> xr.DataArray: + """Subtract a climatology, leaving the anomaly, per grid cell. + + Bins into pentads where the sampling supports it, then calendar months, + then falls back to the record mean. Pentads rather than literal calendar + dates: OM4 puts 73 samples in every calendar year, so a leap year carries + one six-day step and every later sample lands on a date no common year + visits. Grouping on the date itself would give those samples a climatology + made of themselves, and an anomaly of zero. + + This is the field counterpart of `series_without_seasonal_cycle`, and the + deseasonalizing half of `detrended_deseasonalized`. + """ + if field.sizes.get(time_dim, 0) == 0: + return field + time_index = pd.DatetimeIndex(field[time_dim].values) + + if _use_pentad_climatology(time_index): + pentad = xr.DataArray( + pentad_index(time_index), + dims=[time_dim], + coords={time_dim: field[time_dim]}, + name="pentad", + ) + with_pentad = field.assign_coords(pentad=pentad) + anomaly = with_pentad.groupby("pentad") - with_pentad.groupby("pentad").mean( + time_dim, skipna=True + ) + return anomaly.drop_vars("pentad") if "pentad" in anomaly.coords else anomaly + if has_repeated_calendar_months(time_index): + return field.groupby(f"{time_dim}.month") - field.groupby( + f"{time_dim}.month" + ).mean(time_dim, skipna=True) + return field - field.mean(time_dim, skipna=True) + + +def detrended_deseasonalized( + field: xr.DataArray, time_dim: str = "time" +) -> xr.DataArray: + """Remove the seasonal cycle and then a linear trend, per grid cell. + + The residual is the quantity of interest: it isolates the variability a + model either reproduces or smooths away, with the trend and seasonal cycle, + which are easy to get right, taken out first. + """ + return detrend_linear_dataarray( + without_seasonal_cycle(field, time_dim=time_dim), time_dim=time_dim + ) + + +def residual_variance_map( + field: xr.DataArray, time_dim: str = "time", min_samples: int = MIN_RESIDUAL_SAMPLES +) -> xr.DataArray: + """Variance of the detrended, deseasonalized residual at each grid cell. + + Population variance (`ddof=0`), matching the reference implementation. + + Cells with too few valid samples return NaN rather than a number. Removing + a per-bin climatology from a cell that holds roughly one sample per bin + leaves exactly zero by construction, and detrending cannot raise it -- so + an undersampled cell would otherwise report "perfectly still ocean", + indistinguishable from a genuinely quiescent point. That zero then enters + the map RMSE and the pattern correlation as though the model's real + variability there were pure error. + """ + residual = detrended_deseasonalized(field, time_dim=time_dim) + variance = residual.var(time_dim, skipna=True, ddof=0) + valid = cast(xr.DataArray, np.isfinite(field)).sum(time_dim) + return variance.where(valid >= min_samples) + + +def series_without_seasonal_cycle(series: pd.Series) -> pd.Series: + """Remove a pentad or monthly climatology, falling back to the mean.""" + values = series.astype(float) + index = pd.DatetimeIndex(values.index) + if _use_pentad_climatology(index): + pentad = pd.Series(pentad_index(index), index=values.index) + climatology = values.groupby(pentad).mean() + return values - pentad.map(climatology) + if has_repeated_calendar_months(index): + climatology = values.groupby(index.month).mean() + seasonal = pd.Series(index.month, index=values.index).map(climatology) + return values - seasonal + return values - values.mean() + + +def series_without_linear_trend(series: pd.Series) -> pd.Series: + """Remove a least-squares linear trend from a scalar time series.""" + values = series.astype(float) + raw = np.asarray(values.to_numpy(), dtype=float) + finite = np.isfinite(raw) + if finite.sum() < 2: + return values - values.mean() + + # Elapsed days, matching `detrend_linear_dataarray`, so the scalar and map + # paths return the same residual for the same series. + stamps = pd.DatetimeIndex(values.index) + x = (stamps - stamps[0]).days.to_numpy(dtype=float) + slope, intercept = np.polyfit(x[finite], raw[finite], 1) + trend = pd.Series(slope * x + intercept, index=values.index) + return values - trend + + +def series_linear_trend_per_year(series: pd.Series) -> float: + """Ordinary-least-squares slope in units per year.""" + values = series.astype(float) + raw = np.asarray(values.to_numpy(), dtype=float) + finite = np.isfinite(raw) + if values.empty or finite.sum() < 2: + return float("nan") + index = pd.DatetimeIndex(values.index) + x = (index - index[0]).days.to_numpy(dtype=float) / 365.25 + slope, _ = np.polyfit(x[finite], raw[finite], 1) + return float(slope) + + +def series_residual_variance(series: pd.Series) -> float: + """Population variance of the detrended, deseasonalized scalar series.""" + residual = series_without_linear_trend(series_without_seasonal_cycle(series)) + values = np.asarray(residual.values, dtype=float) + finite = values[np.isfinite(values)] + if finite.size == 0: + return float("nan") + return float(np.var(finite)) + + +# -------------------------------------------------------------------------- +# Horizontal regridding +# -------------------------------------------------------------------------- + + +def _wrap_lon(field: xr.DataArray, lon_dim: str = "lon") -> xr.DataArray: + """Pad a global longitude axis with its wrapped neighbours, or pass through. + + Regional grids are returned untouched: their edges are real boundaries, and + inventing data beyond them would be worse than dropping a column. + """ + values = np.asarray(field[lon_dim].values, dtype=float) + if not _spans_globe(values): + return field + left = field.isel({lon_dim: [-1]}).assign_coords({lon_dim: [values[-1] - 360.0]}) + right = field.isel({lon_dim: [0]}).assign_coords({lon_dim: [values[0] + 360.0]}) + return xr.concat([left, field, right], dim=lon_dim) + + +def _clamped_to(model_lat: np.ndarray, target_lat: xr.DataArray) -> xr.DataArray: + """Target latitudes pulled inside the model's range, keeping their labels. + + Interpolation leaves NaN wherever the target sits outside the model's + convex hull. Unlike longitude there is nothing to wrap to at a pole, so the + choice is between dropping those rows -- by an amount that depends on model + resolution -- and carrying the nearest model row outward. The gap is at + most half a model cell, and dropping them silently changes which ocean the + metric describes. + """ + clamped = np.clip( + np.asarray(target_lat.values, dtype=float), + float(model_lat.min()), + float(model_lat.max()), + ) + return xr.DataArray(clamped, dims=target_lat.dims, coords=target_lat.coords) + + +def model_field_on_obs_grid( + model_field: xr.DataArray, + obs_field: xr.DataArray, +) -> xr.DataArray: + """Interpolate a model field onto an observation product's horizontal grid. + + Horizontal only. Callers must have already established that the time axes + match exactly (see `require_exact_time_match`); nothing here resamples time. + """ + rename = {} + if "y" in model_field.dims: + rename["y"] = "lat" + if "x" in model_field.dims: + rename["x"] = "lon" + if rename: + model_field = model_field.rename(rename) + model_field = normalize_lon(model_field, "lon").sortby("lat").sortby("lon") + # A global grid has no true longitude edge, but its outermost centers still + # sit inside the observation product's, so plain interpolation would drop + # the seam columns. Wrapping one column from each end makes them interior. + model_field = _wrap_lon(model_field, "lon") + # Latitude cannot wrap, so clamp the target instead: a 1 degree model tops + # out at +-89.5 where OISST reaches +-89.875, and those rows would otherwise + # be dropped by an amount that depends on model resolution. + lat_target = _clamped_to( + np.asarray(model_field["lat"].values, dtype=float), obs_field["lat"] + ) + interpolated = model_field.interp( + lat=lat_target, + lon=obs_field["lon"], + kwargs={"fill_value": np.nan}, + ) + # `interp` labels the result with the clamped values; restore the real ones. + interpolated = interpolated.assign_coords(lat=obs_field["lat"]) + return interpolated.where(np.isfinite(obs_field)) + + +# -------------------------------------------------------------------------- +# Uncertainty +# -------------------------------------------------------------------------- + + +def interannual_rmse_summary( + annual_rmse: np.ndarray, + bootstrap_samples: int = DEFAULT_BOOTSTRAP_SAMPLES, +) -> dict[str, float | int | str]: + """Summarize a set of per-year RMSE values with spread and a bootstrap CI. + + The point estimate is the root of the equal-year mean annual MSE, so each + calendar year carries the same weight regardless of how many samples it + holds. Resampling whole years (rather than individual timestamps) respects + the strong autocorrelation within a year. + """ + values = np.asarray(annual_rmse, dtype=float) + values = values[np.isfinite(values)] + if values.size == 0: + return { + "annual_std": np.nan, + "ci_low": np.nan, + "ci_high": np.nan, + "n_years": 0, + "n_bootstrap": 0, + "uncertainty_method": "unavailable", + "block_aggregate_rmse": np.nan, + "primary_aggregation_method": "root_mean_annual_mse_equal_calendar_years", + } + + annual_mse = values**2 + block_aggregate_rmse = float(np.sqrt(annual_mse.mean())) + annual_std = float(np.std(values, ddof=1)) if values.size > 1 else np.nan + # With two blocks a percentile bootstrap returns exactly [min, max] of two + # numbers. Report no interval rather than one that overstates itself. + if values.size < MIN_BOOTSTRAP_BLOCKS or bootstrap_samples <= 0: + ci_low = ci_high = np.nan + n_bootstrap = 0 + else: + # Fixed seed so identical data gives an identical interval. + rng = np.random.default_rng(0) + sampled = rng.choice( + annual_mse, size=(bootstrap_samples, values.size), replace=True + ) + bootstrap_rmse = np.sqrt(sampled.mean(axis=1)) + ci_low, ci_high = np.percentile(bootstrap_rmse, [2.5, 97.5]).astype(float) + n_bootstrap = int(bootstrap_samples) + + return { + "annual_std": annual_std, + "ci_low": float(ci_low), + "ci_high": float(ci_high), + "n_years": int(values.size), + "n_bootstrap": n_bootstrap, + "uncertainty_method": ( + "calendar_year_block_bootstrap" if n_bootstrap else "insufficient_blocks" + ), + "block_aggregate_rmse": block_aggregate_rmse, + "primary_aggregation_method": "root_mean_annual_mse_equal_calendar_years", + } + + +def rmse_map_with_uncertainty( + error_squared: xr.DataArray, + area: xr.DataArray, + spatial_dims: tuple[str, str], + context: str, + bootstrap_samples: int = DEFAULT_BOOTSTRAP_SAMPLES, + min_year_coverage: float = MIN_YEAR_COVERAGE, +) -> tuple[xr.DataArray, dict[int, float], dict[str, float | int | str]]: + """Per-cell RMSE map, per-year total RMSE, and the uncertainty summary. + + Args: + error_squared: Squared model-minus-observation error, with a time dim. + area: Grid-cell area on the same horizontal grid. + spatial_dims: The two horizontal dim names to reduce over. + context: Human-readable label used in error messages. + bootstrap_samples: Block-bootstrap draws; 0 disables the interval. + min_year_coverage: Least area-weighted fraction of a year's paired + cells that must carry data for that year to contribute a block. + + Returns: + The RMSE map, a mapping of year to that year's area-weighted total + RMSE, and the summary produced by `interannual_rmse_summary`. + """ + if error_squared.sizes.get("time", 0) == 0: + raise ValueError(f"No time samples available for {context}") + + years = require_complete_calendar_years(error_squared["time"], context) + annual_mse_maps = ( + error_squared.groupby("time.year").mean("time", skipna=True).sel(year=years) + ) + # Coverage is measured on the data, not the time axis: a year can keep + # every timestamp and still be mostly NaN, and would then carry a full + # equal-weight block. + finite_samples = ( + cast(xr.DataArray, np.isfinite(error_squared)) + .groupby("time.year") + .sum() + .sel(year=years) + ) + # Denominator is the cells that are ever comparable. Land and the blanked + # equatorial band are permanently NaN, so counting them would score a + # flawless year at roughly the ocean fraction of the planet. + comparable = cast(xr.DataArray, np.isfinite(error_squared)).any("time") + stamps_per_year = ( + xr.ones_like(error_squared.isel({d: 0 for d in spatial_dims}, drop=True)) + .groupby("time.year") + .sum() + .sel(year=years) + ) + total_samples = comparable * stamps_per_year + # The map, the scalar score, and the bootstrap all rest on the same + # equal-year blocks, so the headline number is the one the interval covers. + rmse_map = np.sqrt(annual_mse_maps.mean("year", skipna=True)) + finite = np.isfinite(annual_mse_maps) + weighted_sum = (annual_mse_maps.where(finite) * area).sum(spatial_dims, skipna=True) + weight_sum = area.where(finite).sum(spatial_dims, skipna=True) + annual_rmse = np.sqrt(weighted_sum / weight_sum) + + # Area-weighted fraction of the comparable cells that carry data, per year. + # Weighting by area matters: losing a band of tropical cells is a bigger + # loss of information than the same count near the pole. + covered = (finite_samples * area).sum(spatial_dims) / ( + (total_samples * area).sum(spatial_dims) + ) + + computed = xr.Dataset( + { + "rmse_map": rmse_map, + "annual_rmse": annual_rmse, + "weight_sum": weight_sum, + "covered": covered, + } + ).compute() + + coverage = dict( + zip(years, (float(v) for v in computed["covered"].values), strict=True) + ) + starved = { + year: frac for year, frac in coverage.items() if not frac >= min_year_coverage + } + if starved: + detail = ", ".join(f"{year}: {frac:.1%}" for year, frac in starved.items()) + raise ValueError( + f"{context}: year(s) with too little paired data to carry an " + f"equal-weight block ({detail}; require " + f"{min_year_coverage:.0%}). The score and its bootstrap weight years " + "equally, so a sparse year would distort both. Narrow the window, or " + "check the observation product's coverage for those years." + ) + + annual_values = np.asarray(computed["annual_rmse"].values, dtype=float) + return ( + computed["rmse_map"], + dict(zip(years, (float(v) for v in annual_values), strict=True)), + interannual_rmse_summary(annual_values, bootstrap_samples=bootstrap_samples), + ) diff --git a/samudra/metrics/observations.py b/samudra/metrics/observations.py new file mode 100644 index 0000000000000000000000000000000000000000..87326bde3b76770a35de042560e85d3507de3aa0 --- /dev/null +++ b/samudra/metrics/observations.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Loaders for observation products, and the model-side adapter. + +Three observation products back the metrics: + +- **DUACS** — surface geostrophic velocity from satellite altimetry, used for + the velocity and EKE metrics. +- **OISST** — sea-surface temperature, used for the SST metrics. +- **ARGO-IAP** — gridded temperature profiles, used for ocean heat content. + +Each ships on its own native grid with its own naming, so the loaders resolve +coordinate and variable aliases, put longitude on `[0, 360)`, and attach a +spherical cell area where the product does not carry one. + +The model adapter is deliberately thin. Our `predictions.zarr` is already +written in the OM4 analysis layout with true 2-D `lat`/`lon`, `areacello`, `dz` +and a wet mask (see `samudra/utils/writer.py`), so there is no geometry to +reconstruct — only a presentation change from `(y, x)` dims to the `(lat, lon)` +dims the observation grids use. +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import numpy as np +import xarray as xr + +from samudra.constants import DatasetSpec +from samudra.metrics import kernels +from samudra.utils.location import ResolvedLocation + +if TYPE_CHECKING: + from samudra.config import ObsMetricsConfig + +logger = logging.getLogger(__name__) + +EARTH_RADIUS_M = kernels.EARTH_RADIUS_M + +LAT_ALIASES = ["lat", "latitude", "y", "nav_lat"] +LON_ALIASES = ["lon", "longitude", "x", "nav_lon"] +TIME_ALIASES = ["time", "valid_time"] +# `depth_std` is what the IAP/ARGO product calls its standard-level depth axis. +DEPTH_ALIASES = ["depth", "depth_std", "lev", "level", "z", "DEPTH", "depth_0", "z_t"] + +# DUACS ships both absolute geostrophic velocity (from absolute dynamic +# topography) and the anomaly (from sea-level anomaly). The model derives +# absolute velocity from `zos`, so `absolute` is the like-for-like comparison. +DUACS_VELOCITY_COMPONENTS = { + "absolute": { + "u": ["ugos", "surface_geostrophic_eastward_sea_water_velocity"], + "v": ["vgos", "surface_geostrophic_northward_sea_water_velocity"], + "label": "absolute geostrophic velocity", + }, + "anomaly": { + "u": ["ugosa"], + "v": ["vgosa"], + "label": "geostrophic velocity anomaly", + }, +} + +SST_ALIASES = ["sst", "analysed_sst", "thetao", "temperature"] +ARGO_TEMPERATURE_ALIASES = ["thetao", "temp", "temperature", "T"] + + +def find_coord_name( + obj: xr.Dataset | xr.DataArray, candidates: Sequence[str], required: bool = True +) -> str | None: + """First of `candidates` present as a coordinate or dimension.""" + for name in candidates: + if name in obj.coords or name in obj.dims: + return name + if required: + raise KeyError( + f"Could not find a coordinate among {candidates}. " + f"Available: {list(obj.coords)} / {list(obj.dims)}" + ) + return None + + +def find_var_name(ds: xr.Dataset, candidates: Sequence[str]) -> str: + """First of `candidates` present as a data variable.""" + for name in candidates: + if name in ds.data_vars: + return name + raise KeyError( + f"Could not find a variable among {candidates}. Available: {list(ds.data_vars)}" + ) + + +def spherical_cell_area(lat_vals: np.ndarray, lon_vals: np.ndarray) -> np.ndarray: + """Cell area on a sphere for a rectilinear lat/lon grid, in m^2.""" + lat = np.asarray(lat_vals, dtype=float) + lon = np.asarray(lon_vals, dtype=float) + if lat.ndim != 1 or lon.ndim != 1: + raise ValueError("Expected 1-D lat/lon coordinates") + dlon = np.gradient(np.deg2rad(lon)) + dlat = np.gradient(np.deg2rad(lat)) + dx = EARTH_RADIUS_M * np.cos(np.deg2rad(lat))[:, None] * dlon[None, :] + dy = EARTH_RADIUS_M * dlat[:, None] * np.ones((lat.size, lon.size)) + return np.abs(dx) * np.abs(dy) + + +def standardize(ds: xr.Dataset) -> xr.Dataset: + """Rename to `lat`/`lon`/`time`, coerce time, and normalize longitude.""" + lat_name = find_coord_name(ds, LAT_ALIASES) + lon_name = find_coord_name(ds, LON_ALIASES) + time_name = find_coord_name(ds, TIME_ALIASES) + + ds = ds.assign_coords( + {time_name: kernels.coerce_datetime_values(ds[time_name].values)} + ) + ds = ds.rename({lat_name: "lat", lon_name: "lon", time_name: "time"}) + ds = kernels.normalize_lon(ds, "lon") + ds = ds.sortby("lat").sortby("lon") + return _drop_singleton_dims(ds) + + +def _drop_singleton_dims(ds: xr.Dataset) -> xr.Dataset: + """Squeeze away degenerate non-physical dims some products carry.""" + keep = {"time", "lat", "lon", *DEPTH_ALIASES} + squeeze = [dim for dim, size in ds.sizes.items() if size == 1 and dim not in keep] + return ds.squeeze(squeeze, drop=True) if squeeze else ds + + +def with_cell_area(ds: xr.Dataset) -> xr.Dataset: + """Attach a spherical `area` variable unless the product already has one.""" + if "area" in ds: + return ds + area = spherical_cell_area(ds["lat"].values, ds["lon"].values) + return ds.assign(area=(["lat", "lon"], area)) + + +def open_product(location: ResolvedLocation, name: str) -> xr.Dataset: + """Open and standardize one observation product.""" + logger.info("Opening observation product %s from %s", name, location) + ds = location.open(chunks={}) + ds = standardize(ds) + ds = with_cell_area(ds) + logger.info( + " %s: %s, %s to %s", + name, + dict(ds.sizes), + np.datetime_as_string(ds["time"].values[0], unit="D"), + np.datetime_as_string(ds["time"].values[-1], unit="D"), + ) + return ds + + +def duacs_velocity( + ds: xr.Dataset, kind: str = "absolute" +) -> tuple[xr.DataArray, xr.DataArray]: + """The DUACS eastward/northward geostrophic velocity pair, equator masked.""" + if kind not in DUACS_VELOCITY_COMPONENTS: + raise ValueError( + f"Unknown DUACS velocity kind {kind!r}; " + f"expected one of {sorted(DUACS_VELOCITY_COMPONENTS)}" + ) + config = DUACS_VELOCITY_COMPONENTS[kind] + u = ds[find_var_name(ds, config["u"])] + v = ds[find_var_name(ds, config["v"])] + return ( + kernels.apply_equatorial_mask(u, "lat"), + kernels.apply_equatorial_mask(v, "lat"), + ) + + +def model_on_latlon_grid(ds: xr.Dataset, dataset_spec: DatasetSpec) -> xr.Dataset: + """Present a rollout on `(lat, lon)` dims for comparison with observations. + + Our rollouts carry `(y, x)` dims alongside *2-D* `lat`/`lon` coordinates, so + the coordinates have to be dropped before `y`/`x` can take those names — + otherwise the rename collides. On a rectilinear ("gaussian") grid the `y` + and `x` axes are exactly the latitude and longitude in degrees, so the swap + is lossless. + + Curvilinear grids are refused rather than mishandled: there `y`/`x` are + index space, not degrees, and comparing against observations needs a real + regridding step (xesmf) that is not implemented here yet. + """ + if dataset_spec.grid_type != "gaussian": + raise NotImplementedError( + f"Observation metrics need a rectilinear model grid, but this dataset " + f"has grid_type={dataset_spec.grid_type!r}. On a curvilinear grid the " + "'y'/'x' axes are index space rather than degrees, so comparison " + "requires conservative regridding (xesmf) onto the observation grid. " + "See issues #801 and #809." + ) + + dims = {str(dim) for dim in ds.dims} + if {"y", "x"} <= dims: + # The eval writer's layout: `(y, x)` dims under 2-D lat/lon coords. + # Those coords must go before `y`/`x` can take their names, or the + # rename collides. `areacello` is kept -- it is the model's own cell + # area, better than anything re-derived from the axes. + areacello = ( + ds["areacello"] if "areacello" in ds.coords or "areacello" in ds else None + ) + ds = ds.drop_vars(["lat", "lon", "lat_b", "lon_b"], errors="ignore") + ds = ds.rename({"y": "lat", "x": "lon"}) + if areacello is not None: + renamed_area = areacello.drop_vars( + ["lat", "lon", "lat_b", "lon_b"], errors="ignore" + ).rename({"y": "lat", "x": "lon"}) + ds = ds.assign_coords(areacello=renamed_area) + elif {"lat", "lon"} <= dims: + # The form a `DataSource` hands back, as the OM4 baseline does: + # already 1-D lat/lon, with the 2-D geography on `lat_2d`/`lon_2d`. + # Drop those so alignment sees only the 1-D axes. + ds = ds.drop_vars(["lat_2d", "lon_2d", "lat_b", "lon_b"], errors="ignore") + else: + raise ValueError( + "Expected horizontal dimensions ('y', 'x') as written by the eval " + "rollout, or ('lat', 'lon') as standardised by a DataSource; found " + f"dims {sorted(dims)}" + ) + + ds = kernels.normalize_lon(ds, "lon").sortby("lat").sortby("lon") + + # OM4 and its rollouts carry a cftime axis, which cannot be compared + # against the pandas timestamps used to select windows. Normalise at this + # single entry point so callers need not know the store's calendar. + if "time" in ds.coords: + ds = ds.assign_coords(time=kernels.coerce_datetime_values(ds["time"].values)) + return ds + + +def model_cell_area(ds: xr.Dataset) -> xr.DataArray: + """The rollout's own cell area, falling back to a spherical estimate.""" + if "areacello" in ds.coords or "areacello" in ds: + return ds["areacello"] + logger.warning( + "Rollout carries no 'areacello'; falling back to a spherical area estimate" + ) + return xr.DataArray( + spherical_cell_area(ds["lat"].values, ds["lon"].values), + dims=["lat", "lon"], + coords={"lat": ds["lat"], "lon": ds["lon"]}, + ) + + +def model_depth_thickness(ds: xr.Dataset, dataset_spec: DatasetSpec) -> xr.DataArray: + """Exact native layer thicknesses aligned to the rollout's `lev` axis.""" + lev = ds["lev"] + if "dz" in ds.coords or "dz" in ds: + return ds["dz"] + thickness = np.asarray(dataset_spec.depth_thickness[: lev.size], dtype=float) + return xr.DataArray(thickness, dims=["lev"], coords={"lev": lev}) + + +def open_products( + cfg: ObsMetricsConfig, data_root: ResolvedLocation +) -> dict[str, xr.Dataset]: + """Open and standardize the three observation products.""" + return { + "duacs": open_product(data_root.resolve(cfg.duacs_location), "DUACS"), + "oisst": open_product(data_root.resolve(cfg.oisst_location), "OISST"), + "argo": open_product(data_root.resolve(cfg.argo_iap_location), "ARGO-IAP"), + } diff --git a/samudra/metrics/report.py b/samudra/metrics/report.py new file mode 100644 index 0000000000000000000000000000000000000000..eeccfd7bc16fba23c6145dc51ddfe552cc5048e7 --- /dev/null +++ b/samudra/metrics/report.py @@ -0,0 +1,483 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Driver that turns a finished rollout into observation metrics. + +The output is one tidy long-format `pandas.DataFrame`: one row per +(metric, model, depth layer, period). That frame is the contract. `to_wandb` +flattens the primary rows into scalars for run-to-run comparison, and the whole +frame goes to CSV so the per-year detail behind each uncertainty interval stays +available. + +Every metric follows the same shape, which is why they share so much code: +put the model on the observation grid, form the squared error, reduce it to a +per-cell map plus per-year totals, and area-weight the map to one scalar. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr + +from samudra.metrics import comparisons, kernels +from samudra.utils.wandb import MetricsDict + +logger = logging.getLogger(__name__) + +# Column order for the emitted frame. Kept explicit so the CSV is stable across +# runs even when a metric produces no rows. +COLUMNS = [ + "metric", + "model", + "depth", + "value", + "units", + "period_kind", + "period_start", + "period_end", + "year", + "n_time_samples", + "annual_std", + "ci_low", + "ci_high", + "map_weighted_rmse", + "n_years", + "n_bootstrap", + "uncertainty_method", + "grid_shape", + "n_paired_cells", + "paired_ocean_fraction", +] + +SPATIAL_DIMS = ("lat", "lon") + + +def _grid_shape(field: xr.DataArray) -> str: + return f"{field.sizes.get('lat', 0)}x{field.sizes.get('lon', 0)}" + + +def _rows_for_metric( + *, + metric: str, + model: str, + units: str, + error_squared: xr.DataArray, + area: xr.DataArray, + bootstrap_samples: int, + depth: str | None = None, + pairing: comparisons.Pairing | None = None, +) -> tuple[list[dict[str, Any]], xr.DataArray]: + """Reduce a squared-error field to primary and per-year metric rows. + + Returns the rows and the per-cell RMSE map, so callers that also want the + map (for figures, or for a downstream variance comparison) do not recompute. + """ + context = f"{model} {metric}" + (f" [{depth}]" if depth else "") + rmse_map, annual, summary = kernels.rmse_map_with_uncertainty( + error_squared, + area, + SPATIAL_DIMS, + context, + bootstrap_samples=bootstrap_samples, + ) + # The equal-year block estimate, which is what the bootstrap interval is + # built from, so a value can never fall outside its own interval. The + # alternative below is the same two averages in the opposite order; they + # diverge when a product's coverage varies by year. + total = float(summary["block_aggregate_rmse"]) + map_weighted_total = kernels.area_weighted_map_rmse(rmse_map, area) + time_index = pd.DatetimeIndex(error_squared["time"].values) + + primary = { + "metric": metric, + "model": model, + "depth": depth, + "value": total, + "units": units, + "period_kind": "primary_complete_years", + "period_start": f"{time_index.min():%Y-%m-%d}", + "period_end": f"{time_index.max():%Y-%m-%d}", + "year": np.nan, + "n_time_samples": int(time_index.size), + "grid_shape": _grid_shape(rmse_map), + # Kept alongside so the two aggregations stay comparable: a gap between + # this and `value` means the observation coverage varied by year. + "map_weighted_rmse": map_weighted_total, + **(pairing.as_row() if pairing else {}), + **summary, + } + primary.pop("primary_aggregation_method", None) + primary.pop("block_aggregate_rmse", None) # now reported as `value` + + rows = [primary] + for year, value in annual.items(): + year_index = time_index[time_index.year == year] + rows.append( + { + "metric": metric, + "model": model, + "depth": depth, + "value": value, + "units": units, + "period_kind": "annual", + "period_start": f"{year_index.min():%Y-%m-%d}", + "period_end": f"{year_index.max():%Y-%m-%d}", + "year": int(year), + "n_time_samples": int(year_index.size), + "grid_shape": _grid_shape(rmse_map), + **(pairing.as_row() if pairing else {}), + } + ) + logger.info(" %s = %.6g %s", context, total, units) + return rows, rmse_map + + +def _variance_map_rows( + *, + metric_prefix: str, + model: str, + units: str, + model_field: xr.DataArray, + obs_field: xr.DataArray, + area: xr.DataArray, +) -> list[dict[str, Any]]: + """Residual-variance map comparison: RMSE of the maps, plus pattern correlation. + + The two numbers answer different questions. The RMSE says whether the model + has the right *amount* of residual variability; the pattern correlation says + whether it puts that variability in the right *places*. A model can score + well on one and badly on the other. + + `model_field` must be on the model's **native** grid. Variance is computed + there and only the resulting variance map is regridded, because horizontal + interpolation is linear while variance is quadratic: interpolating first + smooths the field and systematically damps the variance. The effect scales + with how much the regrid coarsens -- negligible for SST (quarter degree to + quarter degree) but a ~6% underestimate for OHC (quarter to half degree). + """ + model_var = kernels.residual_variance_map(model_field) + obs_var = kernels.residual_variance_map(obs_field) + model_var = kernels.model_field_on_obs_grid(model_var, obs_var) + model_var, obs_var = xr.align(model_var, obs_var, join="inner") + + computed = xr.Dataset({"model": model_var, "obs": obs_var}).compute() + model_var, obs_var = computed["model"], computed["obs"] + + rmse = kernels.area_weighted_map_rmse(model_var - obs_var, area) + corr = kernels.area_weighted_pattern_corr(model_var, obs_var, area) + time_index = pd.DatetimeIndex(model_field["time"].values) + shared = { + "model": model, + "depth": None, + "period_kind": "full_overlap", + "period_start": f"{time_index.min():%Y-%m-%d}", + "period_end": f"{time_index.max():%Y-%m-%d}", + "year": np.nan, + "n_time_samples": int(time_index.size), + "grid_shape": _grid_shape(model_var), + **comparisons.Comparison( + metric_prefix, model_var, obs_var, area + ).pairing.as_row(), + } + logger.info( + " %s residual variance: map rmse = %.6g, pattern corr = %.4f", + f"{model} {metric_prefix}", + rmse, + corr, + ) + return [ + { + "metric": f"{metric_prefix}_residual_variance_map_rmse", + "value": rmse, + "units": units, + **shared, + }, + { + "metric": f"{metric_prefix}_residual_variance_pattern_corr", + "value": corr, + "units": "1", + **shared, + }, + ] + + +def _velocity_metrics( + model: str, + rollout: xr.Dataset, + duacs: xr.Dataset, + window: tuple[pd.Timestamp, pd.Timestamp], + bootstrap_samples: int, + velocity_kind: str, +) -> list[dict[str, Any]]: + """Surface geostrophic velocity vector RMSE and instantaneous EKE RMSE vs DUACS.""" + velocity = comparisons.surface_velocity( + rollout, duacs, window, model, velocity_kind + ) + rows, _ = _rows_for_metric( + metric="surface_geostrophic_velocity_vector_total_rmse", + model=model, + units="m s-1", + error_squared=velocity.vector_error_squared, + area=velocity.area, + bootstrap_samples=bootstrap_samples, + pairing=velocity.eastward.pairing, + ) + + eke = velocity.eddy_kinetic_energy() + eke_rows, _ = _rows_for_metric( + metric="instantaneous_surface_eke_total_rmse", + model=model, + units="m2 s-2", + error_squared=eke.error_squared, + area=eke.area, + bootstrap_samples=bootstrap_samples, + pairing=eke.pairing, + ) + return rows + eke_rows + + +def _sst_metrics( + model: str, + rollout: xr.Dataset, + oisst: xr.Dataset, + window: tuple[pd.Timestamp, pd.Timestamp], + bootstrap_samples: int, +) -> list[dict[str, Any]]: + """SST RMSE and residual-variance map diagnostics vs OISST.""" + scored = comparisons.sea_surface_temperature(rollout, oisst, window, model) + rows, _ = _rows_for_metric( + metric="surface_sst_total_rmse", + model=model, + units="degC", + error_squared=scored.error_squared, + area=scored.area, + bootstrap_samples=bootstrap_samples, + pairing=scored.pairing, + ) + + # Residual variance uses the whole rollout, not the scoring window. + full = comparisons.sea_surface_temperature( + rollout, oisst, None, f"{model} SST variance" + ) + rows += _variance_map_rows( + metric_prefix="surface_sst", + model=model, + units="degC2", + model_field=full.native, # native grid; variance is regridded, not the field + obs_field=full.obs, + area=full.area, + ) + return rows + + +def _ohc_metrics( + model: str, + rollout: xr.Dataset, + argo: xr.Dataset, + window: tuple[pd.Timestamp, pd.Timestamp], + bootstrap_samples: int, + model_dz: xr.DataArray, +) -> list[dict[str, Any]]: + """Per-area OHC RMSE by layer, plus upper-700 m residual-variance diagnostics.""" + rows: list[dict[str, Any]] = [] + upper_label = kernels.OHC_LAYERS[0].label + for layer in kernels.OHC_LAYERS: + scored = comparisons.ohc_layer(rollout, argo, layer, window, model, model_dz) + layer_rows, _ = _rows_for_metric( + metric="ohc_per_area_total_rmse", + model=model, + units="J m-2", + error_squared=scored.error_squared, + area=scored.area, + bootstrap_samples=bootstrap_samples, + depth=layer.label, + pairing=scored.pairing, + ) + rows += layer_rows + # Emitted as rows rather than columns: they qualify the layer score + # without being one, and a tidy frame is the place for that. + for metric, value, units in ( + ( + "ohc_bathymetry_disagreement", + scored.bathymetry_disagreement_m, + "m", + ), + ("ohc_partial_column_fraction", scored.model_partial_columns, ""), + ): + rows.append( + { + "metric": metric, + "model": model, + "depth": layer.label, + "value": value, + "units": units, + "period_kind": "full_overlap", + "period_start": layer_rows[0]["period_start"], + "period_end": layer_rows[0]["period_end"], + "grid_shape": layer_rows[0]["grid_shape"], + } + ) + + if layer.label == upper_label: + full = comparisons.ohc_layer( + rollout, + argo, + layer, + None, + f"{model} variance", + model_dz, + complete_years_only=False, + ) + rows += _variance_map_rows( + metric_prefix="ohc_upper700_per_area", + model=model, + units="(J m-2)2", + model_field=full.native, # native grid; variance is regridded + obs_field=full.obs, + area=full.area, + ) + return rows + + +def compute_observation_metrics( + rollouts: dict[str, xr.Dataset], + *, + duacs: xr.Dataset, + oisst: xr.Dataset, + argo: xr.Dataset, + model_dz: dict[str, xr.DataArray], + window: tuple[pd.Timestamp, pd.Timestamp], + bootstrap_samples: int = kernels.DEFAULT_BOOTSTRAP_SAMPLES, + velocity_kind: str = "absolute", +) -> pd.DataFrame: + """Score one or more rollouts against the observation products. + + Args: + rollouts: Label to rollout dataset, already on `(lat, lon)` dims via + `observations.model_on_latlon_grid`. Typically the model under + evaluation plus an OM4 baseline. + duacs: Standardized DUACS product, for velocity and EKE. + oisst: Standardized OISST product, for SST. + argo: Standardized ARGO-IAP product, for ocean heat content. + model_dz: Label to native layer thickness, for exact OHC integration. + window: Inclusive start/end of the primary evaluation period. + bootstrap_samples: Calendar-year block-bootstrap draws; 0 disables. + velocity_kind: Which DUACS geostrophic velocity to compare against. + + Returns: + Tidy frame with one row per (metric, model, depth, period). + """ + frames: list[dict[str, Any]] = [] + for label, rollout in rollouts.items(): + logger.info("Scoring %r against observations", label) + frames += _velocity_metrics( + label, rollout, duacs, window, bootstrap_samples, velocity_kind + ) + frames += _sst_metrics(label, rollout, oisst, window, bootstrap_samples) + frames += _ohc_metrics( + label, rollout, argo, window, bootstrap_samples, model_dz[label] + ) + + frame = pd.DataFrame(frames) + for column in COLUMNS: + if column not in frame.columns: + frame[column] = np.nan + return frame[COLUMNS] + + +# Short, stable W&B key per (metric, depth). Keys have to survive refactors, +# because a renamed key silently breaks every historical run comparison. +_WANDB_KEYS = { + ("surface_geostrophic_velocity_vector_total_rmse", None): "velocity/total_rmse", + ("instantaneous_surface_eke_total_rmse", None): "eke/total_rmse", + ("surface_sst_total_rmse", None): "sst/total_rmse", + ( + "ohc_per_area_total_rmse", + kernels.OHC_LAYERS[0].label, + ): "ohc_0_700/per_area_total_rmse", + ( + "ohc_per_area_total_rmse", + kernels.OHC_LAYERS[1].label, + ): "ohc_700_2000/per_area_total_rmse", + ( + "ohc_bathymetry_disagreement", + kernels.OHC_LAYERS[0].label, + ): "ohc_0_700/bathymetry_disagreement_m", + ( + "ohc_bathymetry_disagreement", + kernels.OHC_LAYERS[1].label, + ): "ohc_700_2000/bathymetry_disagreement_m", + ( + "ohc_partial_column_fraction", + kernels.OHC_LAYERS[0].label, + ): "ohc_0_700/partial_column_fraction", + ( + "ohc_partial_column_fraction", + kernels.OHC_LAYERS[1].label, + ): "ohc_700_2000/partial_column_fraction", + ( + "surface_sst_residual_variance_map_rmse", + None, + ): "sst/residual_variance_map_rmse", + ( + "surface_sst_residual_variance_pattern_corr", + None, + ): "sst/residual_variance_pattern_corr", + ( + "ohc_upper700_per_area_residual_variance_map_rmse", + None, + ): "ohc_upper700/residual_variance_map_rmse", + ( + "ohc_upper700_per_area_residual_variance_pattern_corr", + None, + ): "ohc_upper700/residual_variance_pattern_corr", +} + +# Uncertainty columns worth promoting to their own scalars, so a run comparison +# can show the interval rather than only the point estimate. +_UNCERTAINTY_SUFFIXES = ("annual_std", "ci_low", "ci_high") + + +def to_wandb(frame: pd.DataFrame, primary_model: str) -> MetricsDict: + """Flatten the primary rows into W&B scalars. + + The model under evaluation gets bare `obs/...` keys; every other model in + the frame is namespaced under `obs//...` so the baseline never + shadows the thing being evaluated. + """ + metrics: MetricsDict = {} + scored = frame[ + frame["period_kind"].isin(["primary_complete_years", "full_overlap"]) + ] + + for _, row in scored.iterrows(): + depth = row["depth"] if isinstance(row["depth"], str) else None + suffix = _WANDB_KEYS.get((row["metric"], depth)) + if suffix is None: + continue + model = str(row["model"]) + prefix = "obs" if model == primary_model else f"obs/{_slug(model)}" + metrics[f"{prefix}/{suffix}"] = float(row["value"]) + for column in _UNCERTAINTY_SUFFIXES: + value = row.get(column) + if value is not None and np.isfinite(value): + metrics[f"{prefix}/{suffix}_{column}"] = float(value) + # Per metric: they score different spans, so one figure would describe + # whichever scored the most years. + n_years = row.get("n_years") + if n_years is not None and np.isfinite(n_years): + metrics[f"{prefix}/{suffix}_n_years"] = float(n_years) + + # The window stays in the table and CSV: W&B scalars are floats, and a date + # is not one. There is no single `obs/n_years` -- metrics score different + # spans, so it is reported per metric above. + return metrics + + +def _slug(model: str) -> str: + return "".join(c if c.isalnum() else "_" for c in model.lower()).strip("_") diff --git a/samudra/metrics/run.py b/samudra/metrics/run.py new file mode 100644 index 0000000000000000000000000000000000000000..bb6f69e8662c1531d908d0b76af9d0f12ce1e3ac --- /dev/null +++ b/samudra/metrics/run.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Glue between an eval job and the metric kernels. + +`samudra.eval` calls `run_observation_metrics` once its rollout is on disk; +`python -m samudra.metrics` calls the same function against an existing +`predictions.zarr`, so a metrics-only rerun never repeats a multi-hour rollout. +""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +import pandas as pd +import xarray as xr + +from samudra.constants import DatasetSpec +from samudra.metrics import observations, report +from samudra.utils.data import stack_levels +from samudra.utils.location import ResolvedLocation +from samudra.utils.wandb import MetricsDict + +if TYPE_CHECKING: + from samudra.config import ObsMetricsConfig + +logger = logging.getLogger(__name__) + +METRICS_CSV_NAME = "observation_metrics.csv" + + +def open_predictions(output_dir: Path) -> xr.Dataset: + """Open the rollout an eval job wrote, failing clearly when it is absent.""" + path = Path(output_dir) / "predictions.zarr" + if not path.exists(): + raise FileNotFoundError( + f"No rollout found at {path}. Observation metrics score a rollout that " + "has been written to disk, so the eval job must run with save_zarr=true." + ) + return xr.open_zarr(path, chunks={}) + + +def analysis_ready(data: xr.Dataset, dataset_spec: DatasetSpec) -> xr.Dataset: + """Put a dataset in depth-stacked form, matching the rollout writer's layout.""" + if any( + name in data.data_vars + for name in ( + f"thetao_{i}" for i in range(dataset_spec.num_prognostic_depth_levels) + ) + ): + return stack_levels(data, dataset_spec) + return data + + +@dataclass +class ObservationRun: + """Everything one observation pass produces. + + `eval` wants the scalars, `viz` wants the products and the rollouts to draw + from; both want the frame. Returning all of it means the pass happens once + per job rather than once per consumer. + """ + + products: dict[str, xr.Dataset] + rollouts: dict[str, xr.Dataset] + frame: pd.DataFrame + scalars: MetricsDict + + +def score_rollouts( + obs_cfg: ObsMetricsConfig, + *, + rollouts: dict[str, xr.Dataset], + dataset_spec: DatasetSpec, + data_root: ResolvedLocation, + primary_label: str, + output_dir: Path | None = None, +) -> ObservationRun: + """Open the observation products and score every rollout against them. + + Args: + obs_cfg: Product locations, scoring window, and bootstrap settings. + rollouts: Datasets to score, keyed by the label they are reported under. + Depth-stacked or not; either layout is accepted. + dataset_spec: Spec of the dataset the rollouts came from. + data_root: Root the observation locations resolve against. + primary_label: Which rollout the W&B scalars describe; the others are + baselines, reported under their own keys. + output_dir: When given, the full frame is written here as CSV. + + Returns: + The products, the rollouts on the observation grid, the tidy frame, and + the flattened scalars for W&B. + """ + start = time.perf_counter() + products = observations.open_products(obs_cfg, data_root) + + prepared: dict[str, xr.Dataset] = {} + model_dz: dict[str, xr.DataArray] = {} + for label, data in rollouts.items(): + on_grid = observations.model_on_latlon_grid( + analysis_ready(data, dataset_spec), dataset_spec + ) + prepared[label] = on_grid + model_dz[label] = observations.model_depth_thickness(on_grid, dataset_spec) + + logger.info( + "Computing observation metrics over %s to %s for: %s", + obs_cfg.rmse_start, + obs_cfg.rmse_end, + ", ".join(prepared), + ) + frame = report.compute_observation_metrics( + prepared, + duacs=products["duacs"], + oisst=products["oisst"], + argo=products["argo"], + model_dz=model_dz, + window=obs_cfg.window, + bootstrap_samples=obs_cfg.bootstrap_samples, + velocity_kind=obs_cfg.velocity_kind, + ) + + if output_dir is not None: + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + csv_path = output_dir / METRICS_CSV_NAME + frame.to_csv(csv_path, index=False) + logger.info("Wrote %d metric rows to %s", len(frame), csv_path) + + elapsed = time.perf_counter() - start + logger.info("Observation metrics took %.1f s", elapsed) + + scalars = report.to_wandb(frame, primary_label) + scalars["obs/seconds"] = elapsed + return ObservationRun(products, prepared, frame, scalars) + + +def run_observation_metrics( + obs_cfg: ObsMetricsConfig, + *, + predictions: xr.Dataset, + dataset_spec: DatasetSpec, + data_root: ResolvedLocation, + model_label: str, + baselines: dict[str, xr.Dataset] | None = None, + output_dir: Path | None = None, +) -> tuple[pd.DataFrame, MetricsDict]: + """Score a rollout against observations and return the frame plus W&B scalars. + + The eval-facing shape of `score_rollouts`, which `viz` uses directly because + it also needs the products and the prepared rollouts to draw from. + """ + scored = score_rollouts( + obs_cfg, + rollouts={model_label: predictions, **(baselines or {})}, + dataset_spec=dataset_spec, + data_root=data_root, + primary_label=model_label, + output_dir=output_dir, + ) + return scored.frame, scored.scalars diff --git a/samudra/metrics/spectra.py b/samudra/metrics/spectra.py new file mode 100644 index 0000000000000000000000000000000000000000..127eea80764d75fae2542f7217ffc68a0283ec0b --- /dev/null +++ b/samudra/metrics/spectra.py @@ -0,0 +1,476 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Spectral kernels for observation-based evaluation. + +Spectra answer a question the RMSE metrics cannot: *at which scales* a model +gets the ocean wrong. A rollout can match the observed variance overall and +still put it at the wrong wavelengths, and a model that smooths eddies away +looks fine on a map average while its spectrum falls off early. + +Like `kernels`, everything here is a pure function over arrays. Errors are +reported in log10 space (dex), because power spectra span many orders of +magnitude and a linear residual would be dominated entirely by the largest +scales. + +Implements the spectral metrics of the metrics document [1]. + +[1]: https://docs.google.com/document/d/1cWzIrcajfW-gvoMzV3PU0JDAouVt5xerRmEiBIUbV1s +""" + +from __future__ import annotations + +import logging + +import numpy as np +import pandas as pd +import xarray as xr +from scipy import signal + +logger = logging.getLogger(__name__) + +# Wavenumber bins per spatial spectrum, as a fraction of the shorter grid side. +DEFAULT_BIN_FACTOR = 2 + +# Smallest region side a spatial spectrum is computed for. +MIN_REGION_SIDE = 8 + +METRES_PER_DEGREE = 111_320.0 + +# Cycles per metre to radians per kilometre. +RADIANS_PER_KM = 2.0 * np.pi * 1000.0 + +# Welch segment length. Long enough to resolve the seasonal cycle at 5-day +# sampling, short enough to leave several independent segments in an 8-year +# record. +DEFAULT_NPERSEG = 256 + +# Regions the spectra are reported over: open-ocean boxes away from coasts, +# plus the two western-boundary systems where mesoscale activity is strongest. +SPATIAL_REGIONS: tuple[tuple[str, slice, slice], ...] = ( + ("North Pacific", slice(170, 200), slice(25, 45)), + ("Gulf Stream", slice(300, 320), slice(25, 45)), + ("Agulhas", slice(40, 60), slice(-50, -30)), +) +TEMPORAL_REGIONS: tuple[tuple[str, slice, slice], ...] = ( + ("Global", slice(0, 360), slice(-90, 90)), + ("Gulf Stream", slice(300, 320), slice(25, 45)), + ("Agulhas", slice(40, 60), slice(-50, -30)), +) + + +def _periodic_hann(size: int) -> np.ndarray: + """Hann window with no repeated endpoint (scipy's "sym=False").""" + return 0.5 - 0.5 * np.cos(2.0 * np.pi * np.arange(size) / size) + + +def remove_plane(field: np.ndarray) -> np.ndarray: + """Subtract a least-squares 2-D plane from the last two axes. + + A spatial trend across the box is a boundary effect, not a physical scale. + Left in place it leaks power across every wavenumber once the window is + applied. + """ + *lead, height, width = field.shape + y = np.linspace(-1.0, 1.0, height) + x = np.linspace(-1.0, 1.0, width) + grid_y, grid_x = np.meshgrid(y, x, indexing="ij") + design = np.stack([grid_x.ravel(), grid_y.ravel(), np.ones(grid_x.size)], axis=1) + flat = field.reshape(-1, height * width) + coefficients, *_ = np.linalg.lstsq(design, flat.T, rcond=None) + plane = (design @ coefficients).T.reshape(-1, height, width) + return (flat.reshape(-1, height, width) - plane).reshape(*lead, height, width) + + +def bin_index(values: np.ndarray, edges: np.ndarray) -> np.ndarray: + """Which bin each value falls in, given the full edge array. + + A value landing exactly on an edge belongs to the bin *above* it. On a + regular grid that is not a corner case: whole rows of the mode lattice can + coincide with edges, and the opposite convention moves every one of them + down a bin. + + `right=False` is the numpy spelling of the reference's + `torch.bucketize(..., right=True)`; the two libraries use the flag to mean + opposite things, which is easy to port backwards. + """ + return np.digitize(values, edges[1:-1], right=False) + + +def isotropic_spectrum( + field: np.ndarray, + dx: float = 1.0, + dy: float = 1.0, + num_bins: int | None = None, + detrend: bool = True, +) -> tuple[np.ndarray, np.ndarray]: + """Radially averaged power spectrum of a 2-D field, or a stack of them. + + Args: + field: `(height, width)`, or `(n, height, width)` for a stack. + dx: Grid spacing along the last axis, in the units the wavenumber + should be reciprocal to. + dy: Grid spacing along the second-to-last axis. + num_bins: Wavenumber bins across the full Fourier plane, before the + bins above the lower Nyquist are dropped; defaults to the shorter + side over `DEFAULT_BIN_FACTOR`. Fewer bins come back than are asked + for, which is why the returned wavenumbers must be used rather than + assumed. + detrend: Remove a 2-D plane before windowing. + + Returns: + Bin-center wavenumbers and the isotropic spectrum, in cycles per unit + of `dx`, the latter shaped `(bins,)` or `(n, bins)` to match the input. + + The spectrum is multiplied by its bin-center wavenumber, so a flat line on + log-log axes means equal energy per octave and the area under the curve is + the variance. + """ + array = np.asarray(field, dtype=float) + stacked = array.ndim == 3 + if not stacked: + array = array[None, ...] + if array.ndim != 3: + raise ValueError(f"Expected a 2-D field or a stack of them, got {array.shape}") + + n, height, width = array.shape + if num_bins is None: + num_bins = min(height, width) // DEFAULT_BIN_FACTOR + if num_bins < 2: + raise ValueError( + f"A {height}x{width} region is too small for a spectrum; it yields " + f"{num_bins} wavenumber bin(s), and a single bin describes no slope" + ) + + array = ( + remove_plane(array) + if detrend + else array - array.mean(axis=(-2, -1), keepdims=True) + ) + + # A Hann window in both directions suppresses the discontinuity at the box + # edges, which would otherwise ring across all wavenumbers. Dividing by the + # mean squared window restores the variance the taper removes. + # Periodic ("DFT-even") Hann, not the symmetric np.hanning: an FFT treats + # the box as one period, so the window must not repeat its endpoint. + window = np.outer(_periodic_hann(height), _periodic_hann(width)) + correction = float(np.mean(window**2)) + power = np.abs(np.fft.rfft2(array * window, norm="forward")) ** 2 + spectrum_2d = power / correction * (width * dx) * (height * dy) + + k_x: np.ndarray = np.fft.rfftfreq(width, d=dx) + k_y: np.ndarray = np.fft.fftfreq(height, d=dy) + grid_y, grid_x = np.meshgrid(k_y, k_x, indexing="ij") + # sqrt of the sum of squares rather than np.hypot: on a regular grid many + # modes fall exactly on a bin edge, and the two formulas disagree in the + # last bit, which moves those modes between neighbouring bins. + k_magnitude = np.sqrt(grid_x**2 + grid_y**2) + + # Bins span the whole Fourier plane out to its corner, and those beyond the + # lower Nyquist are then dropped. Sizing the bins to the corner rather than + # to the Nyquist is what keeps bin width independent of how much of the + # plane survives truncation. Nothing above the lower Nyquist is retained: + # past it the two axes disagree about what is resolved. + nyquist = min(1.0 / (2.0 * dx), 1.0 / (2.0 * dy)) + edges = np.linspace(0.0, float(k_magnitude.max()), num_bins + 1) + edges = edges[edges < nyquist] + num_bins = edges.size - 1 + if num_bins < 2: + raise ValueError( + f"A {height}x{width} region is too small for a spectrum; it yields " + f"{max(num_bins, 0)} wavenumber bin(s) below the Nyquist, and a " + f"single bin describes no slope" + ) + centers = 0.5 * (edges[:-1] + edges[1:]) + + flat_k = k_magnitude.ravel() + retained = flat_k <= edges[-1] + index = bin_index(flat_k[retained], edges) + flat_psd = spectrum_2d.reshape(n, -1)[:, retained] + + totals = np.zeros((n, num_bins)) + for row, values in zip(totals, flat_psd, strict=True): + np.add.at(row, index, values) + counts = np.bincount(index, minlength=num_bins).astype(float) + # A bin no mode fell into has no answer, not an answer of zero: zero would + # drop silently off a log axis and out of a log-space curve comparison, + # making "nothing sampled here" indistinguishable from "no energy here". + with np.errstate(invalid="ignore"): + isotropic = np.where(counts > 0, totals / np.maximum(counts, 1.0), np.nan) + isotropic = isotropic * centers + + return centers, isotropic if stacked else isotropic[0] + + +def _open_ocean_block(field: xr.DataArray, dims: tuple[str, str]) -> np.ndarray | None: + """Return the region's values, or None if any cell is missing. + + A 2-D FFT assumes a complete rectangle. Filling land or missing cells puts + step changes in the interior of the box, which a window cannot taper away + and which leak power across every wavenumber -- so the spectral regions are + curated open-ocean boxes and a box that is not fully observed gets no + spectrum at all. + """ + values = np.asarray(field.transpose(*dims).values, dtype=float) + return values if np.isfinite(values).all() else None + + +def region_spectrum( + field: xr.DataArray, + lon_range: slice, + lat_range: slice, + lat_dim: str = "lat", + lon_dim: str = "lon", + name: str = "region", +) -> tuple[np.ndarray, np.ndarray]: + """Isotropic spectrum of a region, averaged over time if the field has it. + + Returns wavenumbers in radians per kilometre, so a curve reads directly as + wavelengths (2*pi/k) and regions at different latitudes stay comparable. + Returns empty arrays, and logs why, for a region that is too small or not + fully observed -- the caller draws the panel as unavailable rather than + showing a spectrum built from filled-in land. + """ + region = field.sel({lon_dim: lon_range, lat_dim: lat_range}) + if ( + region.sizes[lat_dim] < MIN_REGION_SIDE + or region.sizes[lon_dim] < MIN_REGION_SIDE + ): + logger.info( + "no spectrum for %s: box is %dx%d, below the %dx%d minimum", + name, + region.sizes[lat_dim], + region.sizes[lon_dim], + MIN_REGION_SIDE, + MIN_REGION_SIDE, + ) + return np.array([]), np.array([]) + + # Degrees to metres at the region's mid-latitude. One factor for the whole + # box keeps the transform on a regular grid; the boxes are small enough + # that the cosine varies little across them. + lat = np.asarray(region[lat_dim].values, dtype=float) + lon = np.asarray(region[lon_dim].values, dtype=float) + dy = abs(float(np.median(np.diff(lat)))) * METRES_PER_DEGREE + dx = ( + abs(float(np.median(np.diff(lon)))) + * METRES_PER_DEGREE + * float(np.cos(np.deg2rad(lat.mean()))) + ) + if not (dx > 0 and dy > 0): + logger.info("no spectrum for %s: grid spacing is not positive", name) + return np.array([]), np.array([]) + + if "time" in region.dims: + # Every step or none: averaging over whichever steps happened to be + # fully observed would weight the result towards those conditions. + stack = [] + unobserved = [] + for step in range(region.sizes["time"]): + block = _open_ocean_block(region.isel(time=step), (lat_dim, lon_dim)) + if block is None: + unobserved.append(step) + continue + stack.append(block) + if unobserved: + logger.info( + "no spectrum for %s: %d of %d steps have land or missing cells " + "(first at index %d)", + name, + len(unobserved), + region.sizes["time"], + unobserved[0], + ) + return np.array([]), np.array([]) + if not stack: + logger.info("no spectrum for %s: no time steps", name) + return np.array([]), np.array([]) + try: + centers, spectra = isotropic_spectrum(np.stack(stack), dx=dx, dy=dy) + except ValueError as error: + logger.info("no spectrum for %s: %s", name, error) + return np.array([]), np.array([]) + return centers * RADIANS_PER_KM, spectra.mean(axis=0) + + block = _open_ocean_block(region, (lat_dim, lon_dim)) + if block is None: + logger.info("no spectrum for %s: box contains land or missing cells", name) + return np.array([]), np.array([]) + try: + centers, spectrum = isotropic_spectrum(block, dx=dx, dy=dy) + except ValueError as error: + # A box can pass the side-length check and still leave too few bins + # below the lower Nyquist once the spacings are strongly anisotropic. + # The caller draws that as unavailable, like any other region it cannot + # transform, rather than losing the whole figure. + logger.info("no spectrum for %s: %s", name, error) + return np.array([]), np.array([]) + return centers * RADIANS_PER_KM, spectrum + + +def region_mean_series( + field: xr.DataArray, + lon_range: slice, + lat_range: slice, + lat_dim: str = "lat", + lon_dim: str = "lon", + max_points: int = 100_000, +) -> pd.Series: + """Cosine-latitude-weighted mean of a region, as a time series. + + Large global grids are regularly subsampled first: a temporal spectrum + needs the time axis intact, not every cell, and the weighted mean converges + long before the full grid is used. + """ + region = field.sel({lon_dim: lon_range, lat_dim: lat_range}) + if region.sizes[lat_dim] * region.sizes[lon_dim] > max_points: + region = region.isel( + { + lat_dim: slice(None, None, max(1, region.sizes[lat_dim] // 180)), + lon_dim: slice(None, None, max(1, region.sizes[lon_dim] // 360)), + } + ) + + weights = np.cos(np.deg2rad(region[lat_dim])) + weighted = region.weighted(weights.fillna(0.0)) + series = weighted.mean((lat_dim, lon_dim), skipna=True) + return pd.Series( + np.asarray(series.values, dtype=float), + index=pd.DatetimeIndex(series["time"].values), + ) + + +def welch_psd( + series: pd.Series, nperseg: int = DEFAULT_NPERSEG +) -> tuple[np.ndarray, np.ndarray]: + """Welch power spectral density of a scalar series, in cycles per year. + + Returns empty arrays when the series is too short to segment. The zero + frequency is dropped: it carries the mean, not a timescale, and cannot be + plotted on log axes. + """ + values = series.astype(float).dropna() + if values.size < 4: + return np.array([]), np.array([]) + + index = pd.DatetimeIndex(values.index) + steps = np.diff(index.values).astype("timedelta64[s]").astype(float) / 86400.0 + steps = steps[np.isfinite(steps) & (steps > 0)] + if steps.size == 0: + return np.array([]), np.array([]) + + frequencies, power = signal.welch( + values.to_numpy(), + fs=365.25 / float(np.median(steps)), + detrend="linear", + window="hann", + scaling="density", + nperseg=min(nperseg, values.size), + ) + positive = frequencies > 0 + return frequencies[positive], power[positive] + + +def log10_rmse_between_curves( + x_reference: np.ndarray, + y_reference: np.ndarray, + x_other: np.ndarray, + y_other: np.ndarray, +) -> float: + """RMS difference between two spectra in log10 space, in dex. + + The comparison curve is interpolated onto the reference wavenumbers in + log-log space, over the domain both cover. Working in logs measures + multiplicative error, so a factor-of-two miss counts the same whether it + happens at the energetic large scales or the faint small ones -- in linear + units the largest bin would decide the whole score. + """ + x_reference = np.asarray(x_reference, dtype=float) + y_reference = np.asarray(y_reference, dtype=float) + x_other = np.asarray(x_other, dtype=float) + y_other = np.asarray(y_other, dtype=float) + + reference_ok = ( + np.isfinite(x_reference) + & np.isfinite(y_reference) + & (x_reference > 0) + & (y_reference > 0) + ) + other_ok = ( + np.isfinite(x_other) & np.isfinite(y_other) & (x_other > 0) & (y_other > 0) + ) + if not reference_ok.any() or other_ok.sum() < 2: + return float("nan") + + order = np.argsort(x_other[other_ok]) + x_sorted = x_other[other_ok][order] + y_sorted = y_other[other_ok][order] + x_sorted, unique = np.unique(x_sorted, return_index=True) + y_sorted = y_sorted[unique] + if x_sorted.size < 2: + return float("nan") + + domain = reference_ok & (x_reference >= x_sorted[0]) & (x_reference <= x_sorted[-1]) + if not domain.any(): + return float("nan") + + interpolated = np.interp( + np.log10(x_reference[domain]), np.log10(x_sorted), np.log10(y_sorted) + ) + error = np.log10(y_reference[domain]) - interpolated + error = error[np.isfinite(error)] + if error.size == 0: + return float("nan") + return float(np.sqrt(np.mean(error**2))) + + +def interannual_band( + curves: list[tuple[np.ndarray, np.ndarray]], +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]: + """Geometric mean and +/- 1 log-standard-deviation band across yearly spectra. + + Averaging is done on logs because spectra span orders of magnitude: an + arithmetic mean would be pinned to the most energetic year, and a linear + standard deviation would put the lower band below zero, which cannot be + drawn on log axes. + + Returns `(x, mean, lower, upper, n_years)`, all arrays empty when no curve + is usable. + """ + cleaned = [] + for x, y in curves: + x = np.asarray(x, dtype=float) + y = np.asarray(y, dtype=float) + usable = np.isfinite(x) & np.isfinite(y) & (x > 0) & (y > 0) + if usable.sum() >= 2: + order = np.argsort(x[usable]) + cleaned.append((x[usable][order], y[usable][order])) + + empty = np.array([]) + if not cleaned: + return empty, empty, empty, empty, 0 + + x_reference = max((x for x, _ in cleaned), key=lambda arr: arr.size) + rows = np.asarray( + [ + np.interp( + np.log(x_reference), np.log(x), np.log(y), left=np.nan, right=np.nan + ) + for x, y in cleaned + ] + ) + with np.errstate(invalid="ignore"): + log_mean = np.nanmean(rows, axis=0) + log_std = ( + np.nanstd(rows, axis=0, ddof=1) + if rows.shape[0] > 1 + else np.zeros_like(log_mean) + ) + return ( + x_reference, + np.exp(log_mean), + np.exp(log_mean - log_std), + np.exp(log_mean + log_std), + rows.shape[0], + ) diff --git a/samudra/models/__init__.py b/samudra/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..de08a3554eb2e34b977e65b33e7447673614d86e --- /dev/null +++ b/samudra/models/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from .samudra import Samudra +from .samudra_mini import SamudraMini +from .samudra_multi import SamudraMulti + +__all__ = ["Samudra", "SamudraMulti", "SamudraMini"] diff --git a/samudra/models/base.py b/samudra/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..e0ee5329620731e8caf78ba7f6af36a4992146a0 --- /dev/null +++ b/samudra/models/base.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# TODO: Need to return step-wise losses for logging + +import logging + +import torch + +from samudra.constants import Boundary, Prognostic +from samudra.utils.ctx import GridContext + +logger = logging.getLogger(__name__) + +from samudra.datasets import InferenceDataset, TrainData +from samudra.utils.device import get_device +from samudra.utils.output import ModelInferenceOutput + + +class BaseModel(torch.nn.Module): + """Abstract base model for neural ocean emulators. + + Provides shared functionality for all ocean emulator models, including + residual prediction, ocean masking, and gradient detaching for multi-step + autoregressive rollouts. + """ + + def __init__( + self, + in_channels, + out_channels, + hist, + pred_residuals, + last_kernel_size, + pad, + gradient_detach_interval: int, + ) -> None: + super().__init__() + assert last_kernel_size % 2 != 0, "Cannot use even kernel sizes!" + self.in_channels = in_channels + self.out_channels = out_channels + self.N_pad = int((last_kernel_size - 1) / 2) + self.pad: str = pad + self.pred_residuals = pred_residuals + self.hist = hist + self.gradient_detach_interval = gradient_detach_interval + + def forward_once( + self, prognostic: Prognostic, boundary: Boundary, ctx: GridContext + ) -> Prognostic: + raise NotImplementedError() + + def forward( + self, + train_data: TrainData, + loss_fn=None, + ) -> torch.Tensor | list[torch.Tensor]: + outputs: list[torch.Tensor] = [] + loss = torch.tensor(torch.nan) + for step in range(len(train_data)): + if step == 0: + prog_tensor, boundary_tensor = train_data.get_initial_input() + else: + prev_output = outputs[-1] + if ( + self.gradient_detach_interval > 0 + and step % self.gradient_detach_interval == 0 + ): + prev_output = prev_output.detach() + _, boundary_tensor = train_data.get_input(step) + prog_tensor = prev_output + + decodings = self.forward_once(prog_tensor, boundary_tensor, train_data.ctx) + if self.pred_residuals: + pred = prog_tensor + decodings # Residual prediction + else: + pred = decodings # Absolute prediction + + if loss_fn is not None: + if step == 0: + loss = loss_fn( + pred, + train_data.get_label(step), + ) + else: + loss += loss_fn( + pred, + train_data.get_label(step), + ) + + outputs.append(pred) + + if loss_fn is None: + return outputs + else: + return loss + + def inference( + self, + dataset: InferenceDataset, + initial_prognostic: torch.Tensor, + steps_completed=0, + num_steps=None, + epoch=None, + ) -> ModelInferenceOutput: + # `dataset[idx]` returns `(prog, boundary, label)`. + out_shape = (num_steps, *dataset[0][-1].shape[1:]) + + pred_tensor = torch.zeros(out_shape, device=get_device()) + initial_prognostic = initial_prognostic.to(get_device()) + target_time = dataset.get_target_time(steps_completed, num_steps) + + for step in range(num_steps): + logger.info( + f"Inference [epoch {epoch}]: Rollout step {steps_completed + step} " + f"of {steps_completed + num_steps - 1}." + ) + if step == 0: + prog_tensor = initial_prognostic + boundary_tensor = dataset.get_boundary(steps_completed).to( + device=prog_tensor.device + ) + else: + prog_tensor = pred_tensor[step - 1].unsqueeze(0) + boundary_tensor = dataset.get_boundary( + steps_completed + step, + ).to(device=prog_tensor.device) + + decodings = self.forward_once(prog_tensor, boundary_tensor, dataset.ctx) + if self.pred_residuals: + pred = prog_tensor[0].to(device=get_device()) + decodings + else: + pred = decodings + + pred_tensor[step] = pred + + target_tensor = dataset.inference_target( + slice(steps_completed, steps_completed + num_steps) + ).to(device=get_device()) + + IO = ModelInferenceOutput(pred_tensor, target_tensor, target_time) + return IO diff --git a/samudra/models/modules/__init__.py b/samudra/models/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..435fb2bcdc9917d6b54fe4f4bdcb27a978eb849a --- /dev/null +++ b/samudra/models/modules/__init__.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from .activations import CappedGELU, CappedLeakyReLU, ReLU +from .blocks import ( + AvgPool, + BilinearUpsample, + ConvBlock, + ConvNeXtBlock, + CoreBlock, + CoreBlockBuilder, + MaxPool, + TransposedConvUpsample, + UpsamplingBlockBuilder, + ZonallyPeriodicBilinearUpsample, +) +from .decoder import PerceiverDecoder +from .encoder import PerceiverEncoder +from .unet_backbone import UNetBackbone + +__all__ = [ + "AvgPool", + "BilinearUpsample", + "ZonallyPeriodicBilinearUpsample", + "ConvBlock", + "ConvNeXtBlock", + "CoreBlock", + "TransposedConvUpsample", + "CappedGELU", + "CappedLeakyReLU", + "MaxPool", + "PerceiverDecoder", + "PerceiverEncoder", + "ReLU", + "UNetBackbone", +] diff --git a/samudra/models/modules/activations.py b/samudra/models/modules/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..296545deb9fdcfd7ceb6b4dc021974b49ba1e5f8 --- /dev/null +++ b/samudra/models/modules/activations.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: 2023 Matthias Karlbauer, Nathaniel Cresswell-Clay, Thorsten Kurth +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 +# SPDX-License-Identifier: MIT + +# TODO: Enable setting parameters for activation functions +import torch + + +class ReLU(torch.nn.Module): + """ + Implements a ReLU. + """ + + def __init__(self, **kwargs): + """ + :param kwargs: passed to torch.nn.ReLU + """ + super().__init__() + self.relu = torch.nn.ReLU(**kwargs) + + def forward(self, inputs): + x = self.relu(inputs) + return x + + +class CappedLeakyReLU(torch.nn.Module): + """ + Implements a ReLU with capped maximum value. + """ + + def __init__(self, cap_value=10.0, **kwargs): + """ + :param cap_value: float: value at which to clip activation + :param kwargs: passed to torch.nn.LeadyReLU + """ + super().__init__() + self.relu = torch.nn.LeakyReLU(**kwargs) + self.cap = torch.nn.Buffer(torch.tensor(cap_value, dtype=torch.float32)) + + def forward(self, inputs): + x = self.relu(inputs) + x = torch.clamp(x, max=self.cap) + return x + + +class CappedGELU(torch.nn.Module): + """ + Implements a ReLU with capped maximum value. + """ + + def __init__(self, cap_value=10.0, **kwargs): + """ + :param cap_value: float: value at which to clip activation + :param kwargs: passed to torch.nn.LeadyReLU + """ + super().__init__() + self.gelu = torch.nn.GELU(**kwargs) + self.cap = torch.nn.Buffer(torch.tensor(cap_value, dtype=torch.float32)) + + def forward(self, inputs): + x = self.gelu(inputs) + x = torch.clamp(x, max=self.cap) + return x diff --git a/samudra/models/modules/augment_input.py b/samudra/models/modules/augment_input.py new file mode 100644 index 0000000000000000000000000000000000000000..cde02725267c00c8b57fb4010aae47e66671f732 --- /dev/null +++ b/samudra/models/modules/augment_input.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import math + +import torch +from aurora.model.posencoding import lat_lon_meshgrid +from einops import rearrange +from jaxtyping import Float +from torch import nn + +from samudra.constants import Lat, Lon + + +def make_3d_coordinate_grid(lat: Lat, lon: Lon) -> Float[torch.Tensor, "3 H W"]: + """Make 3D Cartesian coordinates on a unit sphere. + + Returns: + ``(3, H, W)`` tensor of ``(x, y, z)`` unit-sphere coordinates. + """ + lat_lon_grid = lat_lon_meshgrid(lat, lon) # [2, H, W] + lat_rad = torch.deg2rad(lat_lon_grid[0]) # [H, W] + lon_rad = torch.deg2rad(lat_lon_grid[1]) # [H, W] + + x = torch.cos(lat_rad) * torch.cos(lon_rad) + y = torch.cos(lat_rad) * torch.sin(lon_rad) + z = torch.sin(lat_rad) + + return torch.stack([x, y, z], dim=0).float() # [3, H, W] + + +class Concat3dCoordinates(nn.Module): + """Add 3d Cartesian Coordinates on a unit sphere to the channel dimension. + + 3D coordinates are structured like so: + x y z + (0, 0, 0) is earth center. + (1, 0, 0) is at lat, lon = (0, 0) + (0, 1, 0) is at lat, lon = (0, 90) + (0, 0, 1) is at the North Pole + + This is known to provide better pole handling than raw lat/lon coordinates, see [1]. + + > Note: This module assumes that the data at each lat/lon is located at the center of each + > grid point! Please ensure this is the case during pre-processing. + + Args: + lat: A vector of latitudes representing the center of the grid point. + lon: A vector of longitudes representing the center of the grid point. + + References: + [1]: https://ar5iv.labs.arxiv.org/html/2410.07472v1#S4.SS9 + """ + + def forward( + self, + fts: Float[torch.Tensor, "batch channel height width"], + resolution: tuple[Lat, Lon], + ) -> Float[torch.Tensor, "batch channel+3 height width"]: + grid = ( + make_3d_coordinate_grid(*resolution) + .unsqueeze(0) + .to(fts.device) + .expand(fts.shape[0], -1, -1, -1) + ) + return torch.cat((fts, grid), dim=1) + + +def fourier_features_2d_dim(num_freq_bands: int) -> int: + """Channel-dim contribution of `FourierFeatures2D(num_freq_bands)`.""" + return 2 * (2 * num_freq_bands + 1) + + +class FourierFeatures2D(nn.Module): + """Concatenate 2D Fourier positional features along the channel dim. + + Matches the `fourier_encode_data=True, input_axis=2` layout from + perceiver-pytorch's Perceiver. Configured encoder Perceiver paths use this + module explicitly so flash and naive implementations encode intra-patch + position equivalently. + + Frequency layout matches `perceiver_pytorch.fourier_encode`: scales are + `linspace(1., max_freq / 2, num_freq_bands)`, applied to positions + in [-1, 1]. + + Input: (..., ph, pw, V) + Output: (..., ph, pw, V + fourier_features_2d_dim(num_freq_bands)) + """ + + def __init__(self, num_freq_bands: int, max_freq: float) -> None: + super().__init__() + self.num_freq_bands = num_freq_bands + self.max_freq = max_freq + + def forward(self, x: torch.Tensor) -> torch.Tensor: + ph, pw = x.shape[-3], x.shape[-2] + ph_pos = torch.linspace(-1.0, 1.0, steps=ph, device=x.device, dtype=x.dtype) + pw_pos = torch.linspace(-1.0, 1.0, steps=pw, device=x.device, dtype=x.dtype) + # (ph, pw, 2) + pos = torch.stack(torch.meshgrid(ph_pos, pw_pos, indexing="ij"), dim=-1) + + scales = torch.linspace( + 1.0, + self.max_freq / 2, + self.num_freq_bands, + device=x.device, + dtype=x.dtype, + ) + # (ph, pw, 2, 1) * (num_bands,) -> (ph, pw, 2, num_bands) + scaled = pos.unsqueeze(-1) * scales * math.pi + enc = torch.cat([scaled.sin(), scaled.cos(), pos.unsqueeze(-1)], dim=-1) + # Flatten (axis, freq) into a single feature dim. + enc = rearrange(enc, "ph pw ax feat -> ph pw (ax feat)") + + # Broadcast across leading batch dims; concat along channel dim. + leading = x.shape[:-3] + enc = enc.expand(*leading, *enc.shape) + return torch.cat((x, enc), dim=-1) diff --git a/samudra/models/modules/blocks.py b/samudra/models/modules/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..8c9435ada90e8d88855b472d0d5a205556385dbb --- /dev/null +++ b/samudra/models/modules/blocks.py @@ -0,0 +1,379 @@ +# SPDX-FileCopyrightText: 2023 Matthias Karlbauer, Nathaniel Cresswell-Clay, Thorsten Kurth +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 +# SPDX-License-Identifier: MIT + +from collections.abc import Callable +from typing import Protocol + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from jaxtyping import Float + +from samudra.models.modules.activations import CappedGELU + + +class PointwiseLinear(torch.nn.Module): + """A 1×1 convolution implemented as nn.Linear. + + Mathematically equivalent to Conv2d(kernel_size=1), but avoids the + non-contiguous gradient strides that 1×1 convs produce, which cause + DDP to copy gradients instead of using zero-copy views. + + This optimization is use in the official ConvNext implementation[0]. + + [0]: https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py#L18 + """ + + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.linear = torch.nn.Linear(in_channels, out_channels) + + def forward( + self, x: Float[torch.Tensor, "B C_in H W"] + ) -> Float[torch.Tensor, "B C_out H W"]: + return self.linear(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + +class TransposedConvUpsample(torch.nn.Module): + def __init__( + self, + in_channels: int = 3, + out_channels: int = 1, + upsampling: int = 2, + activation: Callable[[], torch.nn.Module] = CappedGELU, + ): + super().__init__() + upsampler: list[torch.nn.Module] = [] + # Upsample transpose conv + upsampler.append( + torch.nn.ConvTranspose2d( + in_channels, + out_channels, + kernel_size=upsampling, + stride=upsampling, + padding=0, + ) + ) + + if activation is not None: + upsampler.append(activation()) + self.upsampler = torch.nn.Sequential(*upsampler) + + def forward(self, x): + return self.upsampler(x) + + +class BilinearUpsample(torch.nn.Module): + def __init__(self, upsampling: int = 2, **kwargs): + super().__init__() + self.upsampler = torch.nn.Upsample(scale_factor=upsampling, mode="bilinear") + + def forward(self, x): + return self.upsampler(x) + + +class ZonallyPeriodicBilinearUpsample(torch.nn.Module): + """Bilinear upsampling that enforces periodicity along the x/longitude axis.""" + + def __init__(self, upsampling: int | tuple[int, int] = 2): + super().__init__() + if isinstance(upsampling, int): + upsampling = (upsampling, upsampling) + if tuple(upsampling) != (2, 2): + raise ValueError( + "ZonallyPeriodicBilinearUpsample only supports 2x upsampling" + ) + self.scale_h, self.scale_w = upsampling + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Upsample with periodic padding along longitude to avoid seams and + # keep interpolation aligned with PyTorch's bilinear sampling grid. + width = x.shape[-1] + padded = torch.nn.functional.pad(x, (1, 1, 0, 0), mode="circular") + upsampled = torch.nn.functional.interpolate( + padded, + scale_factor=(self.scale_h, self.scale_w), + mode="bilinear", + align_corners=False, + ) + # Crop out the extra padded columns (scaled by the upsampling factor). + start = self.scale_w + end = start + width * self.scale_w + return upsampled[..., start:end] + + +class DropPath(torch.nn.Module): + """Drop path dropout (for skip connections). + + During training, randomly drops entire samples' skip connections + with probability ``drop_prob``, scaling survivors by 1/(1-p) to preserve + expected values. Implemented via ``nn.Dropout`` applied to a per-sample + mask of ones. + + References: + [0]: Rethinking U-net Skip Connections for Biomedical Image Segmentation (https://arxiv.org/abs/2402.08276) + [1]: Dropout Reduces Underfitting (https://arxiv.org/abs/2303.01500) + """ + + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.dropout = torch.nn.Dropout(p=drop_prob) + + def forward( + self, skip_conn: Float[torch.Tensor, "B C H W"] + ) -> Float[torch.Tensor, "B C H W"]: + if not self.training or self.dropout.p == 0.0: + return skip_conn + # Per-sample mask: (B, 1, 1, 1) broadcasts over C, H, W. + mask = self.dropout( + torch.ones( + skip_conn.shape[0], + 1, + 1, + 1, + device=skip_conn.device, + dtype=skip_conn.dtype, + ) + ) + return skip_conn * mask + + +class AvgPool(torch.nn.Module): + def __init__( + self, + pooling: int = 2, + ): + super().__init__() + self.avgpool = torch.nn.AvgPool2d(pooling) + + def forward(self, x): + return self.avgpool(x) + + +class MaxPool(torch.nn.Module): + def __init__( + self, + pooling: int = 2, + ): + super().__init__() + self.maxpool = torch.nn.MaxPool2d(pooling) + + def forward(self, x): + return self.maxpool(x) + + +class CoreBlock(torch.nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + dilation: int, + pad: str, + upscale_factor: int = 1, + norm: str = "batch", + ): + super().__init__() + assert kernel_size % 2 != 0, "Cannot use even kernel sizes!" + self.in_channels = in_channels + self.out_channels = out_channels + + self.N_in = in_channels + self.N_pad = int((kernel_size + (kernel_size - 1) * (dilation - 1) - 1) / 2) + self.pad = pad + self.upscale_factor = upscale_factor + self.norm = norm + + def forward(self, fts: torch.Tensor) -> torch.Tensor: + raise NotImplementedError() + + +class ConvBlock(CoreBlock): + def __init__( + self, + in_channels: int = 300, + out_channels: int = 1, + kernel_size: int = 3, + dilation: int = 1, + n_layers: int = 1, + activation: Callable[[], torch.nn.Module] = CappedGELU, + pad="circular", + checkpoint_simple: bool = False, + ): + super().__init__(in_channels, out_channels, kernel_size, dilation, pad) + + layers: list[torch.nn.Module] = [] + layers.append( + torch.nn.Conv2d(in_channels, out_channels, kernel_size, dilation=dilation) + ) + layers.append(torch.nn.BatchNorm2d(out_channels)) + layers.append(activation()) + for _ in range(n_layers - 1): + layers.append( + torch.nn.Conv2d( + out_channels, out_channels, kernel_size, dilation=dilation + ) + ) + layers.append(torch.nn.BatchNorm2d(out_channels)) + layers.append(activation()) + + self.layers = nn.ModuleList(layers) + self.checkpoint_simple = checkpoint_simple + + def forward(self, fts: torch.Tensor) -> torch.Tensor: + for layer in self.layers: + if isinstance(layer, nn.Conv2d): + fts = torch.nn.functional.pad( + fts, (self.N_pad, self.N_pad, 0, 0), mode=self.pad + ) + fts = torch.nn.functional.pad( + fts, (0, 0, self.N_pad, self.N_pad), mode="constant" + ) + # conv2d layers are expensive so we save their activations, + # other (simple) layers are cheap, so we don't save their activations. + if self.checkpoint_simple and not isinstance(layer, nn.Conv2d): + fts = torch.utils.checkpoint.checkpoint(layer, fts, use_reentrant=False) + else: + fts = layer(fts) + return fts + + +def _pointwise(use_linear: bool, in_ch: int, out_ch: int) -> torch.nn.Module: + """Create a pointwise (1×1) channel-mixing layer. + + When ``use_linear`` is True, returns a :class:`PointwiseLinear` backed by + ``nn.Linear``, which produces 2-D weight tensors and avoids the + non-contiguous gradient strides that ``Conv2d(kernel_size=1)`` introduces + for degenerate spatial dimensions. This matters for DDP, which otherwise + falls back to copying gradients instead of using zero-copy views. + + The ``nn.Linear`` approach is also used in the official ConvNeXt + implementation, where it is noted to be "slightly faster in PyTorch" [0]. + + [0]: https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py + """ + if use_linear: + return PointwiseLinear(in_ch, out_ch) + return torch.nn.Conv2d(in_ch, out_ch, kernel_size=1, padding="same") + + +class ConvNeXtBlock(CoreBlock): + """ + A convolution block as reported in https://github.com/CognitiveModeling/dlwp-hpx/blob/main/src/dlwp-hpx/dlwp/model/modules/blocks.py. + + This is a modified version of the actual ConvNextblock which is used in the HealPix + paper. Use of dilations here. + + """ + + def __init__( + self, + in_channels: int = 300, + out_channels: int = 1, + kernel_size: int = 3, + dilation: int = 1, + n_layers: int = 1, + activation: Callable[[], torch.nn.Module] = CappedGELU, + pad="circular", + upscale_factor: int = 4, + norm="batch", + checkpoint_simple: bool = False, + pointwise_linear: bool = False, + ): + super().__init__(in_channels, out_channels, kernel_size, dilation, pad) + assert n_layers == 1, "Can only use a single layer here!" + + # Instantiate pointwise linear to increase/decrease channel depth if necessary + if in_channels == out_channels: + self.skip_module = lambda x: x # Identity-function required in forward pass + else: + self.skip_module = _pointwise(pointwise_linear, in_channels, out_channels) + + # Convolution block + convblock: list[torch.nn.Module] = [] + convblock.append( + torch.nn.Conv2d( + in_channels=in_channels, + out_channels=int(in_channels * upscale_factor), + kernel_size=kernel_size, + dilation=dilation, + ) + ) + # BatchNorm + if norm == "batch": + convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor)) + # Instance Norm + elif norm == "instance": + convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor)) + elif norm == "nonorm": + pass + else: + raise NotImplementedError + if activation is not None: + convblock.append(activation()) + convblock.append( + torch.nn.Conv2d( + in_channels=int(in_channels * upscale_factor), + out_channels=int(in_channels * upscale_factor), + kernel_size=kernel_size, + dilation=dilation, + ) + ) + # BatchNorm + if norm == "batch": + convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor)) + # Instance Norm + elif norm == "instance": + convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor)) + elif norm == "nonorm": + pass + else: + raise NotImplementedError + if activation is not None: + convblock.append(activation()) + # Linear postprocessing + convblock.append( + _pointwise( + pointwise_linear, int(in_channels * upscale_factor), out_channels + ) + ) + self.convblock = torch.nn.Sequential(*convblock) + self.checkpoint_simple = checkpoint_simple + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # return self.skip_module(x) + self.convblock(x) + skip = self.skip_module(x) + for layer in self.convblock: + if isinstance(layer, nn.Conv2d) and layer.kernel_size[0] != 1: + x = torch.nn.functional.pad( + x, (self.N_pad, self.N_pad, 0, 0), mode=self.pad + ) + x = torch.nn.functional.pad( + x, (0, 0, self.N_pad, self.N_pad), mode="constant" + ) + if self.checkpoint_simple and not isinstance(layer, nn.Conv2d): + x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False) + else: + x = layer(x) + return skip + x + + +class CoreBlockBuilder(Protocol): + def __call__( + self, + in_channels: int, + out_channels: int, + dilation: int, + n_layers: int, + pad: str, + checkpoint_simple: bool, + ) -> CoreBlock: ... + + +class UpsamplingBlockBuilder(Protocol): + def __call__( + self, in_channels: int, out_channels: int + ) -> BilinearUpsample | TransposedConvUpsample: ... diff --git a/samudra/models/modules/decoder.py b/samudra/models/modules/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..94bddc03a7c167d0edf54fd009292ccf1825d607 --- /dev/null +++ b/samudra/models/modules/decoder.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# Perceiver-based decoder, complementary to encoder.py + +import torch +import torch.nn.functional as F +from aurora.model.fourier import pos_expansion, scale_expansion +from aurora.model.posencoding import pos_scale_enc +from einops import rearrange +from jaxtyping import Float +from torch import nn + +from samudra.constants import Lat, Lon +from samudra.models.modules.augment_input import make_3d_coordinate_grid +from samudra.models.modules.encoder import patch_from + + +class PerceiverDecoder(nn.Module): + """A PerceiverIO-based decoder that maps a latent patch grid to full-resolution output. + + All ``nh * nw`` pos/scale-encoded latent tokens are passed as **data** to + the PerceiverIO[2], and every output pixel position is a **query**. Each + query cross-attends to the full latent representation, giving it global + spatial context — pixels near patch boundaries can attend to neighboring + patches, and the model can learn smooth inter-patch transitions. + + Concretely: + + 1. Add Aurora-style pos/scale encoding to the ``nh * nw`` latent tokens + (telling the model *where on the globe* each patch is). + 2. Pass all encoded latents as **data** to the PerceiverIO: + ``(B, nh * nw, C)``. + 3. Build 3D unit-sphere **queries** ``(x, y, z)`` for every output pixel + from its lat/lon, embed them via a learned linear layer, and feed + them to the PerceiverIO decoder head. + 4. Inside the PerceiverIO: + a. Internal latents cross-attend to the ``nh * nw`` data tokens. + b. The latents refine through several rounds of self-attention. + c. A final cross-attention maps from queries to the refined latents, + producing ``(B, H * W, out_channels)``. + 5. Reshape to ``(B, out_channels, H, W)``. + + **Spatial windowing**: When ``window_patches`` is set, the latent grid + must be evenly divisible by ``window_patches``. The grid is padded — + circular along longitude (so windows near lon=0 see context from + lon≈360) and constant-zero along latitude (poles are true boundaries) + — then ``Tensor.unfold`` extracts fixed-size overlapping windows. + Each block's PerceiverIO call receives the local data context plus + the corresponding pixel queries. Setting ``context_patches=None`` + gives each window full access to all latent tokens (windowed queries, + global data). + + Because pixel queries are unit-sphere coordinates — continuous values + determined by lat/lon, not grid indices — the same PerceiverIO + generalizes across resolutions. + + Args: + in_channels: Number of input channels from the processor. + out_channels: Number of output channels per pixel. + patch_extent: Spatial extent of each patch in degrees (lat, lon). + Used for computing positional and scale encodings on latent tokens. + queries_dim: Embedding dimension for pixel-position queries. + perceiver_io: A PerceiverIO module. ``dim`` must equal ``in_channels``, + ``queries_dim`` must match this decoder's ``queries_dim``, and + ``logits_dim`` must equal ``out_channels``. + window_patches: Side length (in patches) of each spatial decode window. + If ``None``, all patches are used globally (no windowing). + E.g. ``window_patches=8`` means each PerceiverIO call covers an + 8x8 block of patches. + context_patches: Number of extra patch rings around each window to + include as data context. Only used when ``window_patches`` is set. + Default 1 gives each window one ring of neighboring patches beyond + its own block. ``None`` means full context — every window sees all + latent tokens (windowed queries but global data attention). + + References: + [0]: https://github.com/lucidrains/perceiver-pytorch + [1]: https://ar5iv.labs.arxiv.org/html/2405.13063#A2.SS4 + [2]: https://ar5iv.labs.arxiv.org/html/2107.14795 + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + patch_extent: tuple[float, float], + queries_dim: int, + perceiver_io: nn.Module, + window_patches: int | None, + context_patches: int | None, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.patch_extent = patch_extent + if window_patches is None and context_patches is not None: + raise ValueError( + "window_patches must be set in order for context_patches to be set." + ) + self.window_patches = window_patches + self.context_patches = context_patches + + # TODO(#451): The input to these position and scale linear units could be a hparam. + # Same pos/scale linear layers as the encoder, but applied *before* the + # perceiver (the encoder applies them after). + self.pos_embed = nn.Linear(in_channels, in_channels) + self.scale_embed = nn.Linear(in_channels, in_channels) + + # Embed 3D unit-sphere coordinates into queries_dim for the PerceiverIO decoder head. + self.query_embed = nn.Linear(3, queries_dim) + + self.perceiver_io = perceiver_io + + def forward( + self, + x: Float[torch.Tensor, "batch channels nh nw"], + resolution: tuple[Lat, Lon], + ) -> Float[torch.Tensor, "batch {self.out_channels} H W"]: + # nh, nw: number of patches along height and width (the latent grid dims). + B, C, nh, nw = x.shape + lat, lon = resolution + + H, W = len(lat), len(lon) + + pos_patch_h, pos_patch_w = patch_from(self.patch_extent, H, W) + + # --- Add pos/scale encoding to latent tokens (before perceiver, unlike encoder) --- + tokens = rearrange(x, "b c nh nw -> b (nh nw) c") + + pos_encode, scale_encode = pos_scale_enc( + C, + lat, + lon, + (pos_patch_h, pos_patch_w), + pos_expansion=pos_expansion, + scale_expansion=scale_expansion, + ) + pos_encoding = self.pos_embed( + pos_encode.to(dtype=tokens.dtype, device=tokens.device) + ).unsqueeze(0) + scale_encoding = self.scale_embed( + scale_encode.to(dtype=tokens.dtype, device=tokens.device) + ).unsqueeze(0) + tokens = tokens + pos_encoding + scale_encoding + + # --- Build global pixel-position queries --- + # 3D unit-sphere coordinates for every output pixel. + coords = make_3d_coordinate_grid(lat, lon) # (3, H, W) + coords = rearrange(coords, "d h w -> h w d").to( + dtype=x.dtype, device=x.device + ) # (H, W, 3) + queries = self.query_embed( + rearrange(coords, "h w d -> (h w) d") + ) # (H*W, queries_dim) + queries = rearrange( + queries, "(h w) d -> h w d", h=H, w=W + ) # (H, W, queries_dim) + + # --- Decode via PerceiverIO with optional spatial windowing --- + data_grid = rearrange(tokens, "b (nh nw) c -> b nh nw c", nh=nh, nw=nw) + out = self._decode(data_grid, queries, pos_patch_h, pos_patch_w) + + return out + + def _decode( + self, + data_grid: Float[torch.Tensor, "batch nh nw channels"], + queries_grid: Float[torch.Tensor, "H W queries_dim"], + patch_h: int, + patch_w: int, + ) -> Float[torch.Tensor, "batch {self.out_channels} H W"]: + """Decode a latent patch grid into full-resolution pixel output. + + Without windowing, every pixel query attends to every latent token + (global attention). With windowing, the grid is split into spatial + blocks so each PerceiverIO call only covers a local neighborhood, + keeping cost bounded for large latent grids. + """ + B, nh, nw, C = data_grid.shape + H, W, _ = queries_grid.shape + + if self.window_patches is None: + data = rearrange(data_grid, "b nh nw c -> b (nh nw) c") + queries = rearrange(queries_grid, "h w d -> (h w) d") + out = self.perceiver_io(data, queries=queries) # (B, H*W, out_channels) + return rearrange(out, "b (h w) c -> b c h w", h=H, w=W) + + wp = self.window_patches + cp = self.context_patches + + assert nh % wp == 0 and nw % wp == 0, ( + f"Latent grid ({nh}, {nw}) must be divisible by window_patches={wp}" + ) + + n_blocks_h = nh // wp + n_blocks_w = nw // wp + block_ph = wp * patch_h # pixel height per query block + block_pw = wp * patch_w # pixel width per query block + + # --- Prepare data windows --- + if cp is None: + # Full context: every window sees all latent tokens. + full_data = rearrange(data_grid, "b nh nw c -> b (nh nw) c") + elif cp == 0: + # No context padding — unfold with exact window size. + data = rearrange(data_grid, "b nh nw c -> b c nh nw") + data_windows = data.unfold(2, wp, wp).unfold(3, wp, wp) + else: + # Pad: circular along longitude (last dim), zero along latitude. + data = rearrange(data_grid, "b nh nw c -> b c nh nw") + data = F.pad(data, (cp, cp, 0, 0), mode="circular") + data = F.pad(data, (0, 0, cp, cp), mode="constant", value=0) + win_size_h = wp + 2 * cp + win_size_w = wp + 2 * cp + data_windows = data.unfold(2, win_size_h, wp).unfold(3, win_size_w, wp) + # data_windows shape (when cp is not None): + # (B, C, n_blocks_h, n_blocks_w, win_h, win_w) + + # --- Decode each spatial block --- + out = data_grid.new_zeros(B, H, W, self.out_channels) + + for bi in range(n_blocks_h): + for bj in range(n_blocks_w): + if cp is None: + local_data = full_data + else: + local_data = rearrange( + data_windows[:, :, bi, bj], "b c h w -> b (h w) c" + ) + + qi_start = bi * block_ph + qj_start = bj * block_pw + local_queries = queries_grid[ + qi_start : qi_start + block_ph, + qj_start : qj_start + block_pw, + ] + local_queries = rearrange(local_queries, "h w d -> (h w) d") + + local_out = self.perceiver_io(local_data, queries=local_queries) + local_out = rearrange( + local_out, "b (h w) c -> b h w c", h=block_ph, w=block_pw + ) + out[ + :, + qi_start : qi_start + block_ph, + qj_start : qj_start + block_pw, + :, + ] = local_out + + return rearrange(out, "b h w c -> b c h w") diff --git a/samudra/models/modules/encoder.py b/samudra/models/modules/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..19c649330a50b91dd78671b9082fd727314d3332 --- /dev/null +++ b/samudra/models/modules/encoder.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# Sources inspired by the following implementations: +# - https://github.com/microsoft/aurora/blob/main/aurora/model/patchembed.py +# - https://github.com/microsoft/aurora/blob/main/aurora/model/encoder.py +# - https://github.com/lucidrains/vit-pytorch + +import torch +from aurora.model.fourier import pos_expansion, scale_expansion +from aurora.model.posencoding import pos_scale_enc +from einops import rearrange +from jaxtyping import Float +from torch import nn + +from samudra.constants import Input, Lat, Lon + + +def patch_from( + patch_extent: tuple[float, float], input_height: int, input_width: int +) -> tuple[int, int]: + """Calculate the patch size in lat/lng pixels (or coords) from the patch spatial extent and input grid size.""" + lat_spacing = 180.0 / input_height # Full sphere is 180 degrees (pole to pole) + lon_spacing = 360.0 / input_width # Full circle is 360 degrees + + # Calculate patch size to match target extent + patch_h = int(round(patch_extent[0] / lat_spacing)) + patch_w = int(round(patch_extent[1] / lon_spacing)) + + return patch_h, patch_w + + +class PerceiverEncoder(nn.Module): + """A perceiver-based encoder for Samudra's flattened data (a whole column of the ocean, with history). + + We adopt some of Aurora's positional encodings[1], which uses log-spaced fourier features with geometry-informed + wavelengths. These encode 2d positions (the average latitude and longitude of each patch) as well as grid cell area + (measured in km^2) for each token before it enters the processor. + + > Note: We assume that data along the lat/lon coordinates are positioned at the center of each grid point! Please + > ensure this is the case at the data processing time. + + This encoder is designed to make the same number of patches with the same spatial extents across different scales + of input data (input data may vary in resolution of lat/lng grid). To accomplish this with a single perceiver model, + our `forward` call requires supplementary information: the resolution (a pair of Lat/Lon tensors), which is used to + make consistent positional encodings for patches across different scales. While higher resolution scales will + contain more data per patch, the patch will refer to the same physical area on Earth as all other scales. + + Args: + in_channels (int): the number of input channels (roughly: time x variable x (surface + depths)). + out_channels (int): size of the latent dimension (aka, the embedding dimension). + patch_extent (tuple[float, float]): spatial extent of each patch measured in degrees of lat/lon. + perceiver (nn.Module): the perceiver module implementation to use. + + References: + [1]: https://ar5iv.labs.arxiv.org/html/2405.13063#A2.SS4 + """ + + # TODO(alxmrs): Implement gradient checkpointing + def __init__( + self, + in_channels: int, + out_channels: int, + patch_extent: tuple[float, float], + perceiver: nn.Module, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels: int = out_channels # aka, `embed_dim`. + self.patch_extent = patch_extent + self.perceiver = perceiver + # TODO(#451): The input to these position and scale linear units could be a hparam. + self.pos_embed = nn.Linear(self.out_channels, self.out_channels) + self.scale_embed = nn.Linear(self.out_channels, self.out_channels) + + def forward( + self, x: Input, resolution: tuple[Lat, Lon] + ) -> Float[torch.Tensor, "batch {self.embed_dim} h w"]: + _, V, H, W = x.shape + lat, lon = resolution + patch_h, patch_w = patch_from(self.patch_extent, H, W) + # V is a cross product of variable, level (encoded in vars), and time (has history). + assert V == self.in_channels + # Ensure patch_size is appropriate for the data. + assert H % patch_h == 0, f"{H} % {patch_h} != 0." + assert W % patch_w == 0, f"{W} % {patch_w} != 0." + + # Perceiver experiment ideas: + # 1. leave it as it is: treating each pixel as a token -- i.e. all channels (includes depths) per pixel + # 2. change to original plan, where each float is its own token + # 3. Add a third dim -- ph pw d v -- so each spatial position is a token + x = rearrange( + x, + "b v (h ph) (w pw) -> (b h w) ph pw v", + ph=patch_h, + pw=patch_w, + ) + # NB(alxmrs): This is includes a mean and LayerNorm before linear projection! + x = self.perceiver(x) # (B_H_W, ..., V) -> (B_H_W, out_channels) + + # Make `x` amenable to adding position + scale encoding + x = rearrange( + x, + "(b h w) l -> b (h w) l ", + h=(H // patch_h), + w=(W // patch_w), + ) + + # Calculate and add positional + scale encoding + pos_encode, scale_encode = pos_scale_enc( + self.out_channels, # aka "embed_dim" + lat, + lon, + (patch_h, patch_w), + # TODO(#452): Pos and scale wavelengths range all the way to the whole Earth by default; we could probably + # better tune these for our Oceans modeling use case. + pos_expansion=pos_expansion, + scale_expansion=scale_expansion, + ) + pos_encoding = self.pos_embed( + pos_encode.to(dtype=x.dtype, device=x.device) + ).unsqueeze(0) + scale_encoding = self.scale_embed( + scale_encode.to(dtype=x.dtype, device=x.device) + ).unsqueeze(0) + x = x + pos_encoding + scale_encoding + + # Unpack spatial channels, move channel dimension to correct location. + x = rearrange( + x, + "b (h w) l -> b l h w", + h=(H // patch_h), + w=(W // patch_w), + ) + + return x diff --git a/samudra/models/modules/unet_backbone.py b/samudra/models/modules/unet_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..ef865e15e90876cb7fc40f7a7dfd4a698414bc54 --- /dev/null +++ b/samudra/models/modules/unet_backbone.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import TYPE_CHECKING, assert_never + +import numpy as np +import torch +from torch import nn + +from samudra.models.modules.blocks import ( + BilinearUpsample, + CoreBlock, + CoreBlockBuilder, + DropPath, + TransposedConvUpsample, + UpsamplingBlockBuilder, + ZonallyPeriodicBilinearUpsample, +) +from samudra.utils.train import pairwise + +if TYPE_CHECKING: + from samudra.config import Checkpointing # noqa: F401 + + +class UNetBackbone(nn.Module): + """A configurable, convolutional or ConvNeXt[1] U-Net[2] implementation. + + Args: + ch_width (list[int]): The widths of CNN input channels going down into the U-Net. This module first builds + downsampling CNN blocks before reversing the `ch_widths` to build upsampling CNN blocks. Typically, these + values should be set in monotonically non-decreasing sizes. + dilation (list[int]): List of dilation sizes for CNN blocks. See [3] for a general background. This list must + be one less than the length of `ch_widths`. + n_layers (list[int]): List of the number of CNN layers to be used in each block section of the U-Net. Typically, + this is set to all 1s. This value must match the length of `dilation`. + pad (str): The type of padding to use in all CNN blocks. Passed into `torch.functional.pad`'s `mode` argument. + create_block: A factory method that creates the CoreBlocks for all CNN layers. + downsampling_block (nn.Module): A block that downsamples during the descent of the U-Net. + create_upsampling_block: A factory method that creates upsampling blocks during the ascent of the U-Net. + checkpointing (Checkpointing | None): The current mode for checkpointing (typically "all" or "simple"). None + turns checkpointing off. + + References: + [1]: https://arxiv.org/abs/2201.03545 + [2]: https://arxiv.org/abs/1505.04597 + [3]: https://github.com/vdumoulin/conv_arithmetic/#dilated-convolution-animations. + """ + + def __init__( + self, + in_channels: int, + ch_width: list[int], + dilation: list[int], + n_layers: list[int], + pad: str, + create_block: CoreBlockBuilder, + downsampling_block: nn.Module, + create_upsampling_block: UpsamplingBlockBuilder, + checkpointing: "Checkpointing | None", + drop_path_rate: float = 0.0, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels: int = ch_width[0] + + # Create local copies of config lists that will be reversed + ch_width = [in_channels] + ch_width.copy() + dilation = dilation.copy() + n_layers = n_layers.copy() + self.pad = pad + + match checkpointing: + case "all": + self.checkpoint_all = True + checkpoint_simple = False + case "simple": + self.checkpoint_all = False + checkpoint_simple = True + case None: + self.checkpoint_all = False + checkpoint_simple = False + case _: + assert_never(checkpointing) + + # going down + layers: list[nn.Module] = [] + for i, (a, b) in enumerate(pairwise(ch_width)): + # Core block + layers.append( + create_block( + in_channels=a, + out_channels=b, + dilation=dilation[i], + n_layers=n_layers[i], + pad=pad, + checkpoint_simple=checkpoint_simple, + ) + ) + # Down sampling block + layers.append(downsampling_block) + + # Middle block + layers.append( + create_block( + in_channels=b, + out_channels=b, + dilation=dilation[i], + n_layers=n_layers[i], + pad=pad, + checkpoint_simple=checkpoint_simple, + ) + ) + + # First upsampling + layers.append(create_upsampling_block(in_channels=b, out_channels=b)) + + # Reverse for upsampling path + ch_width.reverse() + dilation.reverse() + n_layers.reverse() + + # going up + for i, (a, b) in enumerate(pairwise(ch_width[:-1])): + layers.append( + create_block( + in_channels=a, + out_channels=b, + dilation=dilation[i], + n_layers=n_layers[i], + pad=pad, + checkpoint_simple=checkpoint_simple, + ) + ) + layers.append(create_upsampling_block(in_channels=b, out_channels=b)) + + # Final conv block + layers.append( + create_block( + in_channels=b, + out_channels=b, # this is the same as self.out_channels + dilation=dilation[i], + n_layers=n_layers[i], + pad=pad, + checkpoint_simple=checkpoint_simple, + ) + ) + + first_block = layers[0] + assert isinstance(first_block, CoreBlock) + self.N_pad = first_block.N_pad + self.layers = nn.ModuleList(layers) + self.num_steps = int(len(ch_width) - 1) + self.drop_path = DropPath(drop_path_rate) + + def forward(self, fts: torch.Tensor) -> torch.Tensor: + skip_inputs: list[torch.Tensor] = [] + for i in range(self.num_steps): + skip_inputs.append(torch.zeros_like(fts)) + count = 0 + for layer in self.layers: + # Circular/Globe padding + if isinstance(layer, nn.Conv2d): + fts = torch.nn.functional.pad( + fts, (self.N_pad, self.N_pad, 0, 0), mode=self.pad + ) + fts = torch.nn.functional.pad( + fts, (0, 0, self.N_pad, self.N_pad), mode="constant" + ) + + # (Maybe) apply checkpointing + if self.checkpoint_all: + fts = torch.utils.checkpoint.checkpoint(layer, fts, use_reentrant=False) # type: ignore + else: + fts = layer(fts) + + # UNet residuals logic (skip connections) + if count < self.num_steps: + if isinstance(layer, CoreBlock): + skip_inputs[count] = fts + count += 1 + elif count >= self.num_steps: + if ( + isinstance(layer, BilinearUpsample) + or isinstance(layer, TransposedConvUpsample) + or isinstance(layer, ZonallyPeriodicBilinearUpsample) + ): + crop = np.array(fts.shape[2:]) + shape = np.array( + skip_inputs[int(2 * self.num_steps - count - 1)].shape[2:] + ) + pads = shape - crop + pads = [ + pads[1] // 2, + pads[1] - pads[1] // 2, + pads[0] // 2, + pads[0] - pads[0] // 2, + ] + fts = nn.functional.pad(fts, pads) + fts += self.drop_path( + skip_inputs[int(2 * self.num_steps - count - 1)] + ) + count += 1 + + return fts diff --git a/samudra/models/samudra.py b/samudra/models/samudra.py new file mode 100644 index 0000000000000000000000000000000000000000..75763cd5afc620fd448d5a436c40c7979f1ecd43 --- /dev/null +++ b/samudra/models/samudra.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn +import torch.utils.checkpoint + +from samudra.constants import Boundary, GridSize, Prognostic +from samudra.models.base import BaseModel +from samudra.models.modules.unet_backbone import UNetBackbone +from samudra.utils.ctx import GridContext +from samudra.utils.device import autocast + + +class Samudra(BaseModel): + """Samudra ocean emulator using a ConvNeXt U-Net backbone. + + Implements the Samudra (and Samudra 2) model architecture for + single-scale ocean emulation. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + pred_residuals: bool, + last_kernel_size: int, + pad: str, + unet: UNetBackbone, + pos_channels: int, + add_3d_coordinates: nn.Module | None, + hist: int, + grid_size: GridSize, + gradient_detach_interval: int, + use_bfloat16: bool, + ): + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + hist=hist, + pred_residuals=pred_residuals, + last_kernel_size=last_kernel_size, + pad=pad, + gradient_detach_interval=gradient_detach_interval, + ) + + if pos_channels > 0: + self.positional_params = nn.Parameter(torch.empty(pos_channels, *grid_size)) + nn.init.normal_(self.positional_params, mean=0.0, std=1e-5) + else: + self.register_parameter("positional_params", None) + + self.add_3d_coordinates = add_3d_coordinates + self.unet = unet + self.decoder = nn.Conv2d(unet.out_channels, out_channels, last_kernel_size) + + self.use_bfloat16 = use_bfloat16 + + def forward_once( + self, prognostic: Prognostic, boundary: Boundary, ctx: GridContext + ) -> Prognostic: + # Samudra is a single-scale model; fuse prognostic + boundary into + # the single channel-stacked input its backbone expects. + fts = torch.cat((prognostic, boundary), dim=1) + + with autocast(enabled=self.use_bfloat16, dtype=torch.bfloat16): + if self.positional_params is not None: + pos = self.positional_params.unsqueeze(0).expand( + fts.shape[0], -1, -1, -1 + ) + fts = torch.cat([fts, pos], dim=1) + + if self.add_3d_coordinates is not None: + fts = self.add_3d_coordinates(fts, ctx.input_resolution_cpu) + + fts = self.unet(fts) + fts = torch.nn.functional.pad( + fts, (self.N_pad, self.N_pad, 0, 0), mode=self.pad + ) + fts = torch.nn.functional.pad( + fts, (0, 0, self.N_pad, self.N_pad), mode="constant" + ) + # TODO(jder): would be nice to keep inputs in bfloat16 and + # have the convolution use float32 internally & in output dtype. + fts = fts.to(torch.float32) + fts = self.decoder(fts) + + return torch.where(ctx.label_mask, fts, 0.0) diff --git a/samudra/models/samudra_mini.py b/samudra/models/samudra_mini.py new file mode 100644 index 0000000000000000000000000000000000000000..d0e625a95a30d6fcaaaacf6ef582c4d90fea317f --- /dev/null +++ b/samudra/models/samudra_mini.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import TYPE_CHECKING + +import torch +from einops import rearrange +from perceiver_pytorch.perceiver_io import PerceiverIO +from perceiver_pytorch.perceiver_pytorch import Attention, FeedForward +from torch import nn +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + apply_activation_checkpointing, +) + +from samudra.constants import Boundary, Prognostic +from samudra.models.base import BaseModel +from samudra.models.modules.augment_input import make_3d_coordinate_grid +from samudra.utils.ctx import GridContext +from samudra.utils.device import autocast + +if TYPE_CHECKING: + from samudra.config import Checkpointing + +_checkpoint_types: tuple[type, ...] = ( + nn.LayerNorm, + FeedForward, + nn.Linear, + PerceiverIO, + Attention, +) + +try: + from flash_attn.modules.block import ( + Block as FlashBlock, # type: ignore[import-not-found] + ) + from flash_perceiver.perceiver import ( + PerceiverBase as FlashPerceiverBase, # type: ignore[import-not-found] + ) + + _checkpoint_types = _checkpoint_types + (FlashPerceiverBase, FlashBlock) +except ImportError: + pass + + +class CoordinateEmbedding(nn.Module): + """Learned embedding for 3D Cartesian coordinates.""" + + def __init__(self, hidden_dim: int, out_dim: int) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(3, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, out_dim), + ) + + def forward(self, coords: torch.Tensor) -> torch.Tensor: + return self.net(coords) + + +class SamudraMini(BaseModel): + """Single PerceiverIO model using one token per lat/lon pixel. + + The input data is flattened to pixel tokens, and both data/query position + information is represented via learned embeddings of 3D unit-sphere + coordinates (x, y, z). + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + pred_residuals: bool, + last_kernel_size: int, + pad: str, + input_embedding_dim: int, + coordinate_embedding_dim: int, + queries_dim: int, + query_chunk_size: int | None, + perceiver_io: nn.Module, + hist: int, + checkpointing: "Checkpointing | None", + gradient_detach_interval: int, + use_bfloat16: bool, + ) -> None: + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + hist=hist, + pred_residuals=pred_residuals, + last_kernel_size=last_kernel_size, + pad=pad, + gradient_detach_interval=gradient_detach_interval, + ) + if query_chunk_size is not None and query_chunk_size <= 0: + raise ValueError( + f"query_chunk_size must be positive when set, got {query_chunk_size}." + ) + + self.input_embed = nn.Linear(in_channels, input_embedding_dim) + self.data_coordinate_embed = CoordinateEmbedding( + coordinate_embedding_dim, input_embedding_dim + ) + self.query_coordinate_embed = CoordinateEmbedding( + coordinate_embedding_dim, queries_dim + ) + self.perceiver_io = perceiver_io + self.query_chunk_size = query_chunk_size + self.use_bfloat16 = use_bfloat16 + + if checkpointing == "all": + apply_activation_checkpointing( + self, + check_fn=lambda m: isinstance(m, _checkpoint_types), + ) + + def _decode( + self, + data_tokens: torch.Tensor, + query_tokens: torch.Tensor, + ) -> torch.Tensor: + if self.query_chunk_size is None: + return self.perceiver_io(data_tokens, queries=query_tokens) + + out_chunks = [] + for start in range(0, query_tokens.shape[0], self.query_chunk_size): + end = start + self.query_chunk_size + out_chunks.append( + self.perceiver_io(data_tokens, queries=query_tokens[start:end]) + ) + return torch.cat(out_chunks, dim=1) + + def forward_once( + self, prognostic: Prognostic, boundary: Boundary, ctx: GridContext + ) -> Prognostic: + # SamudraMini is a single-scale pixel-token model; fuse prognostic + + # boundary into the single channel-stacked input it expects. + fts = torch.cat((prognostic, boundary), dim=1) + B, _, H, W = fts.shape + lat, lon = ctx.input_resolution_cpu + if H != len(lat) or W != len(lon): + raise ValueError( + "Input tensor and resolution size mismatch: " + f"tensor has {(H, W)} but resolution has {(len(lat), len(lon))}." + ) + + with autocast(enabled=self.use_bfloat16, dtype=torch.bfloat16): + # Pixel tokens: one token per (lat, lon), channels are token features. + tokens = rearrange(fts, "b c h w -> b (h w) c") + + coords = make_3d_coordinate_grid(lat, lon).to( + device=fts.device, dtype=fts.dtype + ) + coords = rearrange(coords, "d h w -> (h w) d") + + data_tokens = self.input_embed(tokens) + self.data_coordinate_embed( + coords + ).unsqueeze(0).expand(B, -1, -1) + query_tokens = self.query_coordinate_embed(coords) + + out = self._decode(data_tokens, query_tokens) + + out = out.to(torch.float32) + out = rearrange(out, "b (h w) c -> b c h w", h=H, w=W) + return torch.where(ctx.label_mask, out, 0.0) diff --git a/samudra/models/samudra_multi.py b/samudra/models/samudra_multi.py new file mode 100644 index 0000000000000000000000000000000000000000..dbcca07e25390ef983f3845df2de1dd91a481597 --- /dev/null +++ b/samudra/models/samudra_multi.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import TYPE_CHECKING + +import torch +from perceiver_pytorch import Perceiver +from perceiver_pytorch.perceiver_pytorch import Attention, FeedForward +from torch import nn +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + apply_activation_checkpointing, +) + +from samudra.constants import Boundary, Prognostic +from samudra.models.base import BaseModel +from samudra.models.modules import PerceiverDecoder, PerceiverEncoder +from samudra.models.modules.unet_backbone import UNetBackbone +from samudra.utils.ctx import GridContext +from samudra.utils.device import autocast + +if TYPE_CHECKING: + from samudra.config import Checkpointing + +_checkpoint_types: tuple[type, ...] = ( + nn.LayerNorm, + FeedForward, + nn.Linear, + Perceiver, + PerceiverDecoder, + PerceiverEncoder, + UNetBackbone, + Attention, +) + +try: + from flash_attn.modules.block import ( + Block as FlashBlock, # type: ignore[import-not-found] + ) + from flash_perceiver.perceiver import ( + PerceiverBase as FlashPerceiverBase, # type: ignore[import-not-found] + ) + + _checkpoint_types = _checkpoint_types + (FlashPerceiverBase, FlashBlock) +except ImportError: + pass + + +class SamudraMulti(BaseModel): + """Multi-resolution encoder-processor-decoder model. + + Currently, this model is used only as a physical ocean emulator. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + pred_residuals: bool, + last_kernel_size: int, + pad: str, + add_3d_coordinates: nn.Module | None, + encoder: PerceiverEncoder, + processor: UNetBackbone, + decoder: PerceiverDecoder, + hist: int, + checkpointing: "Checkpointing | None", + gradient_detach_interval: int, + use_bfloat16: bool, + ): + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + hist=hist, + pred_residuals=pred_residuals, + last_kernel_size=last_kernel_size, + pad=pad, + gradient_detach_interval=gradient_detach_interval, + ) + + self.maybe_add_3d_coordinates = add_3d_coordinates + self.encoder = encoder + self.processor = processor + self.decoder = decoder + self.use_bfloat16 = use_bfloat16 + + if checkpointing == "all": + apply_activation_checkpointing( + self, + check_fn=lambda m: isinstance(m, _checkpoint_types), + ) + + def forward_once( + self, prognostic: Prognostic, boundary: Boundary, ctx: GridContext + ) -> Prognostic: + # Prognostic and boundary are carried as separate tensors through the + # data pipeline, but this encoder still expects a single concatenated + # input. The dual-perceiver encoder that fuses them at the token level + # (enabling cross-resolution) lands in a follow-up PR. + fts = torch.cat((prognostic, boundary), dim=1) + with autocast(enabled=self.use_bfloat16, dtype=torch.bfloat16): + if self.maybe_add_3d_coordinates is not None: + fts = self.maybe_add_3d_coordinates(fts, ctx.input_resolution_cpu) + fts = self.encoder(fts, ctx.input_resolution_cpu) + fts = self.processor(fts) + + fts = self.decoder(fts, ctx.output_resolution_cpu) + + # Convert back to float32 + # TODO(alxmrs): We actually only support float16 when turned on; this kind of tricks us. + fts = fts.to(torch.float32) + + return torch.where(ctx.label_mask, fts, 0.0) diff --git a/samudra/py.typed b/samudra/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/samudra/stepper.py b/samudra/stepper.py new file mode 100644 index 0000000000000000000000000000000000000000..d1eaa3d0e495d7907590342bd43ecc3d59d1a9ff --- /dev/null +++ b/samudra/stepper.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Time-stepping primitives for training, validation, and inference. + +Provides module-level functions that handle single-step forward passes +(``train_batch``, ``validate_batch``) and multi-step autoregressive +rollouts (``run_rollout``). +""" + +import logging +from collections.abc import Callable +from functools import partial +from os import PathLike + +import torch + +from samudra.aggregator import InferenceEvaluatorAggregator +from samudra.constants import TensorMap +from samudra.datasets import InferenceDataset, TrainData +from samudra.models.base import BaseModel +from samudra.utils.data import Normalize +from samudra.utils.device import get_device +from samudra.utils.output import ModelInferenceOutput, TrainBatchOutput, ValBatchOutput +from samudra.utils.wandb import get_record_to_wandb +from samudra.utils.writer import ZarrWriter + +logger = logging.getLogger(__name__) + + +def train_batch( + model: torch.nn.Module, batch: TrainData, loss_fn: Callable +) -> TrainBatchOutput: + loss_per_channel = model(batch, loss_fn=partial(loss_fn, ctx=batch.ctx)) + loss = torch.mean(loss_per_channel) + return TrainBatchOutput(loss, loss_per_channel) + + +@torch.no_grad() +def validate_batch( + model: BaseModel | torch.nn.parallel.DistributedDataParallel, + batch: TrainData, + loss_fn: Callable, +) -> ValBatchOutput: + assert len(batch) == 1 # Assert we are using one step of input and output + prognostic, boundary = batch.get_input(0) + label = batch.get_label(0) + + outs = model(batch)[0] + loss_per_channel = loss_fn(outs, label, batch.ctx) + loss = torch.mean(loss_per_channel) + # `input_data` in ValBatchOutput is used for val visualization; pass the + # channel-concatenated tensor for continuity with existing consumers. + # We don't mix scales across boundary and prognostic during training. + input_data = torch.cat((prognostic, boundary), dim=1) + return ValBatchOutput(loss, loss_per_channel, input_data, label, outs, batch.ctx) + + +@torch.no_grad() +def run_rollout( + model: BaseModel, + dataset: InferenceDataset, + inf_aggregator: InferenceEvaluatorAggregator, + epoch: int, + output_dir: str | PathLike | None = None, + model_path: str | PathLike | None = None, + num_model_steps_forward: int = 200, + save_zarr: bool = False, + tensor_map: TensorMap | None = None, + normalize: Normalize | None = None, +) -> None: + """Performs inference, which is an auto-regressive rollout.""" + if save_zarr: + if output_dir is None or model_path is None: + raise ValueError( + "output_dir and model_path must be provided if save_zarr is True" + ) + if tensor_map is None or normalize is None: + raise ValueError( + "tensor_map and normalize must be provided if save_zarr is True" + ) + coords = dataset.get_coords_dict() + if num_model_steps_forward > 0: + chunk_size = num_model_steps_forward + else: + chunk_size = 20 + writer = ZarrWriter( + output_dir, + coords=coords, + hist=inf_aggregator.hist, + model_path=model_path, + time_chunk_size=chunk_size, + normalize=normalize, + tensor_map=tensor_map, + ) + else: + writer = None + record_logs = get_record_to_wandb(label="inference") + logger.info(f"Inference [epoch {epoch}]: processing initial prognostic.") + logs = inf_aggregator.record_initial_prognostic( + initial_prognostic=dataset.initial_prognostic.to(get_device()), + ) + record_logs(logs) + num_model_steps = len(dataset) + num_steps_list = [] + + # If num_model_steps_forward is -1, then we are doing a full forward pass + if num_model_steps_forward == -1: + num_steps_list = [num_model_steps] + else: + # Windows of partial forward passes + num_loops = num_model_steps // num_model_steps_forward + if num_loops > 0: + num_steps_list = [num_model_steps_forward] * num_loops + last_model_steps_forward = num_model_steps % num_model_steps_forward + if last_model_steps_forward > 0: + num_steps_list = num_steps_list + [last_model_steps_forward] + else: + num_steps_list = [num_model_steps] + + num_loops = len(num_steps_list) + initial_prognostic = dataset.initial_prognostic + step = 0 + for loop, num_steps in enumerate(num_steps_list): + logger.info( + f"Inference [epoch {epoch}]: loop {loop} of {num_loops - 1}. " + f"Stepping {num_steps} steps forward." + ) + dataset.to(get_device()) + IO: ModelInferenceOutput = model.inference( + dataset, + initial_prognostic=initial_prognostic, + steps_completed=step, + num_steps=num_steps, + epoch=epoch, + ) + # Setting initial prognostic for next loop + initial_prognostic = IO.prediction[-1].unsqueeze(0).clone() + if writer: + logger.info("Writing to zarr...") + writer.record_batch(IO) + writer.write() + + logger.info("Recording logs...") + logs = inf_aggregator.record_batch(IO) + logger.info("Logging to wandb...") + record_logs(logs) + step += num_steps diff --git a/samudra/train.py b/samudra/train.py new file mode 100644 index 0000000000000000000000000000000000000000..36c0b36445f19c41770e89c862a883205f14dc75 --- /dev/null +++ b/samudra/train.py @@ -0,0 +1,1123 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import datetime +import logging +import multiprocessing +import os +import tempfile +import time +import warnings +from collections import OrderedDict +from multiprocessing.context import BaseContext +from pathlib import Path +from typing import Any + +import dask +import torch +import torch.nn as nn +from torch.utils.data import ( + ConcatDataset, + DataLoader, + DistributedSampler, + RandomSampler, +) + +from samudra.aggregator import Aggregator +from samudra.aggregator.loss import ( + get_channel_loss_dict, + get_channel_loss_scale_dict, + get_depth_loss_dict, + get_variable_loss_dict, +) +from samudra.backend import init_train_backend +from samudra.config import TrainConfig, build_loss_fn +from samudra.constants import ( + MAX_TRAIN_MODEL_STEPS_FORWARD, + BoundaryVarNames, + PrognosticVarNames, + TensorMap, +) +from samudra.datasets import ( + InferenceDataset, + InferenceDatasets, + RawTrainData, + TorchTrainDataset, + TrainData, + TrainDataLoader, +) +from samudra.models.base import BaseModel +from samudra.stepper import ( + TrainBatchOutput, + ValBatchOutput, + run_rollout, + train_batch, + validate_batch, +) +from samudra.utils.data import Normalize, get_inference_steps +from samudra.utils.device import using_gpu +from samudra.utils.distributed import ( + all_reduce_mean, + get_world_size, + is_main_process, + set_seed, +) +from samudra.utils.ema import EMATracker +from samudra.utils.logging import ( + MetricLogger, + SmoothedValue, + get_model_summary, + handle_logging, + handle_warnings, +) +from samudra.utils.loss import DynamicLoss, LossFnWithContext +from samudra.utils.samplers import ( + DistributedEquivalenceGroupBatchSampler, + EquivalenceGroupBatchSampler, +) +from samudra.utils.train import ( + CheckpointPaths, + collate_inference_data, + collate_raw_train_data, +) +from samudra.utils.train_progress import TrainProgress +from samudra.utils.wandb import WandBLogger + +logger = logging.getLogger(__name__) + + +def should_log_validation_images(epoch: int, frequency: int) -> bool: + """Return whether to log validation images for a 1-based training epoch.""" + if epoch < 1: + raise ValueError(f"Epoch must be >= 1, got {epoch}") + if frequency < 1: + raise ValueError( + f"Validation image log frequency must be >= 1, got {frequency}" + ) + return (epoch - 1) % frequency == 0 + + +class Trainer: + """Orchestrates the full model training loop. + + Handles initialization, distributed setup, checkpointing, learning rate + scheduling, EMA, and Weights & Biases logging. + """ + + model: BaseModel | nn.parallel.DistributedDataParallel + + def __init__(self, cfg: TrainConfig) -> None: + cfg.prepare_output_dirs() + cfg.save_yaml(cfg.experiment.output_dir / "config.yaml") + + # Backend + self.device, self.distributed = init_train_backend(cfg.backend) + + # Adjust workers and memory pinning based on device + if not using_gpu(): + cfg.pin_mem = False + elif cfg.disk_mode: + cfg.pin_mem = True + + # Distributed mode + dask.config.set(scheduler="synchronous") + + # Set seeds + self.rand_seed = cfg.experiment.rand_seed + set_seed(self.rand_seed) + + self.data_container = cfg.data.build( + data_root=cfg.experiment.resolved_data_root, + ) + + # Getting prognostic and boundary variables + self.dataset_spec = self.data_container.dataset_spec + self.prognostic_var_names: PrognosticVarNames = ( + self.dataset_spec.prognostic_var_names + ) + self.boundary_var_names: BoundaryVarNames = self.dataset_spec.boundary_var_names + self.levels = self.dataset_spec.num_prognostic_depth_levels + + str_prognostics = ", ".join([i for i in self.prognostic_var_names]) + str_boundaries = ", ".join([i for i in self.boundary_var_names]) + + logger.info(f"Prognostic variables: {str_prognostics}") + logger.info(f"Boundary variables: {str_boundaries}") + logger.info(f"Levels: {self.levels}") + + self.N_bound = len(self.boundary_var_names) + self.N_prog = len(self.prognostic_var_names) + + data_num_workers = cfg.data.loading.num_pytorch_workers() + persistent_workers = cfg.data.loading.persistent_pytorch_workers() + + self.mp_context: BaseContext | None = None + if data_num_workers > 0: + self.mp_context = multiprocessing.get_context("spawn") + + self.num_prog_in = int((cfg.data.hist + 1) * self.N_prog) + self.num_boundary_in = int((cfg.data.hist + 1) * self.N_bound) + self.num_in = self.num_prog_in + self.num_boundary_in + self.num_out = self.num_prog_in + + self.tensor_map = TensorMap(dataset_spec=self.dataset_spec).to(self.device) + + logger.info(f"Number of inputs (prognostic + boundary): {self.num_in}") + logger.info(f"Number of outputs (prognostic): {self.num_out}") + + assert isinstance(cfg.data_stride, list) + assert isinstance(cfg.steps, list) + assert isinstance(cfg.step_transition, list) + assert len(cfg.step_transition) == len(cfg.steps) - 1 + max_steps = str(cfg.steps[-1]) + self.str_video = "steps_" + max_steps + "_" + "_Lateral_Data_025_no_smooth" + + # Dataloaders + logger.info(f"Loading data") + self.concurrent_compute = cfg.data.concurrent_compute + + self.primary_src = self.data_container.primary_source + + # We use dask for inference since it has memory issues otherwise. + # TODO(jder): Could rewrite inference dataset like we did for TorchTrainDataset + # see https://github.com/m2lines/Samudra/issues/208 + self.inference_src = self.data_container.inference_source + + self.loader_version = self.data_container.loader_version + + # Aggregation still works on the primary source only. + self.normalize = Normalize( + self.primary_src, + prognostic_var_names=self.prognostic_var_names, + boundary_var_names=self.boundary_var_names, + ) + + self.model = cfg.model.build( + prog_channels=self.num_prog_in, + boundary_channels=self.num_boundary_in, + out_channels=self.num_out, + hist=cfg.data.hist, + srcs=self.data_container.train_sources, + ).to(self.device) + + self.nets_dir = cfg.experiment.nets_dir + self.network = self.model.__class__.__name__ + + # Loss function + self.loss_fn: LossFnWithContext = build_loss_fn( + cfg.loss, + device=self.device, + num_channels=self.N_prog, + pad_mode=cfg.model.pad, + ) + + # Optimizer + self.optimizer = torch.optim.Adam(self.model.parameters(), lr=cfg.learning_rate) + + # Scheduler + self.scheduler = None + if cfg.scheduler: + self.scheduler = cfg.scheduler.build(self.optimizer, cfg.epochs) + + # Initialize WandB + self.wandb_logger = WandBLogger.init_instance() + self.wandb_logger.configure( + cfg.experiment.wandb.mode == "online", is_main_process() + ) + + self.ckpt_paths = CheckpointPaths(self.nets_dir) + + # Check for preemption + if cfg.preemptible: + assert not cfg.finetune, "Finetune is not supported with preemptible" + preempted = os.path.isfile(self.ckpt_paths.latest_checkpoint_path) + if preempted: + cfg.resume_ckpt_path = str(self.ckpt_paths.latest_checkpoint_path) + + # Set up wandb run + self.wandb_id, self.wandb_name = self.wandb_logger.setup_run( + cfg.resume_ckpt_path, + cfg, + data_container=self.data_container, + finetune=cfg.finetune, + ) + + # Log local and global batch sizes for cross-run comparisons. + self.world_size = get_world_size() + local_effective_batch_size = cfg.batch_size * cfg.gradient_accumulation_steps + global_microbatch_size = cfg.batch_size * self.world_size + global_effective_batch_size = local_effective_batch_size * self.world_size + logger.info( + f"Effective batch size: {global_effective_batch_size} " + f"(batch_size={cfg.batch_size} × " + f"gradient_accumulation_steps={cfg.gradient_accumulation_steps} × " + f"world_size={self.world_size})" + ) + if self.is_wandb_enabled(): + self.wandb_logger.log( + { + "config/effective_batch_size": global_effective_batch_size, + "config/local_batch_size": cfg.batch_size, + "config/local_effective_batch_size": local_effective_batch_size, + "config/world_size": self.world_size, + "config/global_microbatch_size": global_microbatch_size, + "config/global_effective_batch_size": global_effective_batch_size, + }, + step=0, + ) + + self.num_batches_seen = 0 + self.best_val_loss = 1e8 + self.best_inf_loss = 1e8 + self.train_progress = TrainProgress() + loaded_checkpoint = False + if cfg.resume_ckpt_path is not None: + if cfg.finetune: + self.load_checkpoint(cfg.resume_ckpt_path, finetune=True) + self.start_epoch = 1 + else: + self.load_checkpoint(cfg.resume_ckpt_path) + if not self.wandb_logger.enabled and is_main_process(): + warnings.warn( + "This checkpoint had wandb enabled, \ + but wandb is not enabled now!" + ) + loaded_checkpoint = True + else: + self.start_epoch = 1 + + # Modify DDP setup based on device + if self.distributed is not None: + self.model = nn.parallel.DistributedDataParallel( + nn.SyncBatchNorm.convert_sync_batchnorm(self.model), + device_ids=[self.distributed.gpu], + ) + + # EMA (must come after DDP setup so parameter names match final self.model) + if not loaded_checkpoint: + self._ema = EMATracker( + self.model, + decay=cfg.ema_decay, + faster_decay_at_start=cfg.faster_decay_at_start, + ) + + # Training + self.epochs = cfg.epochs + self.test_using_ema = cfg.test_using_ema + self.hist: int = cfg.data.hist + self.steps = cfg.steps + self.step_transition = cfg.step_transition + self.save_freq = cfg.save_freq + self.validation_image_log_freq = cfg.validation_image_log_freq + self.output_dir = cfg.experiment.output_dir + self.debug = cfg.debug + self.data_stride: list[int] = cfg.data_stride + self.batch_size: int = cfg.batch_size + self.gradient_accumulation_steps: int = cfg.gradient_accumulation_steps + self.num_workers: int = data_num_workers + self.persistent_workers: bool = persistent_workers + self.pin_mem: bool = cfg.pin_mem + self.inference_epochs = cfg.inference_epochs + self.max_train_model_steps_forward = MAX_TRAIN_MODEL_STEPS_FORWARD // ( + self.hist + 1 + ) + self.normalize_before_mask: bool = cfg.data.normalize_before_mask + self.normalize_fill_value: float = cfg.data.masked_fill_value + self.delayed_loss_estimate: bool = cfg.delayed_loss_estimate + + self.profiler = cfg.profiler.build(self.output_dir, self.device) + self.validation_images_enabled = self._sync_flag_from_main( + self.wandb_logger.enabled + ) + + assert self.tensor_map is not None + + if self.inference_epochs: + if self.inference_src is None: + raise ValueError( + "Inference time is not configured for the first data source" + ) + self.init_inference_stores() + + # Add type annotations for samplers + self.train_sampler: ( + EquivalenceGroupBatchSampler | DistributedEquivalenceGroupBatchSampler + ) + self.val_sampler: ( + EquivalenceGroupBatchSampler | DistributedEquivalenceGroupBatchSampler + ) + self.inference_sampler: DistributedSampler | RandomSampler + + # Add type annotations for loaders + self.train_loader: TrainDataLoader + self.val_loader: TrainDataLoader + self.inference_loader: DataLoader[TrainData] + + def init_inference_stores(self): + assert self.inference_src is not None + num_time_steps = get_inference_steps( + self.inference_src, + hist=self.hist, + ) + inference_dataset = InferenceDataset( + src=self.inference_src, + prognostic_var_names=self.prognostic_var_names, + boundary_var_names=self.boundary_var_names, + hist=self.hist, + normalize_before_mask=self.normalize_before_mask, + masked_fill_value=self.normalize_fill_value, + long_rollout=True, + ) + + inference_data_combined = InferenceDatasets( + [inference_dataset], [num_time_steps] + ) + + if self.distributed is not None: + self.inference_sampler = DistributedSampler( + inference_data_combined, shuffle=True + ) + else: + self.inference_sampler = RandomSampler(inference_data_combined) + + # Create data loaders + self.inference_loader = DataLoader( + inference_data_combined, + batch_size=1, + sampler=self.inference_sampler, + num_workers=self.num_workers, + pin_memory=False, + drop_last=False, + collate_fn=collate_inference_data, + multiprocessing_context=self.mp_context, + ) + + def run(self) -> None: + logger.info(f"Starting training") + + self.wandb_logger.watch(self.model, log="all") + + self.profiler.start() + + start_time = time.perf_counter() + for epoch in range(self.start_epoch, self.epochs + 1): + # Iterative step training + if epoch == self.start_epoch or epoch in self.step_transition: + cur_step = self.get_current_step(epoch) + self.init_data_loaders(cur_step) + + if hasattr(self.train_sampler, "set_epoch"): + self.train_sampler.set_epoch(epoch) + if hasattr(self.val_sampler, "set_epoch"): + self.val_sampler.set_epoch(epoch) + + start_epoch_train_time = time.perf_counter() + train_stats = self.train_one_epoch(epoch) + end_epoch_train_time = time.perf_counter() + val_stats = self.validate_one_epoch(epoch) + end_epoch_val_time = time.perf_counter() + + if -1 in self.inference_epochs or epoch in self.inference_epochs: + inf_stats = self.inference_one_epoch(epoch) + end_epoch_inf_time = time.perf_counter() + else: + inf_stats = {} + end_epoch_inf_time = None + + train_loss = train_stats["train/mean/loss"] + v_loss = val_stats["val/mean/loss"] + inf_loss = inf_stats.get("inference/time_mean_norm/rmse/channel_mean", None) + + logger.info(f"Achieved Train Loss = {train_loss:.3f}") + logger.info(f"Achieved Validation Loss = {v_loss:.3f}") + if inf_loss is not None: + logger.info(f"Achieved Inference Loss = {inf_loss:.3f}") + + if is_main_process(): + self.save_all_checkpoints(epoch, v_loss, inf_loss) + + time_elapsed = time.perf_counter() - start_epoch_train_time + + log_stats = { + **train_stats, + **val_stats, + **inf_stats, + "epoch": epoch, + "epoch_train_seconds": end_epoch_train_time - start_epoch_train_time, + "epoch_validation_seconds": end_epoch_val_time - end_epoch_train_time, + "epoch_total_seconds": time_elapsed, + **self.train_progress.to_metrics(), + } + + if end_epoch_inf_time is not None: + log_stats["epoch_inference_seconds"] = ( + end_epoch_inf_time - end_epoch_val_time + ) + + if is_main_process(): + self.wandb_logger.log(log_stats, step=self.num_batches_seen) + + total_time = time.perf_counter() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + logger.info(f"Training time {total_time_str}") + self.finish() + + def train_one_epoch(self, epoch): + self.model.train(True) + train_aggregator = Aggregator.get_train_aggregator(self.tensor_map) + metric_logger = MetricLogger(delimiter=" ") + metric_logger.add_meter("lr", SmoothedValue(window_size=1, fmt="{value:.6f}")) + header = f"Training Epoch: [{epoch}]" + + total_batches = len(self.train_loader) + + # Ensure gradients are zeroed at the start of the epoch so we don't + # accidentally accumulate leftovers from checkpoint/loading. + self.optimizer.zero_grad() + + # Calculate how many batches will be in the final incomplete accumulation cycle (if any) + remaining_batches = total_batches % self.gradient_accumulation_steps + final_cycle_start = ( + total_batches - remaining_batches + if remaining_batches > 0 + else total_batches + ) + + for data_iter_step, data in enumerate( + metric_logger.log_every(self.train_loader, 1, header) + ): + if self.debug and (data_iter_step + 1) % 5 == 0: + break + + in_final_cycle = ( + data_iter_step + 1 > final_cycle_start + ) and remaining_batches > 0 + + # Determine the actual number of microbatches in this accumulation cycle + if in_final_cycle: + r = remaining_batches + else: + r = self.gradient_accumulation_steps + + if self.num_batches_seen == 0: + get_model_summary(self.model, data, self.debug) + + with self.train_progress.batch( + data, world_size=self.world_size, device=self.device + ) as batch_progress: + TO: TrainBatchOutput = train_batch(self.model, data, self.loss_fn) + + # Scale loss by this accumulation cycle's actual microbatch count. + scaled_loss = TO.loss / r + scaled_loss.backward() + + train_aggregator.record_batch(TO) + + self.num_batches_seen += 1 + + is_last = data_iter_step + 1 == total_batches + should_step = ( + data_iter_step + 1 + ) % self.gradient_accumulation_steps == 0 + optimizer_stepped = should_step or is_last + batch_progress.optimizer_stepped = optimizer_stepped + # Step optimizer after accumulating enough batches or at the end + if optimizer_stepped: + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + 1.0, + error_if_nonfinite=True, + ) + self.optimizer.step() + self.optimizer.zero_grad() + self._ema(model=self.model) + + lr = ( + self.optimizer.param_groups[-1]["lr"] + if self.scheduler is None + else self.scheduler.get_last_lr()[0] + ) + + with torch.no_grad(): + # Reduce losses + loss_value_reduce = all_reduce_mean(TO.loss.detach()) + loss_per_channel_reduce = all_reduce_mean(TO.loss_per_channel.detach()) + metrics = { + "train/batch/loss": loss_value_reduce, + "train/batch/lr": lr, + "train/batch/ema_cur_decay": self._ema.cur_decay.item(), + **batch_progress.to_metrics(), + **self.train_progress.to_metrics(), + **batch_progress.to_throughput_metrics(), + **get_channel_loss_dict( + label="train", + loss_per_channel=loss_per_channel_reduce, + tensor_map=self.tensor_map, + ), + **get_depth_loss_dict( + label="train", + loss_per_channel=loss_per_channel_reduce, + tensor_map=self.tensor_map, + ), + **get_variable_loss_dict( + label="train", + loss_per_channel=loss_per_channel_reduce, + tensor_map=self.tensor_map, + ), + "train/batch/data_load_time": metric_logger.meters[ + "data_load_time" + ].value, + "train/batch/data_wait_time": metric_logger.meters[ + "data_wait_time" + ].value, + } + + if loss_scale_per_channel_fn := getattr( + self.loss_fn, "loss_scale_per_channel", None + ): + loss_scale_per_channel = loss_scale_per_channel_fn() + # Reshape from time-major channels to [hist, var] and + # average along the history dimension. + loss_per_channel = TO.loss_per_channel.reshape( + -1, loss_scale_per_channel.shape[0] + ).mean(dim=0) + + unscaled_loss_per_channel = ( + loss_per_channel / loss_scale_per_channel + ) + unscaled_loss = torch.mean(unscaled_loss_per_channel) + + metrics.update( + { + **get_channel_loss_scale_dict( + label="train", + loss_scale_per_channel=loss_scale_per_channel, + tensor_map=self.tensor_map, + ), + **get_channel_loss_dict( + label="train", + loss_per_channel=unscaled_loss_per_channel, + tensor_map=self.tensor_map, + loss_name="loss_unscaled", + ), + "train/batch/loss_unscaled": unscaled_loss, + } + ) + + if (it_time := metric_logger.meters["iter_time"]).count > 0: + metrics["train/batch/iter_time"] = it_time.value + + self.wandb_logger.log(metrics, step=self.num_batches_seen) + + metric_logger.update(loss=loss_value_reduce.item()) + metric_logger.update(lr=lr) + + self._maybe_update_loss(TO, data) + + self.profiler.after_batch(self.num_batches_seen) + + if self.scheduler is not None: + self.scheduler.step() + + logger.info(f"Aggregating train logs") + return train_aggregator.get_logs() + + def _maybe_update_loss(self, output: TrainBatchOutput, data: TrainData): + if (update := getattr(self.loss_fn, "update", None)) is None: + return + + if self.delayed_loss_estimate: + # Use the already-computed per-channel loss from the training + # rollout to update DynamicLoss scales, avoiding a second forward + # pass. This introduces a delayed estimate but is more efficient. + loss_per_channel = output.loss_per_channel + # Undo the dynamic scaling to recover the raw per-channel loss. + if get_scales := getattr(self.loss_fn, "loss_scale_per_channel", None): + per_channel_scale = get_scales() + raw_loss = ( + loss_per_channel.detach().reshape(-1, per_channel_scale.shape[0]) + / per_channel_scale + ).reshape(-1) + else: + raise RuntimeError( + "no `loss_scale_per_channel` — cannot recover unscaled per-channel loss." + ) + update(raw_loss) + else: + # Run a fresh single-step forward pass so DynamicLoss sees an + # up-to-date, unscaled loss signal + with torch.no_grad(): + single_step_data = TrainData( + data.num_prognostic_channels, data.num_boundary_channels, data.ctx + ) + prog_input, boundary_input, label = data[0] + single_step_data.append(prog_input, boundary_input, label) + pred = self.model(single_step_data) + # Compute the raw (unscaled) per-channel loss via the inner + # loss function, bypassing DynamicLoss scaling. + if not isinstance(self.loss_fn, DynamicLoss): + raise TypeError(f"Expected loss_fn to be DynamicLoss") + raw_loss = self.loss_fn.loss_fn(pred[0], label, ctx=data.ctx) + update(raw_loss) + + def validate_one_epoch(self, epoch): + self.model.eval() + log_validation_images = ( + should_log_validation_images(epoch, self.validation_image_log_freq) + and self.validation_images_enabled + ) + + val_aggregator = Aggregator.get_validation_aggregator( + self.primary_src.metadata, + self.hist, + self.primary_src.spherical_area_weights.to(self.device), + self.num_out, + self.tensor_map, + self.normalize, + include_image_aggregators=log_validation_images, + ) + metric_logger = MetricLogger(delimiter=" ") + header = f"One-Step Validation Epoch: [{epoch}]" + + with torch.no_grad(), self._test_context(): + for data_iter_step, data in enumerate( + metric_logger.log_every(self.val_loader, 1, header) + ): + if self.debug and (data_iter_step + 1) % 5 == 0: + break + + VO: ValBatchOutput = validate_batch(self.model, data, self.loss_fn) + val_aggregator.record_validation_batch(VO) + metric_logger.update(loss=VO.loss) + + logger.info(f"Aggregating validation logs") + return val_aggregator.get_logs(label="val") + + def inference_one_epoch(self, epoch): + self.model.eval() + + with torch.no_grad(), self._test_context(): + for data_iter_step, (inference_dataset, num_steps) in enumerate( + self.inference_loader + ): + # TODO(alxmrs): Aggregator only supports a single scale. + inf_aggregator = Aggregator.get_inline_inference_aggregator( + num_steps, + self.primary_src.metadata, + self.hist, + self.primary_src.spherical_area_weights.to(self.device), + self.primary_src.masks.prognostic.to(self.device), + self.num_out, + self.tensor_map, + self.normalize, + self.prognostic_var_names, + ) + + # TODO(jder): we need the underlying model so we can use forward_once; + # see https://github.com/m2lines/Samudra/issues/51 + run_rollout( + model=self.model.module + if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) + else self.model, + dataset=inference_dataset, + inf_aggregator=inf_aggregator, + epoch=epoch, + num_model_steps_forward=min( + num_steps // 2, self.max_train_model_steps_forward + ), + tensor_map=self.tensor_map, + normalize=self.normalize, + ) + + logger.info(f"Aggregating inference logs") + logs = inf_aggregator.get_summary_logs() + return {f"inference/{k}": v for k, v in logs.items()} + + def get_current_step(self, epoch): + """Determine the current step based on the epoch and transition points. + + Args: + epoch (int): Current epoch number + + Returns: + tuple: (current_step, current_step_idx) + """ + if epoch == self.start_epoch: + # Find initial step based on start epoch + cur_step = None + cur_step_idx = None + for i, epoch_to_transition in enumerate(self.step_transition): + if epoch <= epoch_to_transition: + cur_step = self.steps[i] + cur_step_idx = i + break + if cur_step is None: + cur_step = self.steps[-1] + cur_step_idx = len(self.steps) - 1 + logger.info(f"Starting training at step {cur_step}") + else: + # Transition to next step + cur_step_idx = next( + i for i, e in enumerate(self.step_transition) if e == epoch + ) + cur_step_idx += 1 + cur_step = self.steps[cur_step_idx] + logger.info(f"Transitioning to step {cur_step}") + + return cur_step + + def init_data_loaders(self, cur_step: int) -> None: + """Initialize training and validation data loaders. + + Args: + cur_step: Current training step size + """ + train_datasets = [ + TorchTrainDataset( + src=src, + prognostic_var_names=self.prognostic_var_names, + boundary_var_names=self.boundary_var_names, + hist=self.hist, + steps=cur_step, + normalize_before_mask=self.normalize_before_mask, + masked_fill_value=self.normalize_fill_value, + stride=stride, + concurrent_compute_=self.concurrent_compute, + ) + for stride in self.data_stride + for src in self.data_container.train_sources + ] + + # Validation is always evaluated on the primary source. This keeps the + # validation loss and physical-space metrics comparable across epochs, + # regardless of the set of resolutions used for training. + val_datasets = [ + TorchTrainDataset( + src=self.data_container.val_sources[0], + prognostic_var_names=self.prognostic_var_names, + boundary_var_names=self.boundary_var_names, + hist=self.hist, + steps=1, # current_step set to 1 for validation + normalize_before_mask=self.normalize_before_mask, + masked_fill_value=self.normalize_fill_value, + stride=stride, + concurrent_compute_=self.concurrent_compute, + ) + for stride in self.data_stride + ] + + # Create datasets + match self.loader_version: + case TorchTrainDataset.FLAG: + train_data: torch.utils.data.Dataset[RawTrainData] = ConcatDataset( + train_datasets + ) + + val_data: torch.utils.data.Dataset[RawTrainData] = ConcatDataset( + val_datasets + ) + + case _: + raise NotImplementedError( + f"Loader version {self.loader_version} not supported." + ) + + logger.info("Instantiating torch loaders") + + match self.loader_version: + case TorchTrainDataset.FLAG: + collate_fn = collate_raw_train_data + case _: + raise NotImplementedError( + f"Collate function not defined for loader version " + f"{self.loader_version}" + ) + + # Create batch samplers - branch on distributed vs non-distributed + # Group by resolution so batches stay homogeneous across configured sources. + def group_key(ds): + return ds.prognostic_src.grid_size + + if self.distributed is not None: + # Distributed training + assert self.distributed.world_size is not None + assert self.distributed.rank is not None + train_batch_sampler = DistributedEquivalenceGroupBatchSampler( + datasets=train_datasets, + group_key=group_key, + batch_size=self.batch_size, + num_replicas=self.distributed.world_size, + rank=self.distributed.rank, + shuffle=True, + drop_last=True, + seed=self.rand_seed, + ) + + val_batch_sampler = DistributedEquivalenceGroupBatchSampler( + datasets=val_datasets, + group_key=group_key, + batch_size=self.batch_size, + num_replicas=self.distributed.world_size, + rank=self.distributed.rank, + shuffle=False, + drop_last=False, + seed=self.rand_seed, + ) + else: + # Non-distributed training + train_batch_sampler = EquivalenceGroupBatchSampler.from_datasets( # type: ignore + datasets=train_datasets, + group_key=group_key, + batch_size=self.batch_size, + shuffle=True, + drop_last=True, + seed=self.rand_seed, + ) + + val_batch_sampler = EquivalenceGroupBatchSampler.from_datasets( # type: ignore + datasets=val_datasets, + group_key=group_key, + batch_size=self.batch_size, + shuffle=True, + drop_last=False, + seed=self.rand_seed, + ) + + # Store samplers for set_epoch calls + self.train_sampler = train_batch_sampler + self.val_sampler = val_batch_sampler + + # Create data loaders (same for both distributed and non-distributed) + # When using batch_sampler, don't specify batch_size or sampler + train_dataloader = DataLoader( + train_data, + batch_sampler=train_batch_sampler, + num_workers=self.num_workers, + persistent_workers=self.persistent_workers and self.num_workers > 0, + pin_memory=self.pin_mem, + collate_fn=collate_fn, + multiprocessing_context=self.mp_context, + ) + + val_dataloader = DataLoader( + val_data, + batch_sampler=val_batch_sampler, + num_workers=self.num_workers, + persistent_workers=self.persistent_workers and self.num_workers > 0, + pin_memory=self.pin_mem, + collate_fn=collate_fn, + multiprocessing_context=self.mp_context, + ) + + # Wrap dataloaders to handle GPU post-processing + self.train_loader = TrainDataLoader( + train_dataloader, train_datasets, self.device + ) + self.val_loader = TrainDataLoader(val_dataloader, val_datasets, self.device) + + def save_all_checkpoints(self, epoch: int, v_loss: float, inf_loss: float): + with self._test_context(): + is_best_val_loss = False + if v_loss <= self.best_val_loss: + logger.info( + f"Epoch validation loss ({v_loss:.3f}) is lower than " + f"previous best validation loss ({self.best_val_loss:.3f})." + ) + logger.info( + "Saving lowest validation loss checkpoint to " + f"{self.ckpt_paths.best_validation_checkpoint_path}" + ) + self.best_val_loss = v_loss + is_best_val_loss = True # wait until inference error is updated + if inf_loss is not None and (inf_loss <= self.best_inf_loss): + logger.info( + f"Epoch inference error ({inf_loss:.3f}) is lower than " + f"previous best inference error ({self.best_inf_loss:.3f})." + ) + logger.info( + "Saving lowest inference error checkpoint to " + f"{self.ckpt_paths.best_inference_checkpoint_path}" + ) + self.best_inf_loss = inf_loss + self.save_checkpoint( + epoch, self.ckpt_paths.best_inference_checkpoint_path + ) + if is_best_val_loss: + self.save_checkpoint( + epoch, self.ckpt_paths.best_validation_checkpoint_path + ) + + logger.info( + f"Saving latest checkpoint to {self.ckpt_paths.latest_checkpoint_path}" + ) + self.save_checkpoint(epoch, self.ckpt_paths.latest_checkpoint_path) + if epoch > 0 and epoch % self.save_freq == 0: + path = self.ckpt_paths.latest_checkpoint_path_with_epoch(epoch) + logger.info(f"Saving per-epoch checkpoint to {path}") + self.save_checkpoint(epoch, path) + + logger.info( + f"Saving latest EMA checkpoint to {self.ckpt_paths.ema_checkpoint_path}" + ) + self.save_checkpoint( + epoch, + self.ckpt_paths.ema_checkpoint_path, + for_inference=True, + ) + + def save_checkpoint( + self, + epoch: int, + checkpoint_path: Path, + for_inference: bool = False, + ): + if for_inference: + with self._ema_context(): + model_state_dict = self.model.state_dict() + else: + model_state_dict = self.model.state_dict() + + # Create temporary file in the same directory as the target + temp_dir = os.path.dirname(checkpoint_path) + with tempfile.NamedTemporaryFile(dir=temp_dir, delete=False) as tmp: + temporary_location = tmp.name + checkpoint = { + "model": model_state_dict, + "optimizer": self.optimizer.state_dict(), + "epoch": epoch, + "best_val_loss": self.best_val_loss, + "best_inf_loss": self.best_inf_loss, + "ema": self._ema.get_state(include_ema_params=not for_inference), + "num_batches_seen": self.num_batches_seen, + "train_progress": self.train_progress.state_dict(), + "wandb_id": self.wandb_id, + "wandb_name": self.wandb_name, + } + loss_state: dict[str, Any] | None = None + if state_dict_fn := getattr(self.loss_fn, "state_dict", None): + loss_state = state_dict_fn() + + if loss_state is not None: + checkpoint["loss_fn_state"] = loss_state + if self.scheduler: + checkpoint["scheduler"] = self.scheduler.state_dict() + + torch.save(checkpoint, temporary_location) + os.replace(temporary_location, checkpoint_path) + + def load_checkpoint(self, checkpoint_path, finetune=False): + logger.info(f"Loading checkpoint from {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location=torch.device(self.device)) + + # Remove module prefix from state dict + def remove_module_prefix(state_dict, prefix="module."): + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + name = k.removeprefix(prefix) + new_state_dict[name] = v + return new_state_dict + + # Load model state + model_state_dict = checkpoint["model"] + new_state_dict = remove_module_prefix(model_state_dict) + self.model.load_state_dict(new_state_dict) + + # Load EMA state + model_ema_state_dict = checkpoint["ema"] + new_ema_state_dict = remove_module_prefix(model_ema_state_dict) + if "ema_params" in new_ema_state_dict: + new_ema_state_dict["ema_params"] = remove_module_prefix( + new_ema_state_dict["ema_params"], prefix="module" + ) + self._ema = EMATracker.from_state(new_ema_state_dict, self.model) + + if not finetune: + self.optimizer.load_state_dict(checkpoint["optimizer"]) + if self.scheduler and "scheduler" in checkpoint: + self.scheduler.load_state_dict(checkpoint["scheduler"]) + + if load_state_dict_fn := getattr(self.loss_fn, "load_state_dict", None): + assert "loss_fn_state" in checkpoint, ( + f"Expected to load loss state for {self.loss_fn} but " + "no state found in checkpoint" + ) + load_state_dict_fn(checkpoint["loss_fn_state"]) + + self.start_epoch = checkpoint["epoch"] + 1 + self.wandb_id = checkpoint.get("wandb_id") + self.wandb_name = checkpoint.get("wandb_name") + self.num_batches_seen = checkpoint.get("num_batches_seen", 0) + self.train_progress = TrainProgress.from_state_dict( + checkpoint.get("train_progress") + ) + + logger.info(f"Start Epoch: {self.start_epoch}") + logger.info(f"Wandb id: {self.wandb_id}") + logger.info(f"Wandb name: {self.wandb_name}") + logger.info(f"Optimizer LR: {self.optimizer.param_groups[-1]['lr']}") + + self.best_val_loss = checkpoint["best_val_loss"] + self.best_inf_loss = checkpoint["best_inf_loss"] + + def is_wandb_enabled(self): + return self.wandb_logger.enabled and is_main_process() + + def _sync_flag_from_main(self, flag: bool) -> bool: + """Broadcast a boolean decision from rank 0 to every process.""" + if self.distributed is None: + return flag + + synced = torch.tensor([int(flag)], device=self.device) + torch.distributed.broadcast(synced, src=0) + return bool(synced.item()) + + @contextlib.contextmanager + def _test_context(self): + """ + The context for running validation/inference. + In this context, the stepper uses the EMA model if + `self.test_using_ema` is True. + """ + if self.test_using_ema: + with self._ema_context(): + yield + else: + yield + + @contextlib.contextmanager + def _ema_context(self): + """ + A context where the stepper uses the EMA model. + """ + self._ema.store(parameters=self.model.parameters()) + self._ema.copy_to(model=self.model) + try: + yield + finally: + self._ema.restore(parameters=self.model.parameters()) + + def finish(self): + self.wandb_logger.finish() + + +def main(): + # Load config from YAML + cfg = TrainConfig.from_yaml_and_cli() + cfg.prepare_output_dirs() # we do this first so logging can use them + + handle_logging(cfg.debug, cfg.experiment.output_dir) + handle_warnings() + + trainer = Trainer(cfg) + + try: + trainer.run() + except Exception as e: + logger.exception("Training failed with an exception") + raise e + + +if __name__ == "__main__": + main() diff --git a/samudra/utils/__init__.py b/samudra/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2fba2cd3fa379a616e909780792c8c6bc97f6783 --- /dev/null +++ b/samudra/utils/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/samudra/utils/compare.py b/samudra/utils/compare.py new file mode 100644 index 0000000000000000000000000000000000000000..b113fa91f0c0f80374e077b63297c6555fad47c2 --- /dev/null +++ b/samudra/utils/compare.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import filecmp +import mimetypes +from pathlib import Path + +import numpy as np +from PIL import Image + + +def compare_images(file1, file2, diff_dir=None): + """ + Compare two PNG images based on their RGB content. + + Args: + file1: Path to first image + file2: Path to second image + diff_dir: Optional directory to save difference images + + Returns: + bool: True if images have identical RGB content, False otherwise + """ + try: + # Open images and convert to RGB (ignoring alpha channel if present) + img1 = Image.open(file1).convert("RGB") + img2 = Image.open(file2).convert("RGB") + + # Check if sizes match + if img1.size != img2.size: + return False + + # Convert to numpy arrays and compare + arr1 = np.array(img1) + arr2 = np.array(img2) + + if np.array_equal(arr1, arr2): + return True + + # If different and diff_dir is provided, create difference image + if diff_dir: + # Calculate absolute difference + diff = np.abs(arr1.astype(np.int16) - arr2.astype(np.int16)) + + # Enhance visibility by scaling up small differences + diff_enhanced = np.clip(diff * 10, 0, 255).astype(np.uint8) + + # Create a composite image showing original1, original2, and difference + height, width = arr1.shape[:2] + composite = np.zeros((height, width * 3 + 20, 3), dtype=np.uint8) + composite[:, :width] = arr1 + composite[:, width + 10 : width * 2 + 10] = arr2 + composite[:, width * 2 + 20 :] = diff_enhanced + + # Save the composite image + diff_path = diff_dir / f"diff_{Path(file1).name}" + Image.fromarray(composite).save(diff_path) + + return False + except Exception as e: + print(f"Error comparing images {file1} and {file2}: {e}") + return False + + +def compare_directories(dir1: Path, dir2: Path, diff_dir: Path | None = None): + """ + Compare two directories and return files that don't match. + + Args: + dir1: Path to first directory + dir2: Path to second directory + diff_dir: Optional directory to save difference images + + Returns: + tuple: (list of non-matching files, number of matching files) + """ + dir1 = Path(dir1) + dir2 = Path(dir2) + + if not dir1.exists(): + raise ValueError(f"Directory {dir1} does not exist") + if not dir2.exists(): + raise ValueError(f"Directory {dir2} does not exist") + + # Get all files in both directories + files1 = {f.relative_to(dir1) for f in dir1.rglob("*") if f.is_file()} + files2 = {f.relative_to(dir2) for f in dir2.rglob("*") if f.is_file()} + + # Files only in one directory + only_in_dir1 = files1 - files2 + only_in_dir2 = files2 - files1 + + # Common files + common_files = files1 & files2 + + # Separate different types of non-matching files + content_different = [] + missing_in_dir1 = [] + missing_in_dir2 = [] + matching_count = 0 + + # Compare common files + print(f"Comparing {len(common_files)} common files...") + for f in common_files: + file1 = dir1 / f + file2 = dir2 / f + + # Special handling for image files + file_mime, _encoding = mimetypes.guess_type(f) + if file_mime and file_mime.startswith("image/"): + if compare_images(file1, file2, diff_dir): + matching_count += 1 + else: + content_different.append(f) + else: + # Regular file comparison for non-image files + if filecmp.cmp(file1, file2, shallow=False): + matching_count += 1 + else: + content_different.append(f) + + # Files only in one directory + for f in only_in_dir1: + missing_in_dir2.append(f) + for f in only_in_dir2: + missing_in_dir1.append(f) + + # Print summary + print(f"\n=== Directory Comparison Summary ===") + print(f"Total files in {dir1.name}: {len(files1)}") + print(f"Total files in {dir2.name}: {len(files2)}") + print(f"Common files: {len(common_files)}") + print(f"Matching files: {matching_count}") + + # Print present but different files first + if content_different: + print(f"\nFiles with different content ({len(content_different)}):") + _print_file_list(content_different) + + # Then print missing files + if missing_in_dir2: + print(f"\nFiles only in {dir1.name} ({len(missing_in_dir2)}):") + _print_file_list(missing_in_dir2) + + if missing_in_dir1: + print(f"\nFiles only in {dir2.name} ({len(missing_in_dir1)}):") + _print_file_list(missing_in_dir1) + + # Combine all non-matching for return value + non_matching = content_different + missing_in_dir1 + missing_in_dir2 + + return non_matching, matching_count + + +def _print_file_list(file_list: list[Path]): + for f in file_list[:10]: + print(f" - {f}") + if len(file_list) > 10: + print(f" ... and {len(file_list) - 10} more") + + +if __name__ == "__main__": + import argparse + import sys + import tempfile + + parser = argparse.ArgumentParser() + parser.add_argument("dir1", type=Path) + parser.add_argument("dir2", type=Path) + args = parser.parse_args() + + diff_dir = Path(tempfile.mkdtemp("image-diffs")) + non_matching, _ = compare_directories(args.dir1, args.dir2, diff_dir) + print(f"See image differences in {diff_dir}") + + if len(non_matching) > 0: + sys.exit(1) diff --git a/samudra/utils/ctx.py b/samudra/utils/ctx.py new file mode 100644 index 0000000000000000000000000000000000000000..6d9c2d7c8c64de8af4720ac4505cdfa0f60fbf51 --- /dev/null +++ b/samudra/utils/ctx.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +from typing import Self + +import torch + +from samudra.constants import Lat, Lon, PrognosticMask + + +@dataclasses.dataclass(frozen=True) +class GridContext: + """Grid-level context for model forward passes and loss computation. + + Bundles spatial metadata that travels alongside input tensors during training: + the ocean/land mask for excluding land cells from gradients, and the grid + resolution for latitude-weighted loss calculations. + + Attributes: + label_mask: Boolean mask indicating valid ocean cells for each prognostic + variable. Shape: (num_prognostic_vars, lat, lon). Land cells are False. + input_resolution_cpu: Tuple of (latitude, longitude) coordinate tensors defining + the input spatial grid. Used by encoders for positional encoding. + output_resolution_cpu: Tuple of (latitude, longitude) coordinate tensors defining + the output spatial grid. Used by decoders to determine output shape and + pixel-query positions. Defaults to input_resolution_cpu for single-scale use. + """ + + label_mask: PrognosticMask + input_resolution_cpu: tuple[Lat, Lon] + output_resolution_cpu: tuple[Lat, Lon] + + def to(self, device: torch.device) -> Self: + """Move the label mask to the specified device.""" + return dataclasses.replace( + self, label_mask=self.label_mask.to(device, non_blocking=True) + ) diff --git a/samudra/utils/data.py b/samudra/utils/data.py new file mode 100644 index 0000000000000000000000000000000000000000..b8c26b4975ad581893ff65a4665f63f43e9dd60a --- /dev/null +++ b/samudra/utils/data.py @@ -0,0 +1,1029 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import logging +import re +from collections import defaultdict +from collections.abc import Callable +from functools import cached_property +from typing import TYPE_CHECKING, Literal, Self + +import numpy as np +import torch +import xarray as xr +from einops import rearrange +from jaxtyping import Bool, Float + +if TYPE_CHECKING: + from samudra.config import TimeConfig + +from samudra.constants import ( + BatchTimeSeriesOutput, + BoundaryVarNames, + DatasetSpec, + DictSingleChannelVar, + Grid, + GridMask, + GridSize, + Input, + Lat, + LoaderVersion, + Lon, + Prognostic, + PrognosticMask, + PrognosticVarNames, + SingleTimeSeriesOutput, + TensorMap, + construct_metadata, +) +from samudra.derived_variables import add_derived_variables +from samudra.utils.llc import _preferred_available_var, _var_without_level + +logger = logging.getLogger(__name__) + + +def _var_name_encode_level(var_name: str) -> bool: + """Check if the variable name encodes the level.""" + var_name_encodes_level = re.compile(r"_[0-9]+") + return bool(var_name_encodes_level.search(var_name)) + + +def _is_compact(data: xr.Dataset, means: xr.Dataset, stds: xr.Dataset) -> bool: + return all( + not _var_name_encode_level(str(v)) + for d in [data, means, stds] + for v in d.keys() + ) + + +@dataclasses.dataclass +class Masks: + """A collection of masks to expose the ocean and mask land.""" + + prognostic: PrognosticMask + boundary: GridMask + + def __post_init__(self): + self.prognostic = self.prognostic.bool() + self.boundary = self.boundary.bool() + + def prognostic_with_hist( + self, hist: int + ) -> Bool[GridMask, " prognostic_vars*({hist}+1)"]: + return torch.concat([self.prognostic] * (hist + 1), dim=0) + + +@dataclasses.dataclass +class DataSource: + """Data source for the model.""" + + name: str + data: xr.Dataset + means: xr.Dataset + stds: xr.Dataset + masks: Masks + dataset_spec: DatasetSpec + + @cached_property + def is_compact(self) -> bool: + """Check if the data source is compact.""" + return _is_compact(self.data, self.means, self.stds) + + @cached_property + def resolution(self) -> tuple[Lat, Lon]: + return ( + torch.from_numpy(self.data.lat.values), + torch.from_numpy(self.data.lon.values), + ) + + @cached_property + def grid_size(self) -> GridSize: + res = self.resolution + return res[0].shape[0], res[1].shape[0] + + @cached_property + def spherical_area_weights(self) -> Grid: + return spherical_area_weights(self.data) + + @cached_property + def metadata(self) -> dict: + return construct_metadata(self.data, self.dataset_spec) + + def filter( + self, + var_names: PrognosticVarNames | BoundaryVarNames, + *, + prefix: str, + ) -> Self: + """Filter the data source to only include the specified variables (and levels). + + If the dataset is compact, it will also filter the levels based on the + variable names (which encode the level in the name). + + Args: + var_names: Variable names to filter. + prefix: Prefix for the new data source name. + + Returns: + A new `DataSource` only with the filtered variables and levels. + """ + name = f"{prefix}[{self.name}]" + if self.is_compact: + parsed_var_names, levels = [], [] + for mangled_var_name in var_names: + if not _var_name_encode_level(mangled_var_name): + parsed_var_names.append(mangled_var_name) + continue + tokens = mangled_var_name.split("_") + var_name, level = tokens[0], int(tokens[1]) + + parsed_var_names.append(var_name) + # Build set of total levels + if level not in levels: + levels.append(level) + + data = self.data[parsed_var_names] + means = self.means[parsed_var_names] + stds = self.stds[parsed_var_names] + if levels: + data = data.isel(lev=levels) + means = means.isel(lev=levels) + stds = stds.isel(lev=levels) + + return dataclasses.replace( + self, name=name, data=data, means=means, stds=stds + ) + + data = self.data[var_names] + means = self.means[var_names] + stds = self.stds[var_names] + + return dataclasses.replace(self, name=name, data=data, means=means, stds=stds) + + def map( + self, + func: Callable[ + [xr.Dataset, xr.Dataset, xr.Dataset], + tuple[xr.Dataset, xr.Dataset, xr.Dataset], + ], + *, + suffix: str | None = None, + ) -> Self: + """Map the function over the data source.""" + if suffix is None: + suffix = func.__qualname__ + + data, means, stds = func(self.data.copy(), self.means.copy(), self.stds.copy()) + + return dataclasses.replace( + self, name=f"{self.name}_{suffix}", data=data, means=means, stds=stds + ) + + def map_data( + self, func: Callable[[xr.Dataset], xr.Dataset], *, suffix: str | None = None + ) -> Self: + """Map the function over just data in DataSource.""" + if suffix is None: + suffix = func.__qualname__ + data = func(self.data.copy()) + return dataclasses.replace(self, name=f"{self.name}_{suffix}", data=data) + + def slice(self, time: "TimeConfig") -> Self: + """Slice the data source to only include the specified time slice.""" + data_time_min = self.data.time.values.min() + data_time_max = self.data.time.values.max() + time_start = time.time_slice.start + time_end = time.time_slice.stop + if time_start > data_time_max or time_end < data_time_min: + raise ValueError( + f"Time slice {time} is entirely outside the range of the data " + f"{str(data_time_min)[:10]} to {str(data_time_max)[:10]}" + ) + + if time_start < data_time_min or time_end > data_time_max: + logger.warning( + f"Time slice {time} is partially outside the range of the data " + f"{str(data_time_min)[:10]} to {str(data_time_max)[:10]}" + ) + + data = self.data.sel(time=time.time_slice) + return dataclasses.replace(self, name=f"{time=}[{self.name}]", data=data) + + # TODO(jder): delete this once we've de-duplicated InferenceDataset with TorchTrainDataset + def normalize(self, fill_nan=True, fill_value=0.0) -> xr.Dataset: + """Normalize input data.""" + norm = (self.data - self.means) / self.stds + if fill_nan: + norm = norm.fillna(fill_value) + return norm + + # TODO(jder): delete this once we've de-duplicated InferenceDataset with TorchTrainDataset + def normalize_with( + self, + data: torch.Tensor, + variable_axis: int = 0, + fill_nan=True, + fill_value=0.0, + ) -> torch.Tensor: + """Normalize input data treated as torch Tensors.""" + reshape_vars = [1] * data.ndim + reshape_vars[variable_axis] = -1 + + # TODO(alxmrs): Do we have to reshape twice? + if "lev" in self.means.dims: + means_np = ( + conditional_rearrange( + self.means, + "(variable lev)=var", + concat_dim="var", + ) + .rename({"var": "variable"}) + .to_numpy() + .reshape(-1) + ) + else: + means_np = self.means.to_array().to_numpy().reshape(-1) + if "lev" in self.stds.dims: + stds_np = ( + conditional_rearrange( + self.stds, + "(variable lev)=var", + concat_dim="var", + ) + .rename({"var": "variable"}) + .to_numpy() + .reshape(-1) + ) + else: + stds_np = self.stds.to_array().to_numpy().reshape(-1) + + means = torch.from_numpy(means_np).reshape(reshape_vars) + stds = torch.from_numpy(stds_np).reshape(reshape_vars) + + norm = (data - means) / stds + if fill_nan: + norm = norm.nan_to_num(nan=fill_value) + norm = norm.to(data.dtype) + return norm + + @classmethod + def from_datasets( + cls, + data: xr.Dataset, + means: xr.Dataset, + stds: xr.Dataset, + *, + dataset_spec: DatasetSpec, + prognostic_var_names: PrognosticVarNames, + boundary_var_names: BoundaryVarNames, + name: str = "DataSource", + ) -> Self: + data, means, stds = validate_data( + data, + means, + stds, + dataset_spec=dataset_spec, + boundary_var_names=boundary_var_names, + ) + masks = extract_wet_mask( + data, + prognostic_var_names, + boundary_var_names, + dataset_spec=dataset_spec, + ) + + return cls( + name=name, + data=data, + means=means, + stds=stds, + masks=masks, + dataset_spec=dataset_spec, + ) + + +@dataclasses.dataclass +class OceanData: + """A slice of ocean data (boundary or prognostic) with normalization statistics. + + This dataclass bundles raw tensor data with the statistics needed to normalize it + and the mask needed to handle land/invalid values. It serves as an intermediary + representation used when constructing training `Example`s from raw xarray data. + + The typical workflow is: + 1. Load raw data from a `DataSource` via `from_data_source()` + 2. Slice to the desired time range with `with_time()` + 3. Apply normalization and masking with `normalize_and_mask()` + 4. Flatten time/variable dims to create the final `Input` or `Prognostic` tensor + + Attributes: + data: Raw ocean variable values with shape (batch, time, variable, lat, lon). + means: Per-variable means for normalization, shape (variable,). + stds: Per-variable standard deviations for normalization, shape (variable,). + mask: Boolean mask indicating valid ocean points (True) vs land (False), + broadcast-compatible with the variable dimension. + """ + + data: Float[torch.Tensor, "batch time variable lat lon"] + means: Float[torch.Tensor, " variable"] + stds: Float[torch.Tensor, " variable"] + mask: Bool[torch.Tensor, " variable"] + + @classmethod + def from_data_source( + cls, + data: Float[torch.Tensor, "batch time variable lat lon"], + mask: Float[torch.Tensor, " variable"], + src: DataSource, + ) -> Self: + means_torch = torch.from_numpy(_flatten(src.means)) + stds_torch = torch.from_numpy(_flatten(src.stds)) + return cls(data, means_torch, stds_torch, mask) + + def with_time(self, time_range: slice) -> Self: + """Slice data across the time dimension.""" + return dataclasses.replace(self, data=self.data[:, time_range, :, :, :]) + + def _normalize( + self, + data: Float[torch.Tensor, "batch time var lat lon"], + fill_nan: bool = True, + fill_value: float = 0.0, + ) -> Float[torch.Tensor, "batch time var lat lon"]: + """Normalize input data treated as torch Tensors.""" + norm = (data - self.means.view(1, 1, -1, 1, 1)) / self.stds.view(1, 1, -1, 1, 1) + if fill_nan: + norm = norm.nan_to_num(nan=fill_value) + norm = norm.to(data.dtype) + return norm + + def normalize_and_mask( + self, normalize_before_mask: bool, masked_fill_value: float + ) -> Float[torch.Tensor, "batch time var lat lon"]: + """Normalize and mask tensors.""" + tensor = self.data + if normalize_before_mask: + tensor = self._normalize(tensor) + tensor = torch.where(self.mask, tensor, masked_fill_value) + if not normalize_before_mask: + tensor = self._normalize(tensor) + return tensor + + def to(self, device: torch.device, non_blocking: bool = True) -> Self: + return dataclasses.replace( + self, + data=self.data.to(device, non_blocking=non_blocking), + means=self.means.to(device, non_blocking=non_blocking), + stds=self.stds.to(device, non_blocking=non_blocking), + mask=self.mask.to(device, non_blocking=non_blocking), + ) + + +@dataclasses.dataclass +class DataSourceSplits: + train: DataSource + val: DataSource + inference: DataSource | None + + +@dataclasses.dataclass +class DataContainer: + train_sources: list[DataSource] + val_sources: list[DataSource] + inference_source: DataSource | None + loader_version: LoaderVersion + dataset_spec: DatasetSpec + + # TODO: this is a bit of a footgun now that we have multiple kinds of sources + # and should be removed in favor of the appropriate source above. + @property + def primary_source(self) -> DataSource: + return self.train_sources[0] + + +def conditional_rearrange( + data: xr.Dataset, pattern: str, except_dim="lev", concat_dim="variable" +) -> xr.DataArray: + """Rearrange a Dataset using an einsum notation with and without a dimension. + + When a dataset has variables with a mixture of dimensions and an einsum-like + rearrange is applied on that dataset, it's common that the pattern will combinate + one too many variables. Sometimes, it's desirable to apply the rearrange pattern + on two versions of the data: one including variables with that dimension and one + without, and then concatenate them along a new dimension. + + For example, surface level boundary variables, which only occur at t0, should not be + combinatorially rearranged with depth variables that have multiple time steps. In + such a situation, this function can be used to apply a standard einsum rearrangement + to depth and surface variables, including and excluding variables who have a `time` + dimension, respectively. + + This method is stable: even if it creates a new number of dimensions, it will + preserve the order of the variables in the original dataset. + + Args: + data: The dataset to rearrange. + pattern: The einsum pattern to use for rearranging. + except_dim: The dimension to exclude from the pattern. + concat_dim: The dimension to concatenate along. + + Returns: + The combined, rearranged dataset as a `xarray.DataArray`. + """ + assert except_dim in pattern, f"{except_dim} must be in the pattern." + + all_vars = list(data.keys()) + + vars_with_dim = [v for v in data if except_dim in data[v].dims] + vars_without_dim = [v for v in data if except_dim not in data[v].dims] + + # Some of the `vars_without_dim` may need to appear before or behind `vars_with_dim` + # in the final data array. These lists help preserve the correct order of the vars, + # even after a rearrangement (i.e. merge to two or more dimensions). + back = [ + v + for v in vars_without_dim + if all_vars.index(v) > all_vars.index(vars_with_dim[0]) + ] + front = [ + v + for v in vars_without_dim + if all_vars.index(v) < all_vars.index(vars_with_dim[-1]) + ] + + data_with_dim = ( + data[vars_with_dim] + .to_array() + .einops.rearrange(pattern, dask="allowed") + .drop_vars(concat_dim, errors="ignore") + ) + data_without_dim = ( + data[vars_without_dim] + .to_array() + .einops.rearrange(pattern.replace(except_dim, ""), dask="allowed") + .drop_vars(concat_dim, errors="ignore") + ) + + da = xr.concat([data_with_dim, data_without_dim], dim=concat_dim) + + n_front = len(front) # e.g. n_front=2 + n_center = data_with_dim.sizes[concat_dim] # e.g. n_center=10 + n_back = len(back) # e.g. n_back=3 + + # In the `concat` above, we put all the `data_without_dim` vars at the end. Some of + # these need to be moved to the front, and the rest stays at the back. Here, we + # compute a list of indices that will sort the data in the correct order. + # + # e.g. with the example constants above, order would look like: + # array([10, 11, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 13, 14]) + order = np.concatenate( + ( + # Moves vars to the front + np.roll(np.arange((new_front := n_center + n_front)), n_front), + np.arange(new_front, new_front + n_back), # rest of vars + ) + ) + order_da = xr.DataArray(order, dims=concat_dim) + + return da.sortby(order_da) + + +def _flatten(ds: xr.Dataset) -> np.ndarray: + """Flatten a means/stds dataset into a 1-D array with one entry per channel. + + Variables with a `lev` dim are flattened with `lev` first; variables without + a `lev` dim are flattened as-is. This matches the channel layout of the + prognostic / boundary tensors, which `to_array().reshape(-1)` would otherwise + mis-align by broadcasting surface-only variables across all depth levels. + """ + flattened = [] + for name in ds.data_vars: + var = ds[name] + if "lev" in var.dims: + var = var.transpose("lev", ...) + flattened.append(var.to_numpy().reshape(-1)) + return np.concatenate(flattened) + + +def _level_index_from_var_name(var_name: str) -> int: + suffix = var_name.rsplit("_", maxsplit=1)[-1] + return int(suffix) if suffix.isdigit() else 0 + + +def _mask_var_for_data_var( + data: xr.Dataset, + var_name: str, + *, + dataset_spec: DatasetSpec, +) -> str: + level = _level_index_from_var_name(var_name) + base_var = _var_without_level(var_name) + llc_staggered_masks = { + "U": (f"mask_w_{level}", f"hFacW_{level}"), + "oceTAUX": ("mask_w_0", "hFacW_0"), + "V": (f"mask_s_{level}", f"hFacS_{level}"), + "oceTAUY": ("mask_s_0", "hFacS_0"), + } + if dataset_spec.type == "llc" and base_var in llc_staggered_masks: + if mask_var := _preferred_available_var(data, llc_staggered_masks[base_var]): + return mask_var + return dataset_spec.mask_vars[level] + + +def _mask_array_for_data_var( + data: xr.Dataset, + var_name: str, + *, + dataset_spec: DatasetSpec, +) -> np.ndarray: + mask = data[_mask_var_for_data_var(data, var_name, dataset_spec=dataset_spec)] + if "time" in mask.dims: + mask = mask.isel(time=0) + return mask.to_numpy() + + +def extract_wet_mask( + data: xr.Dataset, + prognostic_var_names: PrognosticVarNames, + boundary_var_names: BoundaryVarNames, + *, + dataset_spec: DatasetSpec, +) -> Masks: + """A mask for where the oceans are. Water is wet.""" + data_ = flatten_masks(data, dataset_spec=dataset_spec) + wet_inp_np = np.stack( + [ + _mask_array_for_data_var(data_, var_name, dataset_spec=dataset_spec) + for var_name in prognostic_var_names + ] + ) + boundary_mask_vars = [ + _mask_var_for_data_var(data_, var_name, dataset_spec=dataset_spec) + for var_name in boundary_var_names + ] + boundary_masks = [ + _mask_array_for_data_var(data_, var_name, dataset_spec=dataset_spec) + for var_name in boundary_var_names + ] + wet_surface_mask_np = ( + np.stack(boundary_masks) + if len(set(boundary_mask_vars)) > 1 + else boundary_masks[0] + ) + + wet_inp = torch.from_numpy(wet_inp_np) + wet_surface = torch.from_numpy(wet_surface_mask_np) + return Masks(wet_inp.bool(), wet_surface.bool()) + + +def flatten_masks( + data: xr.Dataset, + dataset_spec: DatasetSpec, +) -> xr.Dataset: + """Adds level-wise mask variables from the stacked wet mask.""" + data_ = data.copy() + mask_vars = list(dataset_spec.mask_vars) + if mask_vars[0] not in data_.variables: + assert dataset_spec.mask_all_levels_var in data_.variables, ( + "Wet mask cannot be constructed without " + "either the wetmask variable or the level-wise masks" + ) + + wet_mask = data_[dataset_spec.mask_all_levels_var] + for i, mask_var in enumerate(mask_vars): + data_[mask_var] = wet_mask.isel(lev=i) + + data_ = data_.drop_vars(dataset_spec.mask_all_levels_var) + + return data_ + + +def unflatten_masks( + data: xr.Dataset, + dataset_spec: DatasetSpec, +) -> xr.Dataset: + """Adds a stacked wet mask `xarray.DataArray` from level-wise mask variables.""" + data_ = data.copy() + mask_vars = list(dataset_spec.mask_vars) + if dataset_spec.mask_all_levels_var not in data_.variables: + assert mask_vars[0] in data_.variables, "Wet mask must have masks as data vars!" + + wetmask = data_[mask_vars].to_array( + dim="lev", name=dataset_spec.mask_all_levels_var + ) + + data_[dataset_spec.mask_all_levels_var] = wetmask.assign_coords(lev=data_.lev) + data_ = data_.drop_vars(mask_vars) + + return data_ + + +def spherical_area_weights(data: xr.Dataset) -> Grid: + num_lon = data.lon.size + lats = torch.from_numpy(data.lat.to_numpy()) + weights = torch.cos(torch.deg2rad(lats)).repeat(num_lon, 1).t() + weights /= weights.sum() + return weights + + +def spherical_area(data: xr.Dataset) -> Grid: + """ + Compute real grid cell areas on a spherical Earth. + + Uses the spherical geometry formula: + A = R² × Δλ × (sin(φ₂) - sin(φ₁)) + + where: + - R is Earth's radius (6371 km) + - Δλ is the longitude spacing in radians + - φ₁, φ₂ are the latitude bounds of the cell in radians + + Args: + data: Dataset containing lat/lon coordinates + + Returns: + Grid cell areas in m² + """ + R = 6371000 # Earth radius in meters + + lats = data.lat.to_numpy() + lons = data.lon.to_numpy() + + # Compute grid spacing (assuming uniform spacing) + dlat = np.abs(np.diff(lats).mean()) + dlon = np.abs(np.diff(lons).mean()) + + # Convert to radians + dlat_rad = np.deg2rad(dlat) + dlon_rad = np.deg2rad(dlon) + lats_rad = np.deg2rad(lats) + + # Compute cell areas: A = R² × dlon × (sin(lat + dlat/2) - sin(lat - dlat/2)) + areas = np.zeros((len(lats), len(lons))) + for i, lat_rad in enumerate(lats_rad): + lat_north = lat_rad + dlat_rad / 2 + lat_south = lat_rad - dlat_rad / 2 + area = R**2 * dlon_rad * (np.sin(lat_north) - np.sin(lat_south)) + areas[i, :] = area + + return torch.from_numpy(areas) + + +def get_inference_steps(data_source: DataSource, hist: int = 1): + """ + Get the number of inference/rollout steps for the given time configuration. + + Args: + data_source: The data source sliced to the inference time range + hist: How many additional history samples we get per step + + Returns: + num_steps: Total number of rolled-out inferences which fit into the time range + """ + num_steps = data_source.data.time.size + + # Might have extra remaining days, so we remove them + mod = num_steps % (hist + 1) + num_steps = num_steps - mod + return num_steps + + +def convert_tensor_out_to_dict( + tensor_out: torch.Tensor, + *, + tensor_map: TensorMap, +) -> DictSingleChannelVar: + assert tensor_out.ndim == 5 + assert tensor_out.shape[2] == len(tensor_map.prognostic_var_names) + out_dict = {} + for i, var in enumerate(tensor_map.prognostic_var_names): + out_dict[var] = tensor_out[:, :, i] + out_dict.update(add_derived_variables(tensor_out, tensor_map=tensor_map)) + return out_dict + + +def get_aggregator_dicts( + data: Prognostic | Input, + normalize: "Normalize", + tensor_map: TensorMap, + wet: torch.Tensor, + long_rollout: bool, + input_type: Literal["prognostic", "input"] = "prognostic", + num_prognostic_channels: int = 0, + hist: int = 1, +) -> tuple[DictSingleChannelVar, DictSingleChannelVar]: + # Remove boundary data if input + if input_type == "input": + data = data[:, :num_prognostic_channels] + + # Separate history from channels + data_reshaped: SingleTimeSeriesOutput | BatchTimeSeriesOutput + if long_rollout: + # All batches are part of the same rollout during inference + data_reshaped = rearrange( + data, "n (hi c) h w -> (n hi) c h w", hi=hist + 1 + ).unsqueeze(0) # add artificial batch dim + elif data.ndim != 5: + # Batches are independent rollouts during validation + data_reshaped = rearrange(data, "n (hi c) h w -> n hi c h w", hi=hist + 1) + else: + # This case comes up in tests; typically, data is not in the desired shape automatically. + data_reshaped = data + + # Get normalized dict + data_normalized = data_reshaped.clone() + data_normalized = torch.where(wet == 0, float("nan"), data_normalized) + data_dict = convert_tensor_out_to_dict(data_normalized, tensor_map=tensor_map) + # Unnormalize + data_unnorm = normalize.unnormalize_tensor_prognostic( + data_reshaped, fill_value=float("nan") + ) + # Get unnormalized dict + data_unnorm_dict = convert_tensor_out_to_dict(data_unnorm, tensor_map=tensor_map) + return data_dict, data_unnorm_dict + + +def get_anomalies_vars(var_names: BoundaryVarNames) -> tuple[str, ...]: + """Get the variables that need to be computed for anomalies.""" + return tuple([var for var in var_names if var.endswith("_anomalies")]) + + +def compute_anomalies( + data: xr.Dataset, + means: xr.Dataset, + stds: xr.Dataset, + anomalies_vars: tuple[str, ...], +) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + """Compute anomalies for the given variables.""" + for var in anomalies_vars: + base_var = var.replace("_anomalies", "") + if var not in data.variables and base_var in data.variables: + logger.info(f"Computing anomalies for {base_var}") + climatology = ( + data[base_var].groupby("time.dayofyear").mean("time").compute() + ) + # Remove the seasonal cycle (climatology) from the detrended data + day_of_year = data[base_var]["time"].dt.dayofyear + data[var] = ( + data[base_var] - climatology.sel(dayofyear=day_of_year) + ).compute() + data = data.drop_vars(["dayofyear"]) + means[var] = data[var].mean().compute() + stds[var] = data[var].std().compute() + return data, means, stds + + +def with_level_index_vars( + data: xr.Dataset, + dataset_spec: DatasetSpec, +) -> xr.Dataset: + """ + Ensure variable names use a depth level index, not depth level value. + """ + data_copy = data.copy() + + for var in data.variables: + # OM4 data format has variables in the form: var_lev_{depthlevel} + # ex. so_lev_1040_0. We need to convert into var_{depthlevelidx} + var_str = str(var) + if "_lev_" in var_str: + var_split = var_str.split("_lev_") + var = var_split[0] + lev_in_depth = float(var_split[1].replace("_", ".")) + lev_in_depth_idx = dataset_spec.depth_levels.index(lev_in_depth) + data_copy = data_copy.rename({var_str: f"{var}_{lev_in_depth_idx!s}"}) + + return data_copy + + +def with_depth_value_vars( + data: xr.Dataset, + dataset_spec: DatasetSpec, +) -> xr.Dataset: + """Inverse of `with_level_index_vars`: name 3D variables by depth value. + + Renames the depth-resolved prognostic variables (``_``) + back to the OM4 ``_lev_`` form (e.g. ``thetao_0`` -> + ``thetao_lev_2_5``). Which variables to rename is read directly off + ``dataset_spec.prognostic_var_names`` rather than inferred from the data, so + per-level masks (``mask_``) and level-free prognostics (e.g. ``zos``) are + never mistaken for depth-resolved variables by name alone. + """ + renames = {} + for var_name in dataset_spec.prognostic_var_names: + base, _, idx = var_name.rpartition("_") + if not (base and idx.isdigit()): + continue # level-free prognostic variable, e.g. zos + depth = dataset_spec.depth_levels[int(idx)] + depth_str = str(depth).replace(".", "_") + if var_name in data.variables: + renames[var_name] = f"{base}_lev_{depth_str}" + + return data.rename(renames) + + +def with_lat_lon_coords(data: xr.Dataset) -> xr.Dataset: + """Standardize dataset coordinates; prefer "lat"/"lon" over "y"/"x".""" + data_copy = data.copy() + if "lat" not in data_copy.dims: + # Preserve the 2-D geographic coords under a non-colliding name, then + # rename the x/y dims to the 1-D lat/lon we standardize on. + preserve = {n: f"{n}_2d" for n in ("lat", "lon") if n in data_copy.coords} + data_copy = data_copy.rename(preserve).rename({"x": "lon", "y": "lat"}) + + return data_copy + + +def validate_data( + data: xr.Dataset, + means: xr.Dataset, + stds: xr.Dataset, + dataset_spec: DatasetSpec, + boundary_var_names: BoundaryVarNames, +) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + """Validate the data such that we have the correct format for training.""" + is_compact = _is_compact(data, means, stds) + + if is_compact: + data = with_lat_lon_coords(data) + else: + data = ( + data.pipe(flatten_masks, dataset_spec=dataset_spec) + .pipe(with_level_index_vars, dataset_spec=dataset_spec) + .pipe(with_lat_lon_coords) + ) + + # Check if data variables are in the right format + # This check is to ensure we convert data to the correct format + means = with_level_index_vars(means, dataset_spec=dataset_spec) + stds = with_level_index_vars(stds, dataset_spec=dataset_spec) + + # Check if any anomalies are needed to be computed + anomalies_vars = get_anomalies_vars(boundary_var_names) + out = ( + compute_anomalies(data, means, stds, anomalies_vars) + if anomalies_vars + else (data, means, stds) + ) + + return out + + +class Normalize: + def __init__( + self, + src: DataSource, + prognostic_var_names: PrognosticVarNames, + boundary_var_names: BoundaryVarNames, + ) -> None: + """Store normalization parameters and pre-compute numpy arrays.""" + prognostic_src = src.filter(prognostic_var_names, prefix="prognostic") + boundary_src = src.filter(boundary_var_names, prefix="boundary") + self.prognostic_mean = prognostic_src.means + self.prognostic_std = prognostic_src.stds + self.boundary_mean = boundary_src.means + self.boundary_std = boundary_src.stds + self.wet_mask = src.masks.prognostic + self.wet_mask_surface = src.masks.boundary + + # Pre-compute numpy arrays for faster access. When a dataset mixes + # variables with and without a `lev` dim (e.g. thermo_dynamic_all has + # 4 depth-resolved variables plus surface-only `zos`), a plain + # `to_array().reshape(-1)` would broadcast the surface variable across + # all levels, producing too many channels. `_flatten` flattens each + # variable independently, matching the prognostic tensor channel layout. + self._prognostic_mean_np = _flatten(self.prognostic_mean) + self._prognostic_std_np = _flatten(self.prognostic_std) + self._boundary_mean_np = _flatten(self.boundary_mean) + self._boundary_std_np = _flatten(self.boundary_std) + self._wet_mask_np = self.wet_mask.numpy() + + def normalize_tensor_prognostic( + self, data: torch.Tensor, fill_nan=True, fill_value=0.0 + ) -> torch.Tensor: + """Normalize prognostic tensor.""" + tensor_mean = torch.from_numpy(self._prognostic_mean_np).to( + data.device, data.dtype + ) + tensor_std = torch.from_numpy(self._prognostic_std_np).to( + data.device, data.dtype + ) + + expand_var_dim = [1] * data.ndim + expand_var_dim[-3] = -1 + assert data.shape[-3] == self._prognostic_mean_np.shape[0] + tensor_mean = tensor_mean.reshape(expand_var_dim) + tensor_std = tensor_std.reshape(expand_var_dim) + + norm = (data - tensor_mean) / tensor_std + if fill_nan: + norm = norm.nan_to_num(nan=fill_value) + norm = norm.to(data.dtype) + return norm + + def unnormalize_tensor_prognostic( + self, data: torch.Tensor, fill_value=float("nan") + ) -> torch.Tensor: + """Unnormalize prognostic tensor and apply fill value to land cells.""" + tensor_mean = torch.from_numpy(self._prognostic_mean_np).to( + data.device, data.dtype + ) + tensor_std = torch.from_numpy(self._prognostic_std_np).to( + data.device, data.dtype + ) + + expand_var_dim = [1] * data.ndim + expand_var_dim[-3] = -1 + assert data.shape[-3] == self._prognostic_mean_np.shape[0] + tensor_mean = tensor_mean.reshape(expand_var_dim) + tensor_std = tensor_std.reshape(expand_var_dim) + + unnorm = data * tensor_std + tensor_mean + unnorm = torch.where(self.wet_mask.to(data.device) == 0, fill_value, unnorm) + unnorm = unnorm.to(data.dtype) + return unnorm + + def unnormalize_tensor_boundary( + self, data: torch.Tensor, fill_value=float("nan") + ) -> torch.Tensor: + """Unnormalize boundary tensor.""" + tensor_mean = torch.from_numpy(self._boundary_mean_np).to( + data.device, data.dtype + ) + tensor_std = torch.from_numpy(self._boundary_std_np).to(data.device, data.dtype) + + expand_var_dim = [1] * data.ndim + expand_var_dim[-3] = -1 + assert data.shape[-3] == self._boundary_mean_np.shape[0] + tensor_mean = tensor_mean.reshape(expand_var_dim) + tensor_std = tensor_std.reshape(expand_var_dim) + + unnorm = data * tensor_std + tensor_mean + unnorm = torch.where( + self.wet_mask_surface.to(data.device) == 0, fill_value, unnorm + ) + unnorm = unnorm.to(data.dtype) + return unnorm + + +@dataclasses.dataclass +class LoadStats: + """Captures stats about loading a single TrainData object.""" + + load_time_seconds: float + + @classmethod + def accumulated(cls, stats: list["LoadStats"]) -> "LoadStats": + """Accumulate the stats across multiple LoadStats objects in a batch.""" + return cls(sum(s.load_time_seconds for s in stats)) + + +def compact_dataset(ds: xr.Dataset) -> xr.Dataset: + data = ds.copy() + + var_groups = defaultdict(list) + for key in data.keys(): + if "_lev_" in (k := str(key)): + base_name = k.split("_lev_")[0] + var_groups[base_name].append(k) + + def _parse_level(x) -> float: + return float(x.split("_lev_")[1].replace("_", ".")) + + for base_var, vars_ in var_groups.items(): + sorted_vars = sorted(vars_, key=_parse_level) + levels = [_parse_level(var) for var in sorted_vars] + if hasattr(data, "lev"): + levels = data.lev.values + da = xr.concat([data[var] for var in sorted_vars], dim="lev").assign_coords( + lev=("lev", levels) + ) + data[base_var] = da + data = data.drop_vars(vars_) + + return data + + +def stack_levels( + data: xr.Dataset, + dataset_spec: DatasetSpec, +) -> xr.Dataset: + """Reassemble a flattened OM4 dataset into analysis-ready, depth-stacked form. + + Inverts the preprocessing flattening so downstream analysis does not have to: + per-level prognostic channels (``thetao_0`` ...) become ``thetao(lev, ...)`` + and the per-level ``mask_i`` become a single stacked ``wetmask``. Grid + coordinates are preserved. This is the dataset-level counterpart to the eval + writer's reassembly, for putting ground-truth inputs on the same footing as + predictions. + + Implemented as the inverse of `with_level_index_vars` followed by the existing + `compact_dataset` (stacks the ``_lev_`` form) and `unflatten_masks`. + """ + data = with_depth_value_vars(data, dataset_spec) + data = compact_dataset(data) + if dataset_spec.mask_vars[0] in data.variables: + data = unflatten_masks(data, dataset_spec=dataset_spec) + return data diff --git a/samudra/utils/device.py b/samudra/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..621a5b720d4837d740fb9b05035adc7000108918 --- /dev/null +++ b/samudra/utils/device.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch + +_CHOSEN_DEVICE: torch.device | None = None + + +def autocast(enabled: bool, dtype: torch.dtype) -> torch.autocast: + if using_gpu(): + return torch.autocast("cuda", dtype=dtype, enabled=enabled) + else: + return torch.autocast("cpu", dtype=dtype, enabled=enabled) + + +def using_gpu() -> bool: + return get_device().type == "cuda" + + +# TODO(jder): We'd like to remove this (and other global singletons) +# by moving to some training framework + context objects once we have +# benchmarking set up. +def set_device(device: torch.device) -> None: + global _CHOSEN_DEVICE + _CHOSEN_DEVICE = device + + +def get_device() -> torch.device: + """The primary device for inference. + + We avoid setting this as torch.device so we can choose when + to shuttle tensors to the device. + """ + if _CHOSEN_DEVICE is not None: + return _CHOSEN_DEVICE + + if torch.cuda.is_available(): + return torch.device("cuda") + else: + return torch.device("cpu") diff --git a/samudra/utils/distributed.py b/samudra/utils/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..6c15821a429df86985a9366e5326b4c0e14bdcab --- /dev/null +++ b/samudra/utils/distributed.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import builtins +import datetime +import logging +import os +import random + +logger = logging.getLogger(__name__) + +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.distributed as dist + +from samudra.config import DistributedConfig + + +def set_seed(seed): + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + random.seed(seed) + cudnn.benchmark = ( + True # False # Set to True for better performance but lose reproducibility + ) + # cudnn.deterministic = True + # torch.use_deterministic_algorithms(True) + + +builtin_print = builtins.print + + +def print(*args, **kwargs): + force = kwargs.pop("force", False) + force = force or (get_world_size() > 8) + if force: + now = datetime.datetime.now().time() + builtin_print(f"[{now}] ", end="") # print with time stamp + builtin_print(*args, **kwargs) + + +def suppress_prints(is_master): + """This function disables printing when not in master process.""" + if not is_master: + builtins.print = print + + +def suppress_logging(is_master): + """Suppress logging for non-master processes.""" + if not is_master: + # Get root logger + root = logging.getLogger() + root.setLevel(logging.WARNING) + + # Also suppress any existing handlers on the root logger + for handler in root.handlers: + handler.setLevel(logging.WARNING) + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def init_distributed_mode() -> DistributedConfig: + cfg = DistributedConfig() + + if "RANK" in os.environ: + cfg.rank = int(os.environ["RANK"]) + cfg.gpu = int(os.environ["LOCAL_RANK"]) + cfg.world_size = int(os.environ["WORLD_SIZE"]) + cfg.dist_url = "env://" + elif "SLURM_PROCID" in os.environ: + cfg.rank = int(os.environ["SLURM_PROCID"]) + tasks_per_node = int(os.environ.get("SLURM_NTASKS_PER_NODE", 1)) + if tasks_per_node < torch.cuda.device_count(): + raise RuntimeError( + f"SLURM_NTASKS_PER_NODE is {tasks_per_node}, but this node has {torch.cuda.device_count()} GPUs; the remainder will be idle. " + "You probably want to allocate the same number of GPUs and tasks per node, or else use `torchrun` on a single node; please see CONTRIBUTING.md.", + ) + elif tasks_per_node > torch.cuda.device_count(): + raise RuntimeError( + f"SLURM_NTASKS_PER_NODE is {tasks_per_node}, but this node has {torch.cuda.device_count()} GPUs, so there aren't enough " + "GPUs for all the tasks. You probably want to allocate the same number of GPUs and tasks per node; please see CONTRIBUTING.md." + ) + n_nodes = int(os.environ["SLURM_NNODES"]) + cfg.world_size = tasks_per_node * n_nodes + cfg.gpu = cfg.rank % torch.cuda.device_count() + + if "MASTER_ADDR" in os.environ and "MASTER_PORT" in os.environ: + cfg.dist_url = "env://" + else: + cfg.dist_url = "tcp://localhost:40000" + else: + raise RuntimeError( + "Distributed mode requires SLURM or RANK environment variables;" + " did you mean to use `torchrun`?" + ) + + cfg.dist_backend = "nccl" + torch.distributed.init_process_group( + backend=cfg.dist_backend, + init_method=cfg.dist_url, + world_size=cfg.world_size, + rank=cfg.rank, + device_id=cfg.gpu, + ) + torch.cuda.set_device(cfg.gpu) + logger.info( + f"| distributed init (rank {cfg.rank}), gpu {cfg.gpu}, " + f"world_size {cfg.world_size}, dist_url {cfg.dist_url}" + ) + torch.distributed.barrier() + suppress_prints(cfg.rank == 0) + suppress_logging(cfg.rank == 0) + + return cfg + + +def all_reduce_mean(x): + world_size = get_world_size() + if world_size > 1: + dist.all_reduce(x) + x /= world_size + + return x diff --git a/samudra/utils/ema.py b/samudra/utils/ema.py new file mode 100644 index 0000000000000000000000000000000000000000..3f2d113bda7adc95deb580ca94c5d75f09c23bd0 --- /dev/null +++ b/samudra/utils/ema.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: 2022 Machine Vision and Learning Group, LMU Munich +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 +# SPDX-License-Identifier: MIT + +""" +Exponential Moving Average (EMA) module. + +Copied from https://github.com/CompVis/latent-diffusion/blob/main/ldm/modules/ema.py +and modified. + +MIT License + +Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from samudra.models.base import BaseModel + + +class EMATracker: + """ + Exponential Moving Average (EMA) tracker. + + This tracks the moving average of the parameters of a model, and has methods + that can be used to temporarily replace the parameters of the model with its EMA. + """ + + def __init__( + self, + model: torch.nn.Module, + decay: float = 0.999, + faster_decay_at_start: bool = True, + ): + """ + Create a new EMA tracker. + + Args: + model: The model whose parameters should be tracked. + decay: The decay rate of the moving average. + faster_decay_at_start: Whether to use the number of updates to determine + the decay rate. If True, the decay rate will be min(decay, (1 + + num_updates) / (10 + num_updates)). If False, the decay rate + will be decay. + """ + if decay < 0.0 or decay > 1.0: + raise ValueError("Decay must be between 0 and 1") + + self.decay = torch.tensor(decay, dtype=torch.float32) + self.cur_decay = torch.tensor(decay, dtype=torch.float32) + self._faster_decay_at_start = faster_decay_at_start + self.num_updates = torch.tensor(0, dtype=torch.int) + + self._ema_params = {} + + for name, p in model.named_parameters(): + if p.requires_grad: + ema_name = self._get_ema_name(name) + self._ema_params[ema_name] = p.clone().detach().data + + self._stored_params: list[torch.Tensor] = [] + + def _get_ema_name(self, name: str) -> str: + # remove as '.'-character is not allowed in buffers + # And make ourselves agnostic to module vs not + return name.removeprefix("module.").replace(".", "") + + def __eq__(self, other: object) -> bool: + if not isinstance(other, EMATracker): + return False + + def all_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return torch.all(a == b).item() # type: ignore + + return ( + all_equal(self.decay, other.decay) + and all_equal(self.num_updates, other.num_updates) + and self._faster_decay_at_start == other._faster_decay_at_start + and all( + all_equal(self._ema_params[k], other._ema_params[k]) + for k in self._ema_params + ) + and all( + all_equal(a, b) + for a, b in zip(self._stored_params, other._stored_params, strict=True) + ) + ) + + def __call__(self, model: BaseModel | nn.parallel.DistributedDataParallel): + """ + Update the moving average of the parameters. + + Does not mutate the input, only updates the moving average. + + Args: + model: The model whose parameters should be updated. Should be a model + specified identically to the one passed when this object was + instantiated. + """ + decay = self.decay + + self.num_updates += 1 + if self._faster_decay_at_start: + decay = torch.min( + self.decay, (1 + self.num_updates) / (10 + self.num_updates) + ) + self.cur_decay = decay + with torch.no_grad(): + module_parameters = dict(model.named_parameters()) + + for key in module_parameters: + ema_name = self._get_ema_name(key) + if module_parameters[key].requires_grad: + self._ema_params[ema_name] = self._ema_params[ema_name].type_as( + module_parameters[key] + ) + # EMA_params = decay * EMA_params + (1 - decay) * model_params + self._ema_params[ema_name].sub_( + (1.0 - decay) + * (self._ema_params[ema_name] - module_parameters[key]) + ) + elif ema_name in self._ema_params: + raise ValueError( + f"Expected model parameter {key} to require gradient, " + "but it does not" + ) + + def copy_to(self, model: BaseModel | nn.parallel.DistributedDataParallel): + """ + Copy the averaged parameters to the model, overwriting its values. + """ + m_param = dict(model.named_parameters()) + for key in m_param: + if m_param[key].requires_grad: + m_param[key].data.copy_(self._ema_params[self._get_ema_name(key)].data) + else: + assert self._get_ema_name(key) not in self._ema_params + + def store(self, parameters: Iterable[nn.Parameter]): + """ + Save the current parameters for restoring later. + + Args: + parameters: The parameters to be stored for later restoration by `restore` + """ + self._stored_params = [param.clone() for param in parameters] + + def restore(self, parameters: Iterable[nn.Parameter]): + """ + Restore the parameters stored with the `store` method. + + Useful to validate the model with EMA parameters without affecting the + original optimization process. Store the parameters before the + `copy_to` method. After validation (or model saving), use this to + restore the former parameters. + + Args: + parameters: The parameters to be updated with the values stored by `store` + """ + for c_param, param in zip(self._stored_params, parameters, strict=True): + param.data.copy_(c_param.data) + + self._stored_params = [] + + def get_state(self, include_ema_params: bool): + """ + Get the state of the EMA tracker. + + Args: + include_ema_params: Whether to include the EMA parameters in the state. + You probably want this to be True when saving a checkpoint meant to resume + training (i.e. not an EMA checkpoint) or False when saving a checkpoint for + evaluation (i.e. a checkpoint where the model weights are the EMA weights). + + Returns: + The state of the EMA tracker. + """ + state = { + "decay": self.decay, + "num_updates": self.num_updates, + "faster_decay_at_start": self._faster_decay_at_start, + } + if include_ema_params: + state["ema_params"] = self._ema_params + return state + + @classmethod + def from_state(cls, state, model: torch.nn.Module) -> "EMATracker": + """ + Create an EMA tracker from a state. + + Args: + state: The state of the EMA tracker. + model: The model whose parameters should be tracked, used to + initialize the EMA weights if they are not provided in the state. + + Returns: + The EMA tracker. + """ + ema = cls(model, float(state["decay"]), state["faster_decay_at_start"]) + ema.num_updates = state["num_updates"] + if ema_params := state.get("ema_params"): + unexpected_keys = ema_params.keys() - ema._ema_params.keys() + missing_keys = ema._ema_params.keys() - ema_params.keys() + assert not unexpected_keys and not missing_keys, ( + f"EMA parameters keys do not match. " + f"This is likely due to a mismatch between the model and checkpoint. " + f"Unexpected keys: {unexpected_keys}" + ) + ema._ema_params = ema_params + return ema diff --git a/samudra/utils/llc.py b/samudra/utils/llc.py new file mode 100644 index 0000000000000000000000000000000000000000..02f75f46e99c43c1b197b0f15f10cec3fe94e5ee --- /dev/null +++ b/samudra/utils/llc.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import xarray as xr + +from samudra.constants import DatasetSpec + + +def _rename_llc_level_index_vars(ds: xr.Dataset) -> xr.Dataset: + """Rename LLC statistics variables to the common level-index format. + + For example, ``Theta_lev_0`` becomes ``Theta_0``. + """ + rename_map = { + name: f"{var}_{lev}" + for name in ds.variables + if isinstance(name, str) and "_lev_" in name + for var, lev in [name.split("_lev_", maxsplit=1)] + } + return ds.rename(rename_map) if rename_map else ds + + +def _flatten_llc_level_vars( + data: xr.Dataset, + *, + dataset_spec: DatasetSpec, +) -> xr.Dataset: + """Flatten LLC level dimensions into level-indexed variables. + + For example, ``Theta(time, lev, y, x)`` becomes ``Theta_0(time, y, x)``, + ``Theta_50(time, y, x)``, and the remaining configured levels. + """ + data_copy = data.copy() + for name in list(data_copy.data_vars): + if not isinstance(name, str): + continue + if "lev" not in data_copy[name].dims: + continue + + n_levels = data_copy[name].sizes["lev"] + expected_levels = len(dataset_spec.depth_i_levels) + if n_levels != expected_levels: + raise ValueError( + f"Expected {expected_levels} levels for LLC variable {name}, got " + f"{n_levels}" + ) + + for index, lev in enumerate(dataset_spec.depth_i_levels): + data_copy[f"{name}_{lev}"] = data_copy[name].isel(lev=index) + data_copy = data_copy.drop_vars(name) + + return data_copy + + +def _var_without_level(var_name: str) -> str: + """Return the raw LLC variable name for a level-indexed model variable. + + For example, ``Theta_50`` becomes ``Theta`` while ``oceQnet`` is unchanged. + """ + suffix = var_name.rsplit("_", maxsplit=1)[-1] + return var_name.rsplit("_", maxsplit=1)[0] if suffix.isdigit() else var_name + + +def _preferred_available_var( + data: xr.Dataset, candidates: tuple[str, ...] +) -> str | None: + """Select the first candidate present in an LLC dataset. + + For example, candidates ``("mask_w", "hFacW")`` select ``mask_w`` when + both variables are available. + """ + return next((name for name in candidates if name in data.data_vars), None) + + +def _llc_staggered_mask_vars( + data: xr.Dataset, + requested_data_vars: set[str], +) -> set[str]: + """Find staggered masks needed by the requested LLC variables. + + For example, requesting ``U`` from data containing ``mask_w`` adds + ``mask_w`` to the variables retained during canonicalization. + """ + mask_vars = set() + if {"U", "oceTAUX"} & requested_data_vars: + if mask_var := _preferred_available_var(data, ("mask_w", "hFacW")): + mask_vars.add(mask_var) + if {"V", "oceTAUY"} & requested_data_vars: + if mask_var := _preferred_available_var(data, ("mask_s", "hFacS")): + mask_vars.add(mask_var) + return mask_vars + + +def canonicalize_llc_datasets( + data: xr.Dataset, + means: xr.Dataset, + stds: xr.Dataset, + *, + face: int, + i_start: int, + i_end: int, + j_start: int, + j_end: int, + dataset_spec: DatasetSpec, +) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + """Standardize raw LLC inputs to the common non-compact loader layout. + + For example, a selected ``Theta(time, face, k, j, i)`` crop becomes + level-indexed ``Theta_0(time, y, x)``, ``Theta_50(time, y, x)``, and the + remaining configured levels; statistics names undergo the same renaming. + """ + data_copy = data.copy() + + requested_data_vars = { + _var_without_level(var_name) + for var_name in ( + dataset_spec.prognostic_var_names + dataset_spec.boundary_var_names + ) + } + center_mask_var = _preferred_available_var(data_copy, ("mask_c", "hFacC")) + requested_data_vars.update( + [dataset_spec.mask_all_levels_var, *dataset_spec.mask_vars] + ) + if center_mask_var is not None: + requested_data_vars.add(center_mask_var) + requested_data_vars.update(_llc_staggered_mask_vars(data_copy, requested_data_vars)) + data_copy = data_copy[ + [name for name in data_copy.data_vars if name in requested_data_vars] + ] + + if "face" in data_copy.dims or "face" in data_copy.coords: + data_copy = data_copy.sel(face=face, drop=True) + + spatial_indexers = {} + if "i" in data_copy.dims: + spatial_indexers["i"] = slice(i_start, i_end) + if "i_g" in data_copy.dims: + spatial_indexers["i_g"] = slice(i_start, i_end) + if "j" in data_copy.dims: + spatial_indexers["j"] = slice(j_start, j_end) + if "j_g" in data_copy.dims: + spatial_indexers["j_g"] = slice(j_start, j_end) + if spatial_indexers: + data_copy = data_copy.isel(spatial_indexers) + + unstagger_map = { + "U": {"i_g": "i"}, + "V": {"j_g": "j"}, + "oceTAUX": {"i_g": "i"}, + "oceTAUY": {"j_g": "j"}, + "mask_w": {"i_g": "i"}, + "hFacW": {"i_g": "i"}, + "mask_s": {"j_g": "j"}, + "hFacS": {"j_g": "j"}, + } + for var_name, rename_dims in unstagger_map.items(): + if var_name not in data_copy.variables: + continue + used_renames = { + old: new + for old, new in rename_dims.items() + if old in data_copy[var_name].dims + } + if used_renames: + data_copy[var_name] = data_copy[var_name].rename(used_renames) + + remaining_staggered = [ + name + for name in data_copy.data_vars + if "i_g" in data_copy[name].dims or "j_g" in data_copy[name].dims + ] + if remaining_staggered: + raise ValueError( + "LLC canonicalization left staggered variables unresolved: " + f"{remaining_staggered}" + ) + + data_copy = data_copy.drop_vars(["i_g", "j_g"], errors="ignore") + rename_map = { + old: new + for old, new in { + "k": "lev", + "i": "x", + "j": "y", + }.items() + if old in data_copy.dims + or old in data_copy.variables + or old in data_copy.coords + } + if center_mask_var is not None: + rename_map[center_mask_var] = dataset_spec.mask_all_levels_var + if rename_map: + data_copy = data_copy.rename(rename_map) + + means_copy = _rename_llc_level_index_vars(means.copy()) + stds_copy = _rename_llc_level_index_vars(stds.copy()) + data_copy = _flatten_llc_level_vars(data_copy, dataset_spec=dataset_spec) + + return data_copy, means_copy, stds_copy diff --git a/samudra/utils/location.py b/samudra/utils/location.py new file mode 100644 index 0000000000000000000000000000000000000000..847baecfc57a1f94401c4dec27cc9e5907668830 --- /dev/null +++ b/samudra/utils/location.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Annotated, Any, Literal, Self +from urllib.parse import quote, urljoin, urlparse + +import xarray as xr +from pydantic import ( + BaseModel, + BeforeValidator, + WithJsonSchema, + model_serializer, + model_validator, +) + + +class UnresolvedLocation(BaseModel): + """Representation for a raw string in a Location config. + + This is expected to be a relative or absolute path. + It is flexibily interepreted as a url or local path depending + on what kind of location it is resolved against. + """ + + path: str + + @model_validator(mode="after") + def validate_path(self) -> Self: + if urlparse(self.path).scheme: + raise ValueError( + "Absolute urls are not supported, please use a " + "relative path or set type = 's3' or 'local'" + ) + return self + + @model_serializer + def seralize(self) -> Any: + return self.path + + +class ResolvedLocation(ABC): + """A location which is ready to be opened or resolved against.""" + + @abstractmethod + def open(self, chunks: dict[str, int] | None = None) -> xr.Dataset: + pass + + @abstractmethod + def resolve(self, location: "Location") -> "ResolvedLocation": + pass + + def __truediv__(self, other: "Location") -> "ResolvedLocation": + return self.resolve(other) + + +class S3Location(ResolvedLocation, BaseModel): + """An S3 bucket, assuming credentials in your environment. + + Set ``anon: true`` to read a public bucket without credentials (e.g. the + open datasets on OSN); leave it false to sign requests with the usual + ``AWS_*`` environment variables. + + For example: + ```yaml + data_location: + type: s3 + bucket: emulators + path: sd5313/OM4_highres/om4_halfdeg.zarr + ``` + """ + + type: Literal["s3"] = "s3" + endpoint_url: str | None = None + anon: bool = False + bucket: str + path: str + + def open(self, chunks: dict[str, int] | None = None) -> xr.Dataset: + # TODO(jder): could consider passing credentials here + # rather than relying on the environment + storage_options: dict[str, Any] = {"endpoint_url": self.endpoint_url} + if self.anon: + storage_options["anon"] = True + + return xr.open_dataset( + self.url(), + backend_kwargs={"storage_options": storage_options}, + engine="zarr", + chunks=chunks, + ) + + def url(self) -> str: + path = quote(self.path.lstrip("/")) + bucket = quote(self.bucket, safe="") + return f"s3://{bucket}/{path}" + + def resolve(self, location: "Location") -> "ResolvedLocation": + if isinstance(location, UnresolvedLocation): + return S3Location( + endpoint_url=self.endpoint_url, + anon=self.anon, + bucket=self.bucket, + path=urljoin(self.path + "/", location.path), + ) + return location + + def __str__(self) -> str: + return self.url() + + +class LocalLocation(ResolvedLocation, BaseModel): + """A local absolute filesystem path. + + For example: + ```yaml + data_location: + type: local + path: /path/to/data + ``` + """ + + type: Literal["local"] = "local" + path: Path + + @model_validator(mode="after") + def validate_path(self) -> Self: + if not self.path.is_absolute(): + raise ValueError( + "Locations with type: 'local' must be absolute. " + "For relative paths, use a string instead of a structured location. " + "i.e. 'my/relative/path' instead of " + "{type: 'local', path: 'my/relative/path'}" + ) + return self + + def open(self, chunks: dict[str, int] | None = None) -> xr.Dataset: + engine = "netcdf4" if self.path.suffix == ".nc" else "zarr" + return xr.open_dataset(self.path, engine=engine, chunks=chunks) + + def resolve(self, location: "Location") -> "ResolvedLocation": + if isinstance(location, UnresolvedLocation): + return LocalLocation(path=self.path / location.path) + return location + + def __str__(self) -> str: + return str(self.path) + + +def string_to_unresolved(data: Any) -> Any: + """Turns a string into an UnresolvedLocation.""" + # TODO(jder): we could support other fsspec or universal_pathlib URLs here + if isinstance(data, str): + return UnresolvedLocation(path=data) + return data + + +Location = Annotated[ + Annotated[UnresolvedLocation, WithJsonSchema({"type": "string"})] + | S3Location + | LocalLocation, + BeforeValidator(string_to_unresolved), +] diff --git a/samudra/utils/logging.py b/samudra/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..27b214aef82a5a0257ada6eed275dea48eb797cd --- /dev/null +++ b/samudra/utils/logging.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import datetime +import functools +import logging +import resource +import sys +import time +import traceback +import warnings +from collections import defaultdict, deque +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import torch +from torch.nn.parallel import DistributedDataParallel +from torchinfo import summary + +from samudra.constants import Boundary, Prognostic + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from samudra.datasets import TrainData, TrainDataLoader + from samudra.models.base import BaseModel + from samudra.utils.ctx import GridContext + + +def handle_logging(debug: bool, output_dir: Path): + # Set up logging + logger = logging.getLogger() # Use the root logger or specify a name if needed + logger.setLevel(logging.DEBUG if debug else logging.INFO) + fmt = logging.Formatter("%(asctime)s - %(levelname)s - %(module)s - %(message)s") + + # STDOUT handler + stdout_handler = logging.StreamHandler(sys.stdout) + stdout_handler.setLevel(logging.DEBUG if debug else logging.INFO) + stdout_handler.setFormatter(fmt) + logger.addHandler(stdout_handler) + + # Add experiment log file handler + experiment_log_path = output_dir / "experiment.log" + experiment_handler = logging.FileHandler(experiment_log_path) + experiment_handler = logging.FileHandler(experiment_log_path) + experiment_handler.setLevel(logging.INFO) # Capture info and above + experiment_handler.setFormatter(fmt) + logger.addHandler(experiment_handler) + + # Add separate error log file handler + error_log_path = output_dir / "error.log" + error_handler = logging.FileHandler(error_log_path) + error_handler.setLevel(logging.WARNING) # Capture warnings and errors + error_handler.setFormatter(fmt) + logger.addHandler(error_handler) + + +def handle_warnings(): + def warning_handler(message, category, filename, lineno, file=None, line=None): + logger.info("\n=== Warning Details ===") + logger.info(f"Message: {message}") + logger.info(f"Category: {category}") + logger.info(f"File: {filename}") + logger.info(f"Line: {lineno}") + logger.info("\nFull stack trace:") + stack = traceback.extract_stack()[:-1] # Remove current frame + for frame in stack: + logger.info( + f' File "{frame.filename}", line {frame.lineno}, in {frame.name}' + ) + if frame.line: + logger.info(f" {frame.line}") + logger.info("=====================\n") + + warnings.showwarning = warning_handler + + +class SmoothedValue: + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque: deque[float] = deque(maxlen=window_size) + self.total: float = 0.0 + self.count: int = 0 + self.fmt: str = fmt + + def update(self, value, n=1): + self.deque.append(value) + self.count += n + self.total += value * n + + @property + def median(self): + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + return self.total / self.count + + @property + def max(self): + return max(self.deque) + + @property + def value(self): + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +class MetricLogger: + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, n=1, **kwargs): + for k, v in kwargs.items(): + if v is None: + continue + if isinstance(v, torch.Tensor): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v, n=n) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{attr}'" + ) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append(f"{name}: {str(meter)}") + return self.delimiter.join(loss_str) + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every( + self, + data_loader: "TrainDataLoader", + print_freq, + header=None, + ): + i = 0 + if not header: + header = "" + start_time = time.perf_counter() + end = time.perf_counter() + iter_time = SmoothedValue(fmt="{value:.3f}({avg:.3f})", window_size=print_freq) + data_wait_time = SmoothedValue( + fmt="{value:.3f}({avg:.3f})", window_size=print_freq + ) + data_load_time = SmoothedValue( + fmt="{value:.3f}({avg:.3f})", window_size=print_freq + ) + self.meters["iter_time"] = iter_time + self.meters["data_wait_time"] = data_wait_time + self.meters["data_load_time"] = data_load_time + space_fmt = ":" + str(len(str(len(data_loader)))) + "d" + log_msg_list: list[str] = [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + ] + log_msg_list.append("max cpu mem: {cpu_memory:.0f}") + if torch.cuda.is_available(): + log_msg_list.append("max gpu mem: {gpu_memory:.0f}") + log_msg = self.delimiter.join(log_msg_list) + KB = 1024.0 + MB = 1024.0 * 1024.0 + for obj in data_loader: + data_wait_time.update(time.perf_counter() - end) + if obj.load_stats is not None: + data_load_time.update(obj.load_stats.load_time_seconds) + yield obj + iter_time.update(time.perf_counter() - end) + if i % print_freq == 0 or i == len(data_loader) - 1: + eta_seconds = iter_time.global_avg * (len(data_loader) - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + named_metrics = dict( + eta=eta_string, + meters=str(self), + cpu_memory=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / KB, + ) + if torch.cuda.is_available(): + named_metrics["gpu_memory"] = torch.cuda.max_memory_allocated() / MB + + logger.info(log_msg.format(i, len(data_loader), **named_metrics)) + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + i += 1 + end = time.perf_counter() + total_time = time.perf_counter() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + logger.info( + f"{header} Total time: {total_time_str} " + f"({total_time / len(data_loader):.4f} s / it)" + ) + + +class _ForwardOnceWrapper(torch.nn.Module): + """Wrapper that redirects forward() to forward_once() for torchinfo summary.""" + + def __init__( + self, + model: "BaseModel | DistributedDataParallel", + ctx: "GridContext", + ) -> None: + super().__init__() + self._underlying: BaseModel = getattr(model, "module", model) # type: ignore + self._ctx = ctx + + def forward(self, prognostic: Prognostic, boundary: Boundary) -> torch.Tensor: + # Get the underlying model (handles DDP wrapping) + return self._underlying.forward_once(prognostic, boundary, self._ctx) + + +def get_model_summary( + model: "BaseModel | DistributedDataParallel", data: "TrainData | None", debug: bool +) -> None: + model_parameters = filter(lambda p: p.requires_grad, model.parameters()) + params = sum([np.prod(p.size()) for p in model_parameters]) + logger.info(f"Number of parameters: {params}") + depth = 10 if debug else 2 + # we pass verbose = 0 because we log the summary ourselves + if data is not None: + # TrainData is a complex wrapper that torchinfo cannot traverse. + # Extract the initial prognostic + boundary and wrap the model to + # use forward_once. + prog_tensor, boundary_tensor = data.get_initial_input() + wrapper = _ForwardOnceWrapper(model, data.ctx) + logger.info( + summary( + wrapper, + input_data=[prog_tensor, boundary_tensor], + depth=depth, + verbose=0, + ) + ) + else: + logger.info(summary(model, depth=depth, verbose=0)) + + +def elapsed(func=None, *, level: int = logging.INFO): + """Log the time taken to execute a function. + + Implementation inspired by this blog post: + https://pybit.es/articles/decorator-optional-argument/ + + Args: + func (callable): The function to decorate. If None, returns a partial + function with the log_level argument. + level (int): The logging level to use. Default is logging.INFO. + """ + if func is None: + return functools.partial(elapsed, level=level) + + @functools.wraps(func) + def wrapper(*args, **kwargs): + start_time = time.perf_counter() + result = func(*args, **kwargs) + end_time = time.perf_counter() + logger.log( + level, + "%s took %.4f seconds", + func.__qualname__, + end_time - start_time, + ) + return result + + return wrapper diff --git a/samudra/utils/loss.py b/samudra/utils/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..e2c9c56218672b75a4af420ff137f8d12bde635f --- /dev/null +++ b/samudra/utils/loss.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Callable +from typing import Literal, Protocol, assert_never + +import torch +import torch.nn.functional as F +from jaxtyping import Float + +from samudra.utils.ctx import GridContext + +LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + + +class LossFnWithContext(Protocol): + def __call__( + self, + pred: torch.Tensor, + target: torch.Tensor, + ctx: GridContext, + ) -> torch.Tensor: ... + + +LossMetric = Literal[ + "mse", + "mae", + "mse_mae", + "mse_diff_weighted", +] + + +def loss_fn_from_metric(metric: LossMetric) -> LossFnWithContext: + match metric: + case "mse": + loss_fn: LossFn = decomposed_mse + case "mae": + loss_fn = decomposed_mae + case "mse_mae": + loss_fn = decomposed_mse_mae + case "mse_diff_weighted": + loss_fn = decomposed_mse_diff_weighted + case _: + assert_never(metric) + + def loss_fn_with_ctx( + pred: torch.Tensor, + target: torch.Tensor, + ctx: GridContext, + ) -> torch.Tensor: + wet = ctx.label_mask.to(device=pred.device) + pred = pred * wet + target = target * wet + return loss_fn(pred, target) + + return loss_fn_with_ctx + + +def decomposed_mse(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Standard MSE loss (l2) computed per channel.""" + return F.mse_loss(pred, target, reduction="none").mean(dim=(0, 2, 3)) + + +def decomposed_mae(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Standard MAE loss (l1) computed per channel.""" + return F.l1_loss(pred, target, reduction="none").mean(dim=(0, 2, 3)) + + +# TODO(alxmrs): This used to assume that hist=1; it may need to be fixed in the future. +def decomposed_mse_diff_weighted( + pred: torch.Tensor, target: torch.Tensor +) -> torch.Tensor: + """MSE loss with weighted differences.""" + # Compute standard MSE + mse = F.mse_loss(pred, target, reduction="none") + + # Weight the differences more heavily + diff_weight = 2.0 # Adjustable weight factor + diff_mse = ( + F.mse_loss( + pred[:, 1:] - pred[:, :-1], target[:, 1:] - target[:, :-1], reduction="none" + ) + * diff_weight + ) + + # Combine losses + combined_loss = torch.cat([mse[:, :1], diff_mse], dim=1) + return combined_loss.mean(dim=(0, 2, 3)) + + +def decomposed_mse_scaled( + pred: torch.Tensor, target: torch.Tensor, scaling: torch.Tensor +) -> torch.Tensor: + """MSE loss with scaled residuals.""" + scaled_pred = pred * scaling.view(1, -1, 1, 1) + scaled_target = target * scaling.view(1, -1, 1, 1) + return F.mse_loss(scaled_pred, scaled_target, reduction="none").mean(dim=(0, 2, 3)) + + +def decomposed_mse_mae(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Combined MSE and MAE loss.""" + mse = F.mse_loss(pred, target, reduction="none") + mae = F.l1_loss(pred, target, reduction="none") + combined = (mse + mae) / 2 + return combined.mean(dim=(0, 2, 3)) + + +def _spatial_gradients( + tensor: torch.Tensor, *, pad_mode: str +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute forward differences along y and x axes with configurable x padding.""" + grad_y = tensor[:, :, 1:, :] - tensor[:, :, :-1, :] + grad_y = F.pad(grad_y, (0, 0, 0, 1), mode="constant") + + padded_x = F.pad(tensor, (0, 1, 0, 0), mode=pad_mode) + grad_x = padded_x[:, :, :, 1:] - padded_x[:, :, :, :-1] + + return grad_y, grad_x + + +def gradient_l1_loss( + pred: torch.Tensor, target: torch.Tensor, pad_mode: str +) -> torch.Tensor: + """L1 loss on spatial gradients, averaged per channel.""" + pred_grad_y, pred_grad_x = _spatial_gradients(pred, pad_mode=pad_mode) + target_grad_y, target_grad_x = _spatial_gradients(target, pad_mode=pad_mode) + + grad_loss_y = F.l1_loss(pred_grad_y, target_grad_y, reduction="none") + grad_loss_x = F.l1_loss(pred_grad_x, target_grad_x, reduction="none") + + grad_loss = (grad_loss_y.mean(dim=(0, 2, 3)) + grad_loss_x.mean(dim=(0, 2, 3))) / 2 + return grad_loss + + +def decomposed_mae_gradient_weighted( + pred: torch.Tensor, + target: torch.Tensor, + gradient_weight: float, + pad_mode: str = "constant", +) -> torch.Tensor: + """MAE loss with spatial gradient matching penalty.""" + mae_per_channel = decomposed_mae(pred, target) + grad_loss = gradient_l1_loss(pred, target, pad_mode) + return mae_per_channel + gradient_weight * grad_loss + + +class DynamicLoss: + """A loss function that scales each channel to contribute equally to the loss. + + This uses a rolling estimate of the loss of each channel to scale each + channel's loss, discouraging the model from focusing on only a few channels. + + See: https://openathena.slack.com/archives/C08CYM42DT3/p1752275713570969 + """ + + N_WINDOW = 25 + """Rolling window size to average over. (~number of steps)""" + + def __init__( + self, + loss_fn: LossFnWithContext, + *, + limit: float | None, + device: torch.device, + num_channels: int, + ): + self.loss_fn = loss_fn + self._device = device + self._per_channel_scale: Float[torch.Tensor, " var"] = torch.ones( + num_channels, device=self._device + ) + self._limit = limit + + def __call__( + self, + pred: Float[torch.Tensor, "batch hist*var lat lon"], + target: Float[torch.Tensor, "batch hist*var lat lon"], + ctx: GridContext, + ) -> Float[torch.Tensor, " hist*var"]: + loss_with_history_channels: Float[torch.Tensor, " hist*var"] = self.loss_fn( + pred, target, ctx + ) + # Channels are time-major: (hist+1) * var. + scaled_loss_including_history_dimension: Float[torch.Tensor, "hist var"] = ( + loss_with_history_channels.reshape(-1, self._per_channel_scale.shape[0]) + * self._per_channel_scale + ) + return scaled_loss_including_history_dimension.reshape(-1) + + def update( + self, + loss_per_channel: Float[torch.Tensor, " hist*var"], + ) -> None: + """Given the unscaled per-channel loss, update the per-channel scale.""" + # Local import is needed to prevent a circular import error. + from samudra.utils.distributed import all_reduce_mean, get_world_size + + loss = loss_per_channel.detach() + loss = torch.where(loss == 0, 1e-8, loss) + new_target_weights_with_history: Float[torch.Tensor, " hist*var"] = 1.0 / loss + # Reshape from channels * history to channels + # by averaging along the `hist` dimension + new_target_weights: Float[torch.Tensor, " var"] = ( + new_target_weights_with_history.reshape( + -1, self._per_channel_scale.shape[0] + ).mean(dim=0) + ) + + if get_world_size() > 1: + all_reduce_mean(new_target_weights) + + if self._limit is not None: + min_scale = new_target_weights.min() + max_scale = min_scale * self._limit + new_target_weights = new_target_weights.clamp(min_scale, max_scale) + + self._per_channel_scale = ( + self._per_channel_scale * (DynamicLoss.N_WINDOW - 1) + new_target_weights + ) / DynamicLoss.N_WINDOW + + def loss_scale_per_channel(self) -> Float[torch.Tensor, " var"]: + return self._per_channel_scale + + # new methods for saving and loading state + def state_dict(self) -> dict[str, torch.Tensor]: + """Return state dictionary for checkpointing.""" + return {"per_channel_scale": self._per_channel_scale.detach().cpu()} + + def load_state_dict(self, state: dict[str, torch.Tensor]) -> None: + """Load state from ``state_dict``.""" + if "per_channel_scale" in state: + self._per_channel_scale = state["per_channel_scale"].to(self._device) + + +class GradientLoss: + """Combine a base loss with a gradient matching penalty. + + Applies the provided per-channel loss metric then adds an L1 penalty on + spatial gradients, scaled by ``gradient_weight``. + """ + + def __init__( + self, + loss_fn: LossFnWithContext, + *, + gradient_weight: float, + pad_mode: str, + ): + self.loss_fn = loss_fn + self._gradient_weight = gradient_weight + self._pad_mode = pad_mode + + def __call__( + self, + pred: Float[torch.Tensor, "batch hist*var lat lon"], + target: Float[torch.Tensor, "batch hist*var lat lon"], + ctx: GridContext, + ) -> Float[torch.Tensor, " hist*var"]: + base_loss = self.loss_fn(pred, target, ctx) + # Ensure mask is on the same device as pred for gradient computation + wet = ctx.label_mask.to(device=pred.device) + pred = pred * wet + target = target * wet + grad_loss = gradient_l1_loss(pred=pred, target=target, pad_mode=self._pad_mode) + return base_loss + self._gradient_weight * grad_loss diff --git a/samudra/utils/multiton.py b/samudra/utils/multiton.py new file mode 100644 index 0000000000000000000000000000000000000000..128293624b9a417844e2ebf6180684e4ff442aeb --- /dev/null +++ b/samudra/utils/multiton.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC +from typing import Self + + +class MultitonScope: + """ + Context manager for managing the scope of Multiton instances. + + You can create a fresh scope and make it active/inactive, for example: + + with MultitonScope() as scope: + # Create a new scope + instance = SomeMultiton.init_instance() + # Use the instance within this scope + ... + + or it can be disjoint: + + scope = MultitonScope() + with scope: + # Create a new scope + instance = SomeMultiton.init_instance() + + # later on + + with scope: + # gets the instance created above + instance = SomeMultiton.instance() + + """ + + def __init__(self): + self.scope = {} + + def __enter__(self): + # TODO(jder): Could make this thread/async-safe if needed. + self.previous_scope = Multiton._current_scope + Multiton._current_scope = self.scope + return self + + def __exit__(self, exc_type, exc_value, traceback): + assert Multiton._current_scope is self.scope, ( + "MultitonScope.__exit__ called without matching __enter__" + ) + Multiton._current_scope = self.previous_scope + + +class Multiton(ABC): + """ + Allows creation of a set of singletons within some scope, allowing clearing + and restoring them. + """ + + _current_scope: dict = {} + + def __new__(cls): + # Prevent direct instantiation + raise TypeError( + f"{cls} cannot be instantiated directly. Use init_instance() instead." + ) + + @classmethod + def get_instance(cls) -> Self: + """ + Get the instance of the Multiton class for the current scope. + """ + instance = cls._current_scope.get(cls) + if instance is None: + raise ValueError(f"{cls} not initialized") + return instance + + @classmethod + def init_instance(cls, *args, **kwargs) -> Self: + if Multiton._current_scope.get(cls) is not None: + raise ValueError(f"{cls} already initialized") + + instance = super().__new__(cls) + instance._initialize(*args, **kwargs) + Multiton._current_scope[cls] = instance + + return instance + + def _initialize(self, *args, **kwargs): + """ + Override this method to initialize the instance. + """ + raise NotImplementedError( + "Subclasses must override _initialize() to set up the instance." + ) diff --git a/samudra/utils/output.py b/samudra/utils/output.py new file mode 100644 index 0000000000000000000000000000000000000000..8552faa0d7823d10a96b2d91257eedd024e157a3 --- /dev/null +++ b/samudra/utils/output.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import torch +import xarray as xr + +from samudra.utils.ctx import GridContext + + +class TrainBatchOutput: + def __init__(self, loss: torch.Tensor, loss_per_channel: torch.Tensor): + self.loss = loss + self.loss_per_channel = loss_per_channel + + +class ValBatchOutput(TrainBatchOutput): + def __init__( + self, + loss: torch.Tensor, + loss_per_channel: torch.Tensor, + input_data: torch.Tensor, + target_data: torch.Tensor, + gen_data: torch.Tensor, + ctx: GridContext, + ): + super().__init__(loss, loss_per_channel) + assert target_data.shape == gen_data.shape + self.input_data = input_data + self.target_data = target_data + self.gen_data = gen_data + self.ctx = ctx + + +class ModelInferenceOutput: + def __init__( + self, + prediction: torch.Tensor, + target: torch.Tensor, + time: xr.DataArray, + ): + assert prediction.shape == target.shape + self.prediction = prediction + self.target = target + self.time = time diff --git a/samudra/utils/profiler.py b/samudra/utils/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..2992f87f0a42a88d184bd3fbb591977efff8ab71 --- /dev/null +++ b/samudra/utils/profiler.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import torch + + +class Profiler: + def __init__( + self, + output_dir: Path, + cuda_snapshot_frequency: int | None, + ) -> None: + self.output_dir = output_dir + self.cuda_snapshot_frequency = cuda_snapshot_frequency + + def start(self) -> None: + if self.cuda_snapshot_frequency is not None: + torch.cuda.memory._record_memory_history() + + def after_batch(self, num_batches_seen: int) -> None: + if self.cuda_snapshot_frequency is not None: + if num_batches_seen % self.cuda_snapshot_frequency == 0: + torch.cuda.memory._dump_snapshot( + str( + self.output_dir + / f"cuda_memory_snapshot_{num_batches_seen}.pickle" + ) + ) diff --git a/samudra/utils/samplers.py b/samudra/utils/samplers.py new file mode 100644 index 0000000000000000000000000000000000000000..16fe7ce30ecdc06096f289e46725fc83c8f7cc7a --- /dev/null +++ b/samudra/utils/samplers.py @@ -0,0 +1,345 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import itertools +import math +import random +from collections.abc import Callable, Hashable +from typing import TYPE_CHECKING, Self + +from torch.utils.data import BatchSampler, Sampler + +if TYPE_CHECKING: + from samudra.datasets import TorchTrainDataset + +type Batch = list[int] +type DdpStepChunk = tuple[Batch, ...] + + +class _SimpleSubsetSampler(Sampler): + def __init__(self, indices): + super().__init__() + self.indices = indices + + def __iter__(self): + return iter(self.indices) + + def __len__(self): + return len(self.indices) + + +class EquivalenceGroupBatchSampler(Sampler[Batch]): + """Groups indices into equivalence classes, batches within groups, and optionally shuffles. + + This sampler partitions dataset indices into groups. It creates batches within each group, + then chains them together. When shuffle=True, batches are globally shuffled each epoch + to avoid sequential group processing. + + Args: + groups: List of index lists, where each inner list contains indices belonging to + the same equivalence group. + batch_size: Number of samples per batch + shuffle: Whether to shuffle indices within groups and shuffle batches globally + drop_last: Whether to drop incomplete batches at the end of each group + seed: Random seed for deterministic shuffling + """ + + def __init__( + self, + groups: list[list[int]], + batch_size: int, + shuffle: bool = True, + drop_last: bool = False, + *, + seed: int, + ): + super().__init__() + self.groups = groups + self.batch_size = batch_size + self.shuffle = shuffle + self.drop_last = drop_last + self.seed = seed + self.epoch = 0 + + self._samplers = [ + BatchSampler( + _SimpleSubsetSampler(group), + batch_size=self.batch_size, + drop_last=self.drop_last, + ) + for group in self.groups + ] + + def set_epoch(self, epoch: int) -> None: + """Set the epoch for deterministic shuffling.""" + self.epoch = epoch + + @classmethod + def from_dataset_sizes( + cls, + dataset_sizes: list[int], + batch_size: int, + shuffle: bool = True, + drop_last: bool = False, + *, + seed: int, + ) -> Self: + """Create sampler from dataset sizes, treating each as a contiguous group. + + Args: + dataset_sizes: List of individual dataset sizes. Groups are created based on + cumulative boundaries, where each dataset forms its own equivalence group. + batch_size: Number of samples per batch + shuffle: Whether to shuffle indices within groups and shuffle batches globally + drop_last: Whether to drop incomplete batches at the end of each group + seed: Random seed for deterministic shuffling + """ + cumsum = 0 + groups = [] + for size in dataset_sizes: + groups.append(list(range(cumsum, cumsum + size))) + cumsum += size + return cls(groups, batch_size, shuffle, drop_last, seed=seed) + + @classmethod + def from_datasets( + cls, + datasets: list["TorchTrainDataset"], + group_key: Callable[["TorchTrainDataset"], Hashable], + batch_size: int, + shuffle: bool, + drop_last: bool, + *, + seed: int, + ) -> Self: + """Create sampler by grouping datasets using a key function. + + This factory method allows grouping datasets by arbitrary criteria (e.g., resolution, + regardless of other parameters like stride). Datasets with the same key are batched together. + + Args: + datasets: List of TorchTrainDataset instances to group + group_key: Callable that extracts grouping key from a dataset. + batch_size: Number of samples per batch + shuffle: Whether to shuffle indices within groups and shuffle batches globally + drop_last: Whether to drop incomplete batches at the end of each group + seed: Random seed for deterministic shuffling + + Examples: + - lambda ds: (ds._input_src.data.sizes['lat'], ds._input_src.data.sizes['lon']) # group by resolution + - lambda ds: ds._input_src.data.sizes['lat'] # group by latitude size only + + Returns: + EquivalenceGroupBatchSampler configured to group by the provided key + + Example: + >>> # Group datasets by resolution, allowing different strides to be batched together + >>> sampler = EquivalenceGroupBatchSampler.from_datasets( + ... datasets=dataset_list, + ... group_key=lambda ds: ds.prognostic_src.grid_size, + ... batch_size=32, + ... shuffle=True, + ... drop_last=True, + ... seed=15, + ... ) + """ + from collections import defaultdict + + # Group indices by their key + groups: dict[Hashable, list[int]] = defaultdict(list) + + cumsum = 0 + for ds in datasets: + key = group_key(ds) + assert isinstance(key, Hashable), "`group_key` must be hashable." + groups[key].extend(range(cumsum, cumsum + len(ds))) + cumsum += len(ds) + + # Sort by key for deterministic ordering across runs + sorted_groups = sorted(groups.items(), key=lambda x: x[0]) # type: ignore + group_indices = [indices for _, indices in sorted_groups] + + return cls(group_indices, batch_size, shuffle, drop_last, seed=seed) + + def __iter__(self): + rng = random.Random(self.seed + self.epoch) + all_batches: list[Batch] = [] + for group in self.groups: + indices = list(group) + if self.shuffle: + rng.shuffle(indices) + + for start in range(0, len(indices), self.batch_size): + batch = indices[start : start + self.batch_size] + if len(batch) == self.batch_size or not self.drop_last: + all_batches.append(batch) + + if self.shuffle: + rng.shuffle(all_batches) + + yield from all_batches + + def __len__(self): + """Calculate total number of batches across all groups.""" + total_batches = 0 + for sampler in self._samplers: + total_batches += len(sampler) + + return total_batches + + +class DistributedEquivalenceGroupBatchSampler(Sampler[Batch]): + """Distributed version of EquivalenceGroupBatchSampler for multi-GPU training. + + Uses composition to delegate batching logic to EquivalenceGroupBatchSampler, + handling only the distribution and epoch-based shuffling. + + Ensures uniform batch counts across all ranks to prevent DDP hangs at + collective sync points. Each equivalence group is chunked into logical + DDP steps of num_replicas batches, so ranks process the same group at the + same step. For an incomplete per-group step: + - drop_last=True: drops that incomplete group step + - drop_last=False: pads it by duplicating batches from the same group + + > Note: Compared to the non-distributed sampler: this one won't shuffle + > _within_ batches, only between batches, when `shuffle=True`. + + Args: + datasets: List of TorchTrainDataset instances to group + group_key: Callable that extracts grouping key from a dataset + batch_size: Number of samples per batch + num_replicas: Number of distributed workers (world size) + rank: Index of current worker (0 to num_replicas-1) + shuffle: Whether to shuffle batches + drop_last: Whether to drop incomplete batches within each group, and + whether to trim (vs pad) when distributing batches across ranks + seed: Random seed for shuffling (default: 0) + """ + + def __init__( + self, + datasets: list["TorchTrainDataset"], + group_key: Callable[["TorchTrainDataset"], Hashable], + batch_size: int, + num_replicas: int, + rank: int, + shuffle: bool = True, + drop_last: bool = False, + seed: int = 0, + ): + super().__init__() + if num_replicas <= 0: + raise ValueError(f"num_replicas must be positive, got {num_replicas}.") + if rank >= num_replicas or rank < 0: + raise ValueError( + f"Invalid rank {rank}, must be in range [0, {num_replicas})" + ) + + self.num_replicas = num_replicas + self.rank = rank + self.shuffle = shuffle + self.drop_last = drop_last + self.seed = seed + self.epoch = 0 + + # Delegate batching logic to inner sampler (without shuffle for determinism) + self._inner = EquivalenceGroupBatchSampler.from_datasets( + datasets=datasets, + group_key=group_key, + batch_size=batch_size, + shuffle=False, # We handle shuffling with seeded RNG + drop_last=drop_last, + seed=seed, + ) + + def set_epoch(self, epoch: int) -> None: + """Set the epoch for deterministic shuffling across workers.""" + self.epoch = epoch + + def __iter__(self): + # Every rank constructs the same shuffled global chunk order because + # seed, epoch, and input groups are identical across ranks. + rng = random.Random(self.seed + self.epoch) + + # chunks: list[DdpStepChunk] + # A global logical-step schedule. All ranks build this same list, then + # each rank selects its own slot from every chunk below. + chunks: list[DdpStepChunk] = [] + for sampler in self._inner._samplers: + # sampler: BatchSampler for one equivalence group/resolution. + # + # group_chunks: list[DdpStepChunk] + # Each element is one candidate DDP step for this group. The last + # one may be short before the drop/pad handling below. + group_chunks: list[DdpStepChunk] = list( + itertools.batched(sampler, self.num_replicas) + ) + if group_chunks and len(group_chunks[-1]) < self.num_replicas: + if self.drop_last: + # Drop the short step for this group. This avoids a final + # step where only some ranks have same-group work. + group_chunks.pop() + else: + # Pad the incomplete chunk with duplicates from the same + # group so every chunk is homogeneous and full-sized. + # + # partial_chunk: list[Batch] + # The short final DDP step for this group. + partial_chunk: list[Batch] = list(group_chunks[-1]) + # Draw padding from all of this group's batches (not just + # the tail) so padding doesn't systematically over-weight + # the same examples every epoch. + # + # all_group_batches: list[Batch] + # Padding source restricted to this group, so no padded + # DDP step can mix resolutions. + all_group_batches = list( + itertools.chain.from_iterable(group_chunks) + ) + if self.shuffle: + rng.shuffle(all_group_batches) + padding_pool = itertools.cycle(all_group_batches) + while len(partial_chunk) < self.num_replicas: + partial_chunk.append(next(padding_pool)) + # Replace the short tuple[Batch, ...] with a full + # DdpStepChunk. + group_chunks[-1] = tuple(partial_chunk) + if self.shuffle: + # Shuffle this group's DDP steps. Do not shuffle the batches + # within a step, because position in the tuple maps to rank. + rng.shuffle(group_chunks) + chunks.extend(group_chunks) + + if self.shuffle: + # Shuffle DDP steps across groups while keeping each step internally + # homogeneous. + rng.shuffle(chunks) + + # Per-group chunking above already ensures len is a multiple of num_replicas, + # so every consecutive num_replicas batches belong to the same group. + # + # all_batches: list[Batch] + # Flattened global schedule: + # [step0_rank0, step0_rank1, ..., step1_rank0, step1_rank1, ...] + all_batches: list[Batch] = list(itertools.chain.from_iterable(chunks)) + assert len(all_batches) % self.num_replicas == 0 + + # Each worker takes every num_replicas'th batch starting at rank + for i in range(self.rank, len(all_batches), self.num_replicas): + yield all_batches[i] + + def __len__(self): + """Number of batches for this worker (same for all ranks).""" + n = self.num_replicas + total = 0 + for sampler in self._inner._samplers: + group_len = len(sampler) + if self.drop_last: + total += (group_len // n) * n + else: + total += math.ceil(group_len / n) * n + if self.drop_last: + return total // n + else: + return math.ceil(total / n) diff --git a/samudra/utils/schedule.py b/samudra/utils/schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..79bad426a0858316e2b4b9b076521460a70bc787 --- /dev/null +++ b/samudra/utils/schedule.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Literal + +import torch +from pydantic import BaseModel + + +class CosineSchedulerConfig(BaseModel): + """Cosine scheduler; see pytorch CosineAnnealingLR.""" + + type: Literal["cosine"] = "cosine" + target_epochs: int | None = None + + def build( + self, optimizer: torch.optim.Optimizer, epochs: int + ) -> torch.optim.lr_scheduler.LRScheduler: + max_epochs = self.target_epochs if self.target_epochs is not None else epochs + return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs) + + +class CosineWithTailSchedulerConfig(BaseModel): + """Cosine scheduler which goes to tail_lr for the last tail_epochs.""" + + type: Literal["cosine_with_tail"] = "cosine_with_tail" + target_epochs: int | None = None + + tail_lr: float + tail_epochs: int = 10 + + def build( + self, optimizer: torch.optim.Optimizer, epochs: int + ) -> torch.optim.lr_scheduler.LRScheduler: + max_epochs = self.target_epochs if self.target_epochs is not None else epochs + cosine = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=max_epochs - self.tail_epochs, eta_min=self.tail_lr + ) + tail = torch.optim.lr_scheduler.ConstantLR( + optimizer, + factor=self.tail_lr / optimizer.param_groups[0]["lr"], + total_iters=self.tail_epochs, + ) + return torch.optim.lr_scheduler.SequentialLR( + optimizer, + schedulers=[cosine, tail], + milestones=[max_epochs - self.tail_epochs], + ) + + +class CosineWithWarmupConfig(BaseModel): + """Cosine scheduler which goes from warmup_lr to the default lr for the first warmup_epochs.""" + + type: Literal["cosine_with_warmup"] = "cosine_with_warmup" + target_epochs: int | None = None + + warmup_lr: float = 1e-6 + warmup_epochs: int = 5 + + def build( + self, optimizer: torch.optim.Optimizer, epochs: int + ) -> torch.optim.lr_scheduler.LRScheduler: + max_epochs = self.target_epochs if self.target_epochs is not None else epochs + assert len(optimizer.param_groups) == 1, ( + "There can only be one parameter group for the optimizer." + ) + warmup = torch.optim.lr_scheduler.LinearLR( + optimizer, + start_factor=self.warmup_lr / optimizer.param_groups[0]["lr"], + end_factor=1.0, # Reaches full LR + total_iters=self.warmup_epochs, + ) + + assert self.warmup_epochs <= epochs, ( + "'warmup_epochs' is too big; it must be smaller than 'epochs'." + ) + cosine = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=max_epochs - self.warmup_epochs + ) + return torch.optim.lr_scheduler.SequentialLR( + optimizer, + schedulers=[warmup, cosine], + milestones=[self.warmup_epochs], + ) + + +SchedulerConfig = ( + CosineSchedulerConfig | CosineWithTailSchedulerConfig | CosineWithWarmupConfig +) diff --git a/samudra/utils/train.py b/samudra/utils/train.py new file mode 100644 index 0000000000000000000000000000000000000000..37c5aff48bdac7f7161e7ec566aa8f8b3d68c5ac --- /dev/null +++ b/samudra/utils/train.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence +from itertools import tee +from pathlib import Path + +import torch +from xarray_einstats.einops import rearrange # noqa: F401 + +from samudra.datasets import InferenceDataset, RawTrainData +from samudra.utils.data import LoadStats + + +def pairwise(iterable): + # pairwise('ABCDEFG') --> AB BC CD DE EF FG + a, b = tee(iterable) + next(b, None) + return zip(a, b) + + +def collate_raw_train_data(data: Sequence[RawTrainData]) -> RawTrainData: + batched_data = RawTrainData(data[0].dataset_id) + assert all(d.dataset_id == batched_data.dataset_id for d in data), ( + "we don't support heterogenous batches yet" + ) + + steps = len(data[0].raw_data) + for step in range(steps): + input_ = torch.stack([d.raw_data[step][0] for d in data]) + boundary = torch.stack([d.raw_data[step][1] for d in data]) + label = torch.stack([d.raw_data[step][2] for d in data]) + batched_data.insert(input_, boundary, label) + + stats = LoadStats.accumulated( + [d.load_stats for d in data if d.load_stats is not None] + ) + batched_data.load_stats = stats + + return batched_data + + +def collate_inference_data( + data: Sequence[InferenceDataset], +) -> tuple[InferenceDataset, int]: + # TODO: There is probably a better way to do inference batching + assert len(data) == 1, "Inference batch size must be 1" + return data[0][0], data[0][1] + + +class CheckpointPaths: + def __init__(self, checkpoint_dir: Path): + self.checkpoint_dir = checkpoint_dir + + @property + def latest_checkpoint_path(self) -> Path: + return self.checkpoint_dir / "ckpt.pt" + + def latest_checkpoint_path_with_epoch(self, epoch: int) -> Path: + return self.checkpoint_dir / f"ckpt_{epoch}.pt" + + @property + def best_inference_checkpoint_path(self) -> Path: + return self.checkpoint_dir / "best_inference_ckpt.pt" + + @property + def ema_checkpoint_path(self) -> Path: + return self.checkpoint_dir / "ema_ckpt.pt" + + @property + def best_validation_checkpoint_path(self) -> Path: + return self.checkpoint_dir / "best_validation_ckpt.pt" diff --git a/samudra/utils/train_progress.py b/samudra/utils/train_progress.py new file mode 100644 index 0000000000000000000000000000000000000000..1b16e16b6fa63736203153009366aefdd175cb40 --- /dev/null +++ b/samudra/utils/train_progress.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import dataclasses +import time +from collections.abc import Iterator +from typing import Any + +import torch + +from samudra.datasets import TrainData + + +@dataclasses.dataclass +class TrainBatchProgress: + """Comparable progress and timing for one global train-loader batch. + + Counts are global across distributed workers so the resulting ``progress/*`` + metrics can be used as W&B x-axes when comparing runs with different batch + sizes, GPU counts, gradient accumulation, or grid resolutions. + + A sample window is one dataset item from the train loader. A model example + is one supervised autoregressive training step within that window, so a + multi-step rollout contributes multiple model examples per sample window. + Output grid cells are model examples multiplied by output-grid latitude and + longitude cells. Target values are output grid cells multiplied by output + channels. + + Input grid dimensions describe the spatial grid of the tensors presented to + the model at each autoregressive step. They are not a count of input time + levels or rollout steps. Output grid dimensions describe the spatial grid of + each predicted/target field. + + ``TrainProgress`` accumulates the per-batch counts, optimizer update count, + and summed batch wall time across workers. CUDA training synchronizes before + reading batch wall time so ``progress/gpu_seconds`` and throughput metrics + include queued GPU work. + """ + + sample_windows: int + model_examples: int + output_grid_cells: int + target_values: int + input_grid_lat: int + input_grid_lon: int + output_grid_lat: int + output_grid_lon: int + optimizer_stepped: bool = False + batch_seconds: float = 0.0 + + @classmethod + def from_train_data(cls, data: TrainData, world_size: int) -> "TrainBatchProgress": + label = data.get_label(0) + ( + local_batch_size, + output_channels, + output_grid_lat, + output_grid_lon, + ) = label.shape + input_grid_lat = data.ctx.input_resolution_cpu[0].shape[0] + input_grid_lon = data.ctx.input_resolution_cpu[1].shape[0] + + sample_windows = local_batch_size * world_size + model_examples = sample_windows * len(data) + output_grid_cells = model_examples * output_grid_lat * output_grid_lon + target_values = output_grid_cells * output_channels + + return cls( + sample_windows=sample_windows, + model_examples=model_examples, + output_grid_cells=output_grid_cells, + target_values=target_values, + input_grid_lat=input_grid_lat, + input_grid_lon=input_grid_lon, + output_grid_lat=output_grid_lat, + output_grid_lon=output_grid_lon, + ) + + def to_metrics(self) -> dict[str, int]: + return { + "progress/batch_sample_windows": self.sample_windows, + "progress/batch_model_examples": self.model_examples, + "progress/batch_output_grid_cells": self.output_grid_cells, + "progress/batch_target_values": self.target_values, + "progress/input_grid_lat": self.input_grid_lat, + "progress/input_grid_lon": self.input_grid_lon, + "progress/output_grid_lat": self.output_grid_lat, + "progress/output_grid_lon": self.output_grid_lon, + } + + def to_throughput_metrics(self) -> dict[str, float]: + """Compute global per-batch throughput metrics when timing is available.""" + if self.batch_seconds <= 0: + return {} + return { + "throughput/model_examples_per_second": ( + self.model_examples / self.batch_seconds + ), + "throughput/output_grid_cells_per_second": ( + self.output_grid_cells / self.batch_seconds + ), + } + + +@dataclasses.dataclass +class TrainProgress: + """Cumulative training progress counters saved in checkpoints.""" + + sample_windows_seen: int = 0 + model_examples_seen: int = 0 + output_grid_cells_seen: int = 0 + target_values_seen: int = 0 + optimizer_steps: int = 0 + gpu_seconds: float = 0.0 + + @contextlib.contextmanager + def batch( + self, data: TrainData, *, world_size: int, device: torch.device + ) -> Iterator[TrainBatchProgress]: + batch = TrainBatchProgress.from_train_data(data, world_size) + _synchronize_cuda_if_needed(device) + batch_start_time = time.perf_counter() + + yield batch + + _synchronize_cuda_if_needed(device) + batch.batch_seconds = time.perf_counter() - batch_start_time + self.update(batch, world_size=world_size) + + def update(self, batch: TrainBatchProgress, *, world_size: int) -> None: + self.sample_windows_seen += batch.sample_windows + self.model_examples_seen += batch.model_examples + self.output_grid_cells_seen += batch.output_grid_cells + self.target_values_seen += batch.target_values + if batch.optimizer_stepped: + self.optimizer_steps += 1 + self.gpu_seconds += batch.batch_seconds * world_size + + def to_metrics(self) -> dict[str, int | float]: + return { + "progress/sample_windows_seen": self.sample_windows_seen, + "progress/model_examples_seen": self.model_examples_seen, + "progress/output_grid_cells_seen": self.output_grid_cells_seen, + "progress/target_values_seen": self.target_values_seen, + "progress/optimizer_steps": self.optimizer_steps, + "progress/gpu_seconds": self.gpu_seconds, + } + + def state_dict(self) -> dict[str, int | float]: + return dataclasses.asdict(self) + + @classmethod + def from_state_dict(cls, state: dict[str, Any] | None) -> "TrainProgress": + if state is None: + return cls() + fields = {field.name for field in dataclasses.fields(cls)} + return cls(**{key: value for key, value in state.items() if key in fields}) + + +def _synchronize_cuda_if_needed(device: torch.device) -> None: + """Synchronize CUDA kernels before reading wall-clock batch timing.""" + if device.type == "cuda": + torch.cuda.synchronize(device) diff --git a/samudra/utils/wandb.py b/samudra/utils/wandb.py new file mode 100644 index 0000000000000000000000000000000000000000..12a4b5fdf68e43b015634e72117614ce014f2d45 --- /dev/null +++ b/samudra/utils/wandb.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import logging +import os +from collections.abc import Mapping +from typing import TYPE_CHECKING + +logger = logging.getLogger(__name__) + +import matplotlib.pyplot as plt +import numpy as np +import torch +import wandb +from wandb.data_types import WBValue + +from samudra.utils.multiton import Multiton + +# Metrics supported by wandb -- probably there are more possible types too +Metrics = Mapping[str, float | torch.Tensor | WBValue] + +# Same as above but mutable when you're building something up +MetricsDict = dict[str, float | torch.Tensor | WBValue] + +PROVENANCE_CONFIG_KEY = "provenance" + +if TYPE_CHECKING: + from samudra.config import AnyTopLevelConfig + from samudra.utils.data import DataContainer + + +class WandBLogger(Multiton): + def _initialize(self): + self._enabled = False + self._initialized = False + self.run = None + + @property + def enabled(self): + return self._enabled + + def _make_config(self, cfg: "AnyTopLevelConfig", data_container: "DataContainer"): + config = { + f"data_{i}/attrs": src.data.attrs + for i, src in enumerate(data_container.train_sources) + } + config.update(config=cfg.model_dump()) + provenance_env = { + "code_commit": "SAMUDRA_CODE_COMMIT", + "code_repo_url": "SAMUDRA_CODE_REPO_URL", + "code_layer_sha256": "SAMUDRA_CODE_LAYER_SHA256", + "container_git_commit": "SAMUDRA_CONTAINER_GIT_COMMIT", + "container_git_remote_url": "SAMUDRA_CONTAINER_GIT_REMOTE_URL", + "container_image_ref": "SAMUDRA_CONTAINER_IMAGE_REF", + "container_sif_path": "SAMUDRA_CONTAINER_SIF_PATH", + } + provenance = { + name: value + for name, env_var in provenance_env.items() + if (value := os.environ.get(env_var)) + } + if provenance: + config[PROVENANCE_CONFIG_KEY] = provenance + return config + + def setup_run( + self, + checkpoint_path: str | None, + cfg: "AnyTopLevelConfig", + data_container: "DataContainer", + finetune: bool = False, + ): + """Set up a wandb run, either resuming from checkpoint or creating new run. + + Args: + checkpoint_path: Path to checkpoint file, if resuming + cfg: Configuration object + data_container: Data container to log attributes of + finetune: Whether this is a finetuning run + + Returns: + tuple: (wandb_id, wandb_name) + """ + if not checkpoint_path: + return self._init_new_run(cfg, data_container) + + if finetune: + return self._init_new_run(cfg, data_container) + + if not self._enabled: + return None, None + + # Load only on the rank that initializes W&B, and force CPU placement. + # Checkpoints can contain CUDA tensors saved from rank 0; loading them + # here without map_location makes every rank reserve memory on GPU 0 + # before the real per-rank checkpoint load happens. + checkpoint = torch.load(checkpoint_path, map_location="cpu") + wandb_id = checkpoint.get("wandb_id") + wandb_name = checkpoint.get("wandb_name") + del checkpoint + + try: + self.init( + config=self._make_config(cfg, data_container), + name=wandb_name, + dir=cfg.experiment.output_dir, + resume="must", + id=wandb_id, + **cfg.experiment.wandb.model_dump(), + ) + except Exception: + # If resume fails, start new run + self.init( + config=self._make_config(cfg, data_container), + name=wandb_name, + dir=cfg.experiment.output_dir, + **cfg.experiment.wandb.model_dump(), + ) + + return wandb_id, wandb_name + + def _init_new_run(self, cfg: "AnyTopLevelConfig", data_container: "DataContainer"): + """Initialize a new wandb run. + + Args: + cfg: Configuration object + data_container: Data container to log attributes of + Returns: + tuple: (None, generated_name) for new run + """ + wandb_name = cfg.experiment.name + if self._enabled: + self.init( + config=self._make_config(cfg, data_container), + name=wandb_name, + dir=cfg.experiment.output_dir, + **cfg.experiment.wandb.model_dump(), + ) + + wandb_id = self.run.id if self.run else None + else: + wandb_id = None + + return wandb_id, wandb_name + + def finish(self): + """Finish the wandb run.""" + if self._enabled: + wandb.finish() + + def configure(self, enabled: bool, is_main_process: bool): + """Configure whether wandb should be enabled.""" + assert self._enabled is False, "WandB is already initialized" + self._enabled = enabled and is_main_process + + def init(self, **kwargs): + """Initialize wandb run.""" + if self._enabled and not self._initialized: + try: + self.run = wandb.init(**kwargs) + self._initialized = True + except Exception as e: + logger.error(f"Failed to initialize wandb: {e}") + self._enabled = False + + def watch(self, model, **kwargs): + """Watch model parameters and gradients.""" + if self._enabled: + wandb.watch(model, **kwargs) + + def log(self, metrics: Metrics, step: int | None, **kwargs): + """Log metrics to wandb.""" + if self._enabled: + # Really, this should take a mapping, not a dict + # (so it is covariant) but it doens't so we convert + wandb.log(dict(metrics), step=step, **kwargs) + + def Image(self, data, *args, **kwargs): + if isinstance(data, np.ndarray): + data = scale_image(data) + return wandb.Image(data, *args, **kwargs) + + def Video(self, *args, **kwargs): + return wandb.Video(*args, **kwargs) + + def Table(self, *args, **kwargs): + return wandb.Table(*args, **kwargs) + + def Histogram(self, *args, **kwargs): + return wandb.Histogram(*args, **kwargs) + + def log_inference_metrics( + self, + loss_value: float, + loss_per_channel: torch.Tensor, + outputs: list, + depth_indices: dict, + var_indices: dict, + depth_set: list, + var_set: list, + step: int, + predictions: torch.Tensor | None = None, + targets: torch.Tensor | None = None, + targets_unnormalized: torch.Tensor | None = None, + model_pred_unnormalized: torch.Tensor | None = None, + ): + """Log validation metrics including depth and variable specific metrics.""" + if not self._enabled: + return + + self.log({"eval/total_eval_loss_per_batch": loss_value}, step=step) + + # Loss per channel + for i, var in enumerate(outputs): + self.log({"eval/per_channel/" + var: loss_per_channel[i].item()}, step=step) + + # Loss per depth + if all(x is not None for x in [depth_indices, depth_set]): + for d in depth_set: + self.log( + { + "eval/depth/depth_" + str(d) + "_loss": torch.mean( + loss_per_channel[depth_indices[d]] + ).item() + }, + step=step, + ) + + # Loss per input variable + if all(x is not None for x in [var_indices, var_set]): + for k in var_set: + self.log( + { + "eval/per_var/" + k + "_loss": torch.mean( + loss_per_channel[var_indices[k]] + ).item() + }, + step=step, + ) + + # Plot predictions vs targets + if ( + predictions is not None + and targets is not None + and targets_unnormalized is not None + and model_pred_unnormalized is not None + ): + for i, var in enumerate(outputs): + fig = plt.figure(figsize=(10, 5)) + plt.plot( + range(targets.shape[0]), + targets_unnormalized[:, :, :, i].mean(dim=(1, 2)), + label="Target", + ) + min_val, max_val = plt.ylim() + plt.plot( + range(predictions.shape[0]), + model_pred_unnormalized[:, :, :, i].mean(dim=(1, 2)), + label="Prediction", + ) + + # Adjust y-limits based on variable type + if "thetao" in var: + plt.ylim(min_val - 0.25, max_val + 0.25) + elif "so" in var: + plt.ylim(min_val - 0.2, max_val + 0.2) + elif "KE" in var: + plt.ylim(min_val - 0.5, max_val + 0.5) + + plt.title(var) + plt.legend() + self.log({f"eval/plots/{var}": wandb.Image(fig)}, step=step) + plt.close() + + +def scale_image(image_data): + """ + Given an array of scalar data, rescale the data to the range [0, 255]. + """ + data_min = np.nanmin(image_data) + data_max = np.nanmax(image_data) + + image_data = 255 * (image_data - data_min) / (data_max - data_min) + image_data = np.minimum(image_data, 255) + image_data = np.maximum(image_data, 0) + image_data[np.isnan(image_data)] = 0 + + return image_data + + +def get_record_to_wandb(label: str = ""): + wandb = WandBLogger.get_instance() + step = 0 + + def record_logs(logs): + nonlocal step + for j, log in enumerate(logs): + if len(log) > 0: + if label != "": + log = {f"{label}/{k}": v for k, v in log.items()} + wandb.log(log, step=step + j) + step += len(logs) + + return record_logs diff --git a/samudra/utils/writer.py b/samudra/utils/writer.py new file mode 100644 index 0000000000000000000000000000000000000000..f2a7dac04f9e42cb523a3bd7de89efcda01b2c22 --- /dev/null +++ b/samudra/utils/writer.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import os + +import numpy as np +import torch +import xarray as xr +from einops import rearrange + +from samudra.constants import TensorMap +from samudra.utils.data import Normalize, stack_levels +from samudra.utils.output import ModelInferenceOutput + + +class ZarrWriter: + """Writes model prediction outputs to Zarr format for downstream analysis.""" + + def __init__( + self, + output_dir: str | os.PathLike, + coords: dict[str, xr.DataArray], + hist: int, + model_path: str | os.PathLike, + time_chunk_size: int, + normalize: Normalize, + tensor_map: TensorMap, + ): + self.pred_path = os.path.join(output_dir, "predictions.zarr") + + if os.path.exists(self.pred_path): + raise FileExistsError( + f"Predictions already exist at {self.pred_path}. Please choose a unique experiment name, output directory, or delete the existing predictions." + ) + + self.hist = hist + self.buffer: torch.Tensor | None = None + self.time_buffer: xr.DataArray | None = None + self.coords = coords + self.model_path = model_path + self.time_chunk_size = time_chunk_size + + self.normalize = normalize + self.tensor_map = tensor_map + + def record_batch(self, IO: ModelInferenceOutput): + pred_tensor = IO.prediction + pred_time = IO.time + pred_tensor = rearrange( + pred_tensor, "n (hi c) h w -> (n hi) c h w", hi=self.hist + 1 + ) + pred_tensor = self.normalize.unnormalize_tensor_prognostic( + pred_tensor, fill_value=0.0 + ) + if self.buffer is None: + self.buffer = pred_tensor + else: + self.buffer = torch.cat([self.buffer, pred_tensor], dim=0) + + if self.time_buffer is None: + self.time_buffer = pred_time + else: + self.time_buffer = xr.concat([self.time_buffer, pred_time], dim="time") + + def write(self): + # Write to zarr + if self.buffer is None: + raise ValueError("No tensor to write") + + if self.time_buffer is None: + raise ValueError("No time buffer to write") + + buffer = self.buffer.cpu().numpy() # (time, channel, y, x) + + # Lay the flat per-level channels out as a Dataset, then reassemble into + # depth-stacked variables with `stack_levels` -- the same helper used to + # put ground-truth inputs into analysis-ready form, so predictions match. + coords = self._output_coords() + coords["time"] = self.time_buffer + per_level = xr.Dataset( + { + name: (["time", "y", "x"], buffer[:, channel, :, :]) + for channel, name in enumerate(self.tensor_map.prognostic_var_names) + }, + coords=coords, + ) + ds = stack_levels(per_level, self.tensor_map.dataset_spec) + ds = ds.transpose("time", "lev", "y", "x", ...) + ds.attrs["model_path"] = str(self.model_path) + ds = ds.chunk({"time": self.time_chunk_size}) + if os.path.exists(self.pred_path): + ds.to_zarr(self.pred_path, mode="a", append_dim="time") + else: + ds.to_zarr( + self.pred_path, + mode="w", + encoding={var: {"compressor": None} for var in ds.data_vars}, + ) + + # Reset + self.buffer = None + self.time_buffer = None + + def _output_coords(self) -> dict: + """Build output coordinates that match the ground-truth OM4 dataset layout. + + "Ground-truth layout" means the on-disk shape of the preprocessed OM4 data + (the same data the model is scored against), so that predictions and truth + can be compared without any reshaping. Concretely that layout uses ``y``/``x`` + horizontal dims with 2D ``lat``/``lon`` coordinates on top of them (not 1D + ``lat``/``lon`` dims), carries the depth axis ``lev``, and keeps the grid + metadata needed for physical analysis (``areacello``, ``dz``, + ``ocean_fraction``, cell bounds ``lat_b``/``lon_b``). + """ + src = self.coords + y_vals = np.asarray(src["lat"].values) # 1-D latitude axis + x_vals = np.asarray(src["lon"].values) # 1-D longitude axis + ny, nx = y_vals.size, x_vals.size + + spec = self.tensor_map.dataset_spec + n_levels = spec.num_prognostic_depth_levels + + # 2-D lat/lon on the (y, x) grid, matching the ground-truth layout. Prefer + # the real coordinates the source preserved (`lat_2d`/`lon_2d`), which are + # correct on any grid. When they are absent we may only reconstruct them by + # broadcasting the 1-D axes on a "gaussian" (rectilinear) grid, where that is + # exact. On a curvilinear grid (e.g. tripolar) broadcasting fabricates wrong + # geometry, so we refuse rather than silently emit bad coordinates. + if "lat_2d" in src and "lon_2d" in src: + lat2d = np.asarray(src["lat_2d"].values) + lon2d = np.asarray(src["lon_2d"].values) + elif spec.grid_type == "gaussian": + lat2d = np.broadcast_to(y_vals[:, None], (ny, nx)).copy() + lon2d = np.broadcast_to(x_vals[None, :], (ny, nx)).copy() + else: + raise ValueError( + f"Cannot build 2-D lat/lon for grid_type={spec.grid_type!r}: the " + "source coords carry no real 'lat_2d'/'lon_2d', and broadcasting the " + "1-D axes is only valid on a 'gaussian' (rectilinear) grid. Preserve " + "the true 2-D coordinates through preprocessing for curvilinear grids." + ) + + coords: dict[str, tuple] = { + "y": ("y", y_vals), + "x": ("x", x_vals), + "lat": (("y", "x"), lat2d), + "lon": (("y", "x"), lon2d), + "lev": ("lev", np.array(spec.depth_levels[:n_levels])), + } + + rename = {"lat": "y", "lon": "x"} + for name, da in src.items(): + if name in ("time", "lat", "lon", "lat_2d", "lon_2d", "lev", "y", "x"): + continue + # Keep depth-resolved metadata (dz, ocean_fraction) consistent with the + # emitted levels. A shallow model spec (e.g. thermo_dynamic_5) predicts + # fewer levels than the source's full-depth grid coords carry, so slice + # those to the emitted `lev` count instead of conflicting with it. + if "lev" in da.dims: + da = da.isel(lev=slice(0, n_levels)) + dims = tuple(rename.get(str(d), str(d)) for d in da.dims) + coords[name] = (dims, np.asarray(da.values)) + + return coords diff --git a/samudra/viz/__init__.py b/samudra/viz/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f21bc7e41dd88a4cd1b5bec17093ed188b9c0ca5 --- /dev/null +++ b/samudra/viz/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from samudra.viz.config import VizConfig +from samudra.viz.core import Viz, VizRun + +__all__ = ["Viz", "VizConfig", "VizRun"] diff --git a/samudra/viz/__main__.py b/samudra/viz/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..ae4494465d29a911c6e4f41f2211a4196befaedb --- /dev/null +++ b/samudra/viz/__main__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +from samudra.viz.config import VizConfig, main + +main(VizConfig.from_yaml_and_cli()) diff --git a/samudra/viz/config.py b/samudra/viz/config.py new file mode 100644 index 0000000000000000000000000000000000000000..7253d2227d61c0cf80c7daa8027db88e49eb5dde --- /dev/null +++ b/samudra/viz/config.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import functools +import logging +import time +from functools import cached_property +from pathlib import Path +from typing import Annotated + +from pydantic import BaseModel, BeforeValidator, Field, WithJsonSchema + +from samudra.config import DataConfig, ObsMetricsConfig, Om4TimeConfig +from samudra.config_base import TopLevelConfig +from samudra.utils.location import LocalLocation, Location, ResolvedLocation +from samudra.utils.logging import handle_logging +from samudra.viz.core import Viz, VizRun + + +@functools.cache +def _all_steps() -> set[str]: + return set(_ordered_steps()) + + +@functools.cache +def _ordered_steps() -> list[str]: + return list( + name.removeprefix("step_") for name in Viz.__dict__ if name.startswith("step_") + ) + + +def _check_step(v: str) -> str: + if v not in _all_steps(): + raise ValueError( + f"Invalid step: '{v}', expected one of: {', '.join(_ordered_steps())}" + ) + return v + + +VizStep = Annotated[ + str, + BeforeValidator(_check_step), + WithJsonSchema({"type": "string", "enum": _ordered_steps()}), +] + + +class VizRunConfig(BaseModel): + name: str + location: Location + variables: list[str] = Field( + default_factory=lambda: ["thetao", "so", "uo", "vo", "tos", "zos"] + ) + + def build(self, data_root: ResolvedLocation) -> VizRun: + return VizRun( + name=self.name, + data=data_root.resolve(self.location).open(chunks={}), + variables=self.variables, + ) + + +class VizConfig(TopLevelConfig): + base_output_dir: Path + name: str + dataset_name: str + runs: list[VizRunConfig] + data_root: Location | None = None + data: DataConfig | None = Field( + default=None, + description="Optional data source, matching train/eval. When set, ground " + "truth defaults to its primary source's data_location, so " + "`--data @data/om4_demo.yaml` works the same across train, eval, and viz. " + "An explicit `groundtruth_location` overrides it.", + ) + groundtruth_location: Location | None = Field( + default=None, + description="Ground-truth rollout data. Optional when `data` is set (then it " + "defaults to that source's data_location); otherwise required.", + ) + basins_location: Location + # TODO(jder): we could extract this from the run data? + groundtruth_time_range: Om4TimeConfig = Field( + description="Dates from the rollout (not same as eval *input* dates; these are the dates the output is produced for during eval)" + ) + observations: ObsMetricsConfig | None = Field( + default=None, + description=( + "Observation products to draw figures against. The same block " + "`samudra.eval` takes, so one file configures both and the two " + "agree by construction. Omit it to skip the observation steps." + ), + ) + steps: list[VizStep] | None = Field( + default=None, + description=f"Which steps to run; leave empty to run all steps. Possible values are: {', '.join(_ordered_steps())}", + ) + not_steps: list[VizStep] = Field( + default_factory=lambda: [], + description="Steps to *not* run, takes precedence over `steps` (see that key for possible steps).", + ) + debug: bool = Field(default=False, description="") + + @cached_property + def output_path(self) -> Path: + return Path(self.base_output_dir) / self.name + + def _groundtruth_location(self) -> Location: + """Ground-truth location: explicit if given, else the primary data source.""" + if self.groundtruth_location is not None: + return self.groundtruth_location + if self.data is not None: + return self.data.sources[0].data_location + raise ValueError( + "viz needs ground-truth data: set `groundtruth_location`, or pass a " + "data source (e.g. --data @data/om4_demo.yaml)." + ) + + def build(self, default_root: ResolvedLocation) -> Viz: + if self.data_root is None: + data_root = default_root + else: + data_root = default_root.resolve(self.data_root) + + groundtruth_rollout = data_root.resolve(self._groundtruth_location()).open( + chunks={} + ) + + return Viz( + # TODO(jder): change to Path + str(self.output_path), + self.dataset_name, + [run.build(data_root) for run in self.runs], + data_root.resolve(self.basins_location).open(), + groundtruth_rollout, + self.groundtruth_time_range.time_slice, + observations=self.observations, + data_root=data_root, + ) + + +logger = logging.getLogger(__name__) + + +def main(cfg: VizConfig): + cfg.output_path.mkdir(parents=True, exist_ok=True) + handle_logging(cfg.debug, cfg.output_path) + cfg.save_yaml(cfg.output_path / "config.yaml") + + logger.info(f"Writing results to {cfg.output_path}") + + viz = cfg.build(LocalLocation(path=Path.cwd())) + + steps = [s for s in cfg.steps or _ordered_steps() if s not in cfg.not_steps] + logger.info(f"Running steps: {', '.join(steps)}") + + # TODO(jder): could use a ProcessPoolExecutor here, but steps currently + # are not exactly independent (some write to pred_dict which others read, + # there's some appending to a metrics file.) + for step in steps: + _run_step(viz, step) + + +def _run_step(viz: Viz, step: VizStep): + logger.info(f"Running step {step}") + start = time.perf_counter() + try: + getattr(viz, f"step_{step}")() + end = time.perf_counter() + logger.info(f"Step {step} took {end - start:.2f} seconds") + except Exception as e: + logger.error(f"Step {step} failed: {e}") + raise diff --git a/samudra/viz/core.py b/samudra/viz/core.py new file mode 100644 index 0000000000000000000000000000000000000000..ead349cfa42ce7e8aedf2f625379c87aebec1742 --- /dev/null +++ b/samudra/viz/core.py @@ -0,0 +1,4100 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import gc +import glob +import logging +import os +import re +import sys +import warnings +from collections.abc import Iterable +from copy import deepcopy +from pathlib import Path +from subprocess import PIPE, STDOUT, Popen +from typing import TYPE_CHECKING, Any + +import cartopy.crs as ccrs # type: ignore +import cartopy.feature as cfeature # type: ignore +import cmocean as cm # type: ignore +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr +from dask.array.core import Array as DaskArray +from dask.base import compute +from dask.delayed import delayed +from dask.diagnostics.progress import ProgressBar +from matplotlib.ticker import FixedLocator, MaxNLocator, ScalarFormatter +from tqdm.auto import tqdm + +from samudra.constants import DatasetSpec, build_om4_spec +from samudra.metrics.run import score_rollouts +from samudra.utils.data import ( + spherical_area, + spherical_area_weights, + with_level_index_vars, +) +from samudra.utils.location import ResolvedLocation +from samudra.viz import observations as obs_figures +from samudra.viz.norms import percentile_norm, symmetric_percentile_norm + +if TYPE_CHECKING: + from samudra.config import ObsMetricsConfig + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class VizRun: + name: str + data: xr.Dataset + variables: list[str] + + +class Viz: + """Generates maps, time series, and probability density plots from evaluation outputs.""" + + def __init__( + self, + output_path: str, + dataset_name: str, + runs: list["VizRun"], + basins: xr.Dataset, + groundtruth_rollout: xr.Dataset, + time_range: slice, + observations: "ObsMetricsConfig | None" = None, + data_root: ResolvedLocation | None = None, + ): + pred_dict: dict[str, dict[str, Any]] = {} + for run in runs: + pred_dict[run.name] = { + "name": run.name, + "data": run.data, + "ls": run.variables, + } + + key1 = runs[0].name + # TODO: Support non-OM4 dataset specs in visualization. + self.dataset_spec = build_om4_spec() + levels = len(self.dataset_spec.depth_levels) + + groundtruth_rollout = groundtruth_rollout.sel(time=time_range) + + if "y" in groundtruth_rollout.coords: + groundtruth_rollout = groundtruth_rollout.drop_vars( + ["lat", "lon"], errors="ignore" + ) + groundtruth_rollout = groundtruth_rollout.rename({"y": "lat", "x": "lon"}) + + groundtruth_rollout = groundtruth_rollout.assign( + areacello=(["lat", "lon"], spherical_area_weights(groundtruth_rollout)) + ) + + # Compute real grid cell areas for physical calculations + groundtruth_rollout["areacello_spherical"] = ( + ["lat", "lon"], + spherical_area(groundtruth_rollout), + ) + + # This function processes the ds_groundtruth and predictions for plotting + # The predictions are loaded into pred_dict + data, pred_dict = process_data( + groundtruth_rollout, pred_dict, dataset_spec=self.dataset_spec + ) + + last_index = len(data.time) - 1 + self.time_indices = [0, last_index // 2, last_index] + + var_list = { + "vo": r"$v$ $( m/s )$", + "uo": r"$u$ $( m/s )$", + "thetao": r"$T$ $( ^\circ C )$", + "tos": r"$T$ $( ^\circ C )$", + "so": r"$so$ $( psu )$", + "zos": r"$zos$ $( m )$", + "KE": r"$KE$ $( J/m^2 )$", + "OHC": r"$OHC$ $Anomaly$ $( ZJ )$", + } + + # Create folder paths + self.timeseries_path = os.path.join(output_path, f"Timeseries") + if not os.path.isdir(self.timeseries_path): + os.makedirs(self.timeseries_path) + + self.ohc_path = os.path.join(output_path, f"OHC") + if not os.path.isdir(self.ohc_path): + os.makedirs(self.ohc_path) + + self.temp_path = os.path.join(output_path, f"Temperature") + if not os.path.isdir(self.temp_path): + os.makedirs(self.temp_path) + + self.salinity_path = os.path.join(output_path, f"Salinity") + if not os.path.isdir(self.salinity_path): + os.makedirs(self.salinity_path) + + self.pdfs_path = os.path.join(output_path, f"PDFs") + if not os.path.isdir(self.pdfs_path): + os.makedirs(self.pdfs_path) + + self.enso_path = os.path.join(output_path, f"ENSO") + if not os.path.isdir(self.enso_path): + os.makedirs(self.enso_path) + + self.metrics_path = os.path.join(output_path, f"Metrics") + if not os.path.isdir(self.metrics_path): + os.makedirs(self.metrics_path) + + self.movie_path = os.path.join(output_path, f"Movies") + if not os.path.isdir(self.movie_path): + os.makedirs(self.movie_path) + + clist = ["#ff807a", "#1e8685", "#ffb579", "#63c8ab"] + + atlantic_mask0 = basins["basin_atlantic"] + atlantic_mask = atlantic_mask0.where(atlantic_mask0["lat"] >= -32) + atlantic_mask = process_mask(data, atlantic_mask) + pacific_mask0 = basins["basin_pacific"] + pacific_mask = pacific_mask0.where( + pacific_mask0["lat"] >= -32 + ) # TODO: include this -32 masking in the basin data + pacific_mask = process_mask(data, pacific_mask) + indian_ocean_mask0 = basins["basin_indian"] + indian_ocean_mask = indian_ocean_mask0.where(indian_ocean_mask0["lat"] >= -32) + indian_ocean_mask = process_mask(data, indian_ocean_mask) + southern_ocean_mask0 = basins["basin_southern"] + southern_ocean_mask = process_mask(data, southern_ocean_mask0) + arctic_mask0 = basins["basin_arctic"] + arctic_ocean_mask = process_mask(data, arctic_mask0) + + self.basin_masks = xr.Dataset( + { + "Atlantic": atlantic_mask, + "Pacific": pacific_mask, + "Southern": southern_ocean_mask, + "Indian": indian_ocean_mask, + "Arctic": arctic_ocean_mask, + } + ) + + # Compute profile means + with ProgressBar(): + logger.info("Computing profile for ground truth " + dataset_name) + profile_groundtruth = profile_mean(data).load() + + for k in pred_dict.keys(): + logger.info("Computing profile for prediction " + k) + pred_dict[k]["profile_prediction"] = profile_mean( + pred_dict[k]["ds_prediction"] + ).load() + + self.data: xr.Dataset = data + self.profile_groundtruth: xr.Dataset = profile_groundtruth + self.pred_dict: dict[str, dict[str, Any]] = pred_dict + self.dataset_name: str = dataset_name + self.clist: list[str] = clist + self.var_list: dict[str, str] = var_list + self.levels: int = levels + self.key1: str = key1 + self.output_path: str = output_path + + # The observation steps work from the rollouts as written, not from the + # `pred_dict` the other steps share: they reduce through + # `samudra.metrics`, which needs the same input `samudra.eval` gets if + # the figures are to carry the same numbers. + self.observations = observations + self.obs_data_root = data_root + self.raw_runs: list[VizRun] = list(runs) + self.obs_path = os.path.join(output_path, "Figures_observations") + self._obs_cache: tuple | None = None + + def step_timeseries_plots(self): + ### Plotting timeseries for each variable for each level + for v in ["uo", "vo", "thetao", "so", "zos"]: + print(f"Plotting {v} timeseries...") + if not os.path.isdir(os.path.join(self.timeseries_path, f"{v}_timeseries")): + os.makedirs(os.path.join(self.timeseries_path, f"{v}_timeseries")) + + plt.clf() + plt.rcParams.update({"font.size": 20}) + plt.figure(figsize=(18, 10)) + + if v == "zos": + self.profile_groundtruth[v].plot(label=self.dataset_name, c="k") # type: ignore + for i, k in enumerate(self.pred_dict.keys()): + if v in self.pred_dict[k]["ls"]: + self.pred_dict[k]["profile_prediction"][v].plot( + label=self.pred_dict[k]["name"], c=self.clist[i] + ) + min_val, max_val = plt.ylim() + plt.ylim(min_val - 0.05, max_val + 0.05) + plt.xlabel("Time") + plt.ylabel(self.var_list[v]) + plt.legend() + plt.savefig( + os.path.join(self.timeseries_path, f"{v}_timeseries/0.png"), + bbox_inches="tight", + dpi=600, + ) + plt.close() + else: + for lev in range(self.levels): + plt.clf() + plt.rcParams.update({"font.size": 20}) + plt.figure(figsize=(18, 10)) + self.profile_groundtruth[v].isel(lev=lev).plot( # type: ignore + label=self.dataset_name, c="k" + ) + min_val, max_val = plt.ylim() + for i, k in enumerate(self.pred_dict.keys()): + if v in self.pred_dict[k]["ls"]: + self.pred_dict[k]["profile_prediction"][v].isel( + lev=lev + ).plot(label=self.pred_dict[k]["name"], c=self.clist[i]) + # if v == "thetao": + # plt.ylim(min_val - 0.25, max_val + 0.25) + # elif v == "so": + # plt.ylim(min_val - 0.2, max_val + 0.2) + + plt.xlabel("Time") + plt.ylabel(self.var_list[v]) + plt.legend() + plt.savefig( + os.path.join(self.timeseries_path, f"{v}_timeseries/{lev}.png"), + bbox_inches="tight", + dpi=600, + ) + plt.close() + + def step_short_timeseries_plots(self): + # Short Timeseries plots + shallow_levels = [2.5, 775] + vars = ["thetao"] + num_shallow_levels = len(shallow_levels) + + plt.rcParams.update({"font.size": 14}) + + self.key1 = list(self.pred_dict.keys())[0] + num_plots = 0 + for var in vars: + if "lev" in self.pred_dict[self.key1]["ds_prediction"][var].coords: + num_plots += num_shallow_levels # One plot per level + else: + num_plots += 1 # One plot for scalar variables + + # TODO(jder): Set grid size dynamically based on the number of required plots + cols = 2 + rows = 1 + + fig, axes = plt.subplots(rows, cols, figsize=(16, 3)) + axes = axes.flatten() # Flatten the 2D array of axes for easy access + + plot_idx = 0 # Initialize plot index to track subplot positions + + # Loop over each variable to plot its time series + for v in vars: + if v == "zos": + ax = axes[plot_idx] + + # Ground truth plot + self.profile_groundtruth[v].plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + min_val = self.profile_groundtruth[v].min() + max_val = self.profile_groundtruth[v].max() + + for i, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["profile_prediction"][v].plot( + ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i] + ) + + ax.set_ylim(min_val - 0.05, max_val + 0.05) + ax.set_title(f"{v}") + ax.set_xlabel("Time") + ax.set_ylabel(self.var_list[v]) + + plot_idx += 1 + + else: + for lev in shallow_levels: + ax = axes[plot_idx] + + # Ground truth plot + self.profile_groundtruth[v].sel(lev=lev).plot( # type: ignore + ax=ax, label=self.dataset_name, c="k" + ) + mins, maxs = plt.ylim() + + for i, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["profile_prediction"][v].sel(lev=lev).plot( + ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i] + ) + + # Adjust y-axis limits + if v == "thetao": + if lev > 100: + # ax.set_ylim(mins - 0.02, maxs + 0.02) + pass + # else: + # ax.set_ylim(mins - 0.25, maxs + 0.25) + elif v == "so": + ax.set_ylim(mins - 0.2, maxs + 0.2) + + ax.set_title(f"{lev}m" + r" $\theta_O$", fontsize=14) + ax.set_xlabel("Time") + ax.set_ylabel(r"$\theta_O$ [$\degree C$]") + + plot_idx += 1 + if plot_idx >= rows * cols: + break + + # Adjust layout to avoid overlap and place the legend outside the plot + fig.tight_layout(rect=(0, 0, 0.85, 0.96)) + + # Create a single legend for all plots + handles, labels = ax.get_legend_handles_labels() + fig.legend( + handles, labels, loc="lower center", bbox_to_anchor=(0.455, 0.9), ncols=3 + ) + + # Remove any empty subplots + for idx in range(num_plots, len(axes)): + fig.delaxes(axes[idx]) + + # Save the figure + output_file = os.path.join( + self.timeseries_path, "temperature_timeseries_grid_shallow_both.png" + ) + plt.savefig(output_file, bbox_inches="tight", dpi=600) + + def step_shallow_timeseries_grid_plots(self): + shallow_levels = [2.5, 775, 2400] # Define shallow depth self.levels + + plt.rcParams.update({"font.size": 14}) + + variables = [ + "so", + "uo", + "vo", + "thetao", + ] # List of variables for rows: salinity, zonal velocity, meridional velocity + cols = len( + shallow_levels + ) # Number of columns corresponds to shallow self.levels + rows = len(variables) # One row per variable + + fig, axes = plt.subplots( + rows, cols, figsize=(16, 12), gridspec_kw={"wspace": 0.27, "hspace": 0.2} + ) # Adjust figure size for more rows + axes = axes.reshape( + rows, cols + ) # Reshape axes for easy access by row and column + + # Define labels for each variable + var_labels = { + "so": "S [psu]", + "uo": "uo [m/s]", + "vo": "vo [m/s]", + "thetao": r"$\theta_O$ [$\degree C$]", + } + + # Loop over each variable and plot profiles for each shallow level + for row_idx, var in enumerate(variables): + for col_idx, lev in enumerate(shallow_levels): + ax = axes[row_idx, col_idx] # Access subplot by row and column + + if var == "thetao" and (lev == 2.5 or lev == 775): + fig.delaxes(ax) + continue + + # Ground truth plot for each variable at the specified level + self.profile_groundtruth[var].sel(lev=lev).plot( # type: ignore + ax=ax, label=self.dataset_name, c="k" + ) + mins, maxs = ax.get_ylim() + + for i, k in enumerate(self.pred_dict.keys()): + if var in self.pred_dict[k]["ls"]: + self.pred_dict[k]["profile_prediction"][var].sel(lev=lev).plot( + ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i] + ) + + # Adjust y-axis limits based on the variable + + if var == "thetao": + # if lev > 2000: + # ax.set_ylim(mins - 0.01, maxs + 0.01) + # elif lev > 100: + # ax.set_ylim(mins - 0.01, maxs + 0.01) + # else: + # ax.set_ylim(mins - 0.25, maxs + 0.25) + ax.set_title(f"{lev}m " + r"$\theta_O$", fontsize=14) + elif var == "so": # Salinity + # if lev > 2000: + # ax.set_ylim(mins - 0.001, maxs + 0.001) + # elif lev > 100: + # ax.set_ylim(mins - 0.004, maxs + 0.004) + # else: + # ax.set_ylim(mins - 0.1, maxs + 0.1) + ax.set_title(f"{lev}m $S$", fontsize=14) + handles, labels = ax.get_legend_handles_labels() + if var in self.pred_dict[k]["ls"]: + if var == "uo": # Zonal velocity + # if lev > 2000: + # ax.set_ylim(mins - 0.0003, maxs + 0.0003) + # elif lev > 100: + # ax.set_ylim(mins - 0.0005, maxs + 0.0005) + # else: + # ax.set_ylim(mins - 0.005, maxs + 0.005) + ax.set_title(f"{lev}m $uo$", fontsize=14) + elif var == "vo": # Meridional velocity + # if lev > 2000: + # ax.set_ylim(mins - 0.0002, maxs + 0.0002) + # elif lev > 100: + # ax.set_ylim(mins - 0.0004, maxs + 0.0004) + # else: + # ax.set_ylim(mins - 0.005, maxs + 0.005) + ax.set_title(f"{lev}m $vo$", fontsize=14) + + # ax.ticklabel_format(axis='y', style='sci', scilimits=(0,0)) + ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False)) + + # Set y-axis label only on the leftmost column + if col_idx == 0 or var == "thetao": + ax.set_ylabel(var_labels[var]) # Set y-axis label based on variable + else: + ax.set_ylabel("") + # ax.set_yticklabels([]) # Hide y-axis tick labels for other columns + + # Set x-axis label and ticks only on the bottommost row + if row_idx == rows - 1: + ax.set_xlabel("Time") + else: + ax.set_xlabel("") + ax.set_xticklabels([]) + + # Adjust layout to avoid overlap and place the legend outside the plot + # fig.tight_layout(rect=[0, 0, 0.85, 0.96]) + + # Create a single legend for all plots + fig.legend( + handles, labels, loc="lower center", bbox_to_anchor=(0.5, 0.90), ncols=3 + ) + + # Save the figure with an updated filename + output_file = os.path.join( + self.timeseries_path, "timeseries_grid_shallow_all_vars.png" + ) + plt.savefig(output_file, bbox_inches="tight", dpi=600) + # plt.show() + # plt.close(fig) + + def step_temp_timeseries_shallow_grid_plots(self): + shallow_levels = [2.5, 105, 550] + + num_shallow_levels = len(shallow_levels) + + plt.rcParams.update({"font.size": 12}) + num_plots = 0 + self.key1 = list(self.pred_dict.keys())[0] + for var in ["thetao"]: + if "lev" in self.pred_dict[self.key1]["ds_prediction"][var].coords: + num_plots += num_shallow_levels # One plot per level + else: + num_plots += 1 # One plot for scalar variables + + # Set grid size dynamically based on the number of required plots + cols = 3 # Number of columns + rows = int(np.ceil(num_plots / cols)) + + fig, axes = plt.subplots(rows, cols, figsize=(20, rows * 3)) + axes = axes.flatten() # Flatten the 2D array of axes for easy access + + plot_idx = 0 # Initialize plot index to track subplot positions + + # Loop over each variable to plot its time series + for v in ["thetao"]: + if v == "zos": + # Handle 'zos' separately (no self.levels) + ax = axes[plot_idx] + self.profile_groundtruth[v].plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + + for i, k in enumerate(self.pred_dict.keys()): + if v in self.pred_dict[k]["ls"]: + self.pred_dict[k]["profile_prediction"][v].plot( + ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i] + ) + + # Adjust y-axis limits and formatting + min_val, max_val = ax.get_ylim() + ax.set_ylim(min_val - 0.05, max_val + 0.05) + ax.set_title(f"{v}", fontsize=14) + ax.set_xlabel("Time") + ax.set_ylabel(self.var_list[v]) + + plot_idx += 1 # Move to the next subplot + + else: + # For other variables, loop over each level + for lev in shallow_levels: + ax = axes[plot_idx] + self.profile_groundtruth[v].sel(lev=lev).plot( # type: ignore + ax=ax, label=self.dataset_name, c="k" + ) + + min_val, max_val = ax.get_ylim() + for i, k in enumerate(self.pred_dict.keys()): + if v in self.pred_dict[k]["ls"]: + self.pred_dict[k]["profile_prediction"][v].sel( + lev=lev + ).plot( + ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i] + ) + + # Adjust y-axis limits for specific variables + if v == "thetao": + ax.set_ylim(min_val - 0.25, max_val + 0.25) + elif v == "so": + ax.set_ylim(min_val - 0.2, max_val + 0.2) + + lev_str = ( + self.pred_dict[k]["profile_prediction"][v] + .sel(lev=lev) + .lev.item() + ) + ax.set_title( + (r"$\theta_O$" + f" at {lev_str}m" + r" ($\degree C$)"), + fontsize=14, + ) + ax.set_xlabel("Time") + ax.set_ylabel(self.var_list[v]) + + plot_idx += 1 # Move to the next subplot + + if plot_idx >= rows * cols: + break # Stop if the grid is full + + # Adjust layout to avoid overlap and place the legend outside the plot + fig.tight_layout( + rect=(0, 0, 0.85, 0.96) + ) # Leave space on the right for the legend + + # Create a single legend for all plots + handles, labels = ax.get_legend_handles_labels() + fig.legend( + handles, + labels, + loc="center left", + ncols=3, + bbox_to_anchor=(0.34, -0.03), + fontsize=12, + ) + + # fig.suptitle("Time Series Plots", fontsize=24) + + # Remove any empty subplots + for idx in range(num_plots, len(axes)): + fig.delaxes(axes[idx]) + + # Save the figure + output_file = os.path.join( + self.timeseries_path, "temp_timeseries_grid_shallow_skipped.png" + ) + plt.savefig(output_file, bbox_inches="tight", dpi=600) + + def step_global_thetao_time_series(self): + var = "thetao" + Days_to_Eq = 0 + + plt.rcdefaults() + plt.rcParams.update({"font.size": 12}) + fig, ax = plt.subplots( + 1, 1, figsize=(16, 3), gridspec_kw={"wspace": 0.25, "hspace": 0.5} + ) + + thetao = ( + self.data["thetao"] + .weighted(self.data["areacello"] * self.data["dz"]) + .mean(["x", "y", "lev"]) + ) + thetao = thetao.rename(r"$\theta_O$") + thetao = thetao.assign_attrs(units=r"$\degree C$") + + for i, k in enumerate(self.pred_dict.keys()): + thetao_pred = ( + self.pred_dict[k]["ds_prediction"][var] + .weighted(self.data["areacello"] * self.data["dz"]) + .mean(["x", "y", "lev"]) + ) + thetao_pred.plot(ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i]) + coeffs_ = np.polyfit( + np.arange(thetao_pred[Days_to_Eq:].size), thetao_pred[Days_to_Eq:], 1 + ) + print(f"{self.pred_dict[k]['name']}: {coeffs_[0] * 73}") + + thetao.plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + coeffs_OHC_ground_trend = np.polyfit( + np.arange(thetao[Days_to_Eq:].size), thetao[Days_to_Eq:], 1 + ) + print("OHC: ", coeffs_OHC_ground_trend[0] * 73) + + ax.legend(ncol=3) + ax.set_title("") + # ax.set_ylim([3.230, 3.245]) + plt.savefig( + os.path.join(self.timeseries_path, f"Global_Thetao_Timeseries"), + bbox_inches="tight", + dpi=600, + ) + # plt.show() + + def step_global_salinity_timeseries_plots(self): + var = "so" + Days_to_Eq = 0 + + plt.rcdefaults() + plt.rcParams.update({"font.size": 12}) + fig, ax = plt.subplots( + 1, 1, figsize=(16, 3), gridspec_kw={"wspace": 0.25, "hspace": 0.5} + ) + + salinity = ( + self.data["so"] + .weighted(self.data["areacello"] * self.data["dz"]) + .mean(["x", "y", "lev"]) + ) + salinity = salinity.rename("S") + salinity = salinity.assign_attrs(units="psu") + + for i, k in enumerate(self.pred_dict.keys()): + salinity_pred = ( + self.pred_dict[k]["ds_prediction"][var] + .weighted(self.data["areacello"] * self.data["dz"]) + .mean(["x", "y", "lev"]) + ) + salinity_pred.plot(ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i]) + coeffs_ = np.polyfit( + np.arange(salinity_pred[Days_to_Eq:].size), + salinity_pred[Days_to_Eq:], + 1, + ) + print(f"{self.pred_dict[k]['name']}: {coeffs_[0] * 73}") + + salinity.plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + coeffs_salinity_ground_trend = np.polyfit( + np.arange(salinity[Days_to_Eq:].size), salinity[Days_to_Eq:], 1 + ) + print("OHC: ", coeffs_salinity_ground_trend[0] * 73) + + ax.legend(ncol=3, loc="upper left") + # ax.set_ylim([34.723, 34.728]) + ax.set_ylim((34.72, 34.73)) + ax.set_title("") + ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False)) + plt.savefig( + os.path.join(self.timeseries_path, f"Global_Salinity_Timeseries"), + bbox_inches="tight", + dpi=600, + ) + + def step_ohc_noanomaly_plots(self): + c_p = 3850 # J/(kg C) + rho_0 = 1025 # kg/m^3 + f = open(os.path.join(self.output_path, "compare_info.txt"), "a") + + plt.rcdefaults() + fig, ax = plt.subplots( + 1, 1, figsize=(10, 3), gridspec_kw={"wspace": 0.25, "hspace": 0.5} + ) + plt.rcParams.update({"font.size": 9}) + + OHC = ( + (self.data["thetao"] * c_p * rho_0) + * self.data["areacello"] + * self.data["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + OHC = OHC - OHC.isel(time=0) + OHC = OHC.rename("OHC") + OHC = OHC.assign_attrs(units="ZJ") + + for i, k in enumerate(self.pred_dict.keys()): + OHC_pred = ( + (self.pred_dict[k]["ds_prediction"]["thetao"] * c_p * rho_0) + * self.pred_dict[k]["ds_prediction"]["areacello"] + * self.pred_dict[k]["ds_prediction"]["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + OHC_pred = OHC_pred - OHC_pred.isel(time=0) + OHC_pred = OHC_pred.rename("OHC") + OHC_pred = OHC_pred.assign_attrs(units="ZJ") + OHC_pred.plot(ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i]) + coeffs_OHC_pred_trend = np.polyfit(np.arange(OHC_pred.size), OHC_pred, 1) + (pos,) = ax.plot( + OHC_pred.time.data, + np.arange(OHC_pred.size) * coeffs_OHC_pred_trend[0] + + coeffs_OHC_pred_trend[1], + c=self.clist[i], + ls="--", + ) + # ax[0].annotate(f'{coeffs_OHC_pred_trend[0]:.2e}', + # xy=(pos.get_xdata()[-1], pos.get_ydata()[-1]), + # xytext=(pos.get_xdata()[-2], pos.get_ydata()[-2]), + # fontsize=9, color=self.clist[i]) + f.write( + f"\nOHC {self.pred_dict[k]['name']} Trend Slope : {coeffs_OHC_pred_trend[0]}" + ) + self.pred_dict[k]["OHC_slope"] = coeffs_OHC_pred_trend[0] + + OHC.plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + coeffs_OHC_trend = np.polyfit(np.arange(OHC.size), OHC, 1) + (pos,) = ax.plot( + OHC.time.data, + np.arange(OHC.size) * coeffs_OHC_trend[0] + coeffs_OHC_trend[1], + c="k", + ls="--", + ) + # ax[0].annotate(f'{coeffs_OHC_trend[0]:.2e}', + # xy=(pos.get_xdata()[0], pos.get_ydata()[0]), + # xytext=(pos.get_xdata()[1], pos.get_ydata()[1]), + # fontsize=9, color='k') + f.write(f"\nOHC GT Trend Slope : {coeffs_OHC_trend[0]}") + # ax.legend() + ax.set_title("") + handles, labels = ax.get_legend_handles_labels() + fig.legend( + handles, labels, loc="lower center", bbox_to_anchor=(0.5, 0.89), ncols=3 + ) + + f.close() + plt.savefig( + os.path.join(self.ohc_path, "OHC_ref0_noanomaly"), + bbox_inches="tight", + dpi=600, + ) + + def ohc_anomaly_global(self, data: xr.Dataset) -> xr.DataArray: + c_p = 3850 # J/(kg C) + rho_0 = 1025 # kg/m^3 + + # Use real areacello for physical calculations + areacello = data["areacello_spherical"] + + OHC = ((data["thetao"] * c_p * rho_0) * areacello * data["dz"]).sum( + ["x", "y", "lev"] + ) / 1e21 + + OHC = remove_climatology(OHC) + + OHC = OHC.rename("OHC Anomaly") + OHC = OHC.assign_attrs(units="ZJ") + return OHC + + def linear_fit(self, var: xr.DataArray) -> tuple[float, float]: + coeffs_trend = np.polyfit(np.arange(var.size), var, 1) + return coeffs_trend[0], coeffs_trend[1] + + def step_ohc_plots(self): + f = open(os.path.join(self.output_path, "compare_info.txt"), "a") + + plt.rcdefaults() + fig, ax = plt.subplots( + 1, 1, figsize=(10, 3), gridspec_kw={"wspace": 0.25, "hspace": 0.5} + ) + plt.rcParams.update({"font.size": 9}) + + OHC = self.ohc_anomaly_global(self.data) + + for i, k in enumerate(self.pred_dict.keys()): + OHC_pred = self.ohc_anomaly_global(self.pred_dict[k]["ds_prediction"]) + OHC_pred.plot(ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i]) # type: ignore + coeffs_OHC_pred_trend = self.linear_fit(OHC_pred) + (pos,) = ax.plot( + OHC_pred.time.data, + np.arange(OHC_pred.size) * coeffs_OHC_pred_trend[0] + + coeffs_OHC_pred_trend[1], + c=self.clist[i], + ls="--", + ) + f.write( + f"\nOHC {self.pred_dict[k]['name']} Trend Slope : {coeffs_OHC_pred_trend[0]}" + ) + self.pred_dict[k]["OHC_slope"] = coeffs_OHC_pred_trend[0] + + OHC.plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + coeffs_OHC_trend = self.linear_fit(OHC) + (pos,) = ax.plot( + OHC.time.data, + np.arange(OHC.size) * coeffs_OHC_trend[0] + coeffs_OHC_trend[1], # type: ignore + c="k", + ls="--", + ) + f.write(f"\nOHC GT Trend Slope : {coeffs_OHC_trend[0]}") + GT_ohc_slope = coeffs_OHC_trend[0] + # ax.legend() + ax.set_title("") + handles, labels = ax.get_legend_handles_labels() + fig.legend( + handles, labels, loc="lower center", bbox_to_anchor=(0.5, 0.89), ncols=3 + ) + + f.close() + plt.savefig(os.path.join(self.ohc_path, "OHC"), bbox_inches="tight", dpi=600) + + return GT_ohc_slope + + def step_depthwise_ohc_plots(self): + plt.rcParams.update({"font.size": 14}) + + Days_to_Eq = 0 + c_p = 3850 # J/(kg C) + rho_0 = 1025 # kg/m^3 + fig, ax = plt.subplots( + 3, 1, figsize=(10, 7.5), gridspec_kw={"wspace": 0.25, "hspace": 0.5} + ) + plt.rcParams.update({"font.size": 9}) + + f = open(os.path.join(self.output_path, "compare_info.txt"), "a") + + # Upper - GT + OHC_truth_upper = ( + (self.data["thetao"].sel(lev=slice(0, 700)) * c_p * rho_0) + * self.data["areacello"] + * self.data["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + OHC_truth_upper = remove_climatology(OHC_truth_upper) + OHC_truth_upper.plot(ax=ax[0], label=self.dataset_name, c="k") # type: ignore + coeffs_OHC_ground_trend = np.polyfit( + np.arange(OHC_truth_upper[Days_to_Eq:].size), + OHC_truth_upper[Days_to_Eq:], + 1, + ) + (pos,) = ax[0].plot( + OHC_truth_upper[Days_to_Eq:].time.data, + np.arange(OHC_truth_upper[Days_to_Eq:].size) * coeffs_OHC_ground_trend[0] + + coeffs_OHC_ground_trend[1], + c="k", + ls="--", + ) + # ax[0].annotate(f'{coeffs_OHC_ground_trend[0]:.2e}', + # xy=(pos.get_xdata()[0], pos.get_ydata()[0]), + # xytext=(pos.get_xdata()[1], pos.get_ydata()[1]), + # fontsize=9, color='k') + f.write(f"\nUpper - GT Trend Slope : {coeffs_OHC_ground_trend[0]}") + GT_upper = coeffs_OHC_ground_trend[0] + upper_trend_truth = coeffs_OHC_ground_trend[0] * 73 + print("upper_trend_truth: ", upper_trend_truth) + + # Upper - Pred + for i, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["OHC_pred_upper"] = ( + ( + self.pred_dict[k]["ds_prediction"]["thetao"].sel(lev=slice(0, 700)) + * c_p + * rho_0 + ) + * self.pred_dict[k]["ds_prediction"]["areacello"] + * self.pred_dict[k]["ds_prediction"]["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + self.pred_dict[k]["OHC_pred_upper"] = remove_climatology( + self.pred_dict[k]["OHC_pred_upper"] + ) + self.pred_dict[k]["OHC_pred_upper"] = self.pred_dict[k][ + "OHC_pred_upper" + ].rename("0-0.7km") + self.pred_dict[k]["OHC_pred_upper"] = self.pred_dict[k][ + "OHC_pred_upper" + ].assign_attrs(units="ZJ") + self.pred_dict[k]["coeffs_OHC_pred_trend_upper"] = np.polyfit( + np.arange(self.pred_dict[k]["OHC_pred_upper"][Days_to_Eq:].size), + self.pred_dict[k]["OHC_pred_upper"][Days_to_Eq:], + 1, + ) + self.pred_dict[k]["OHC_pred_upper"].plot( + ax=ax[0], label=self.pred_dict[k]["name"], c=self.clist[i] + ) + (pos,) = ax[0].plot( + self.pred_dict[k]["OHC_pred_upper"][Days_to_Eq:].time.data, + np.arange(self.pred_dict[k]["OHC_pred_upper"][Days_to_Eq:].size) + * self.pred_dict[k]["coeffs_OHC_pred_trend_upper"][0] + + self.pred_dict[k]["coeffs_OHC_pred_trend_upper"][1], + c=self.clist[i], + ls="--", + ) + # ax[0].annotate(f'{self.pred_dict[k]["coeffs_OHC_pred_trend_upper"][0]:.2e}', + # xy=(pos.get_xdata()[-1], pos.get_ydata()[-1]), + # xytext=(pos.get_xdata()[-2], pos.get_ydata()[-2]), + # fontsize=9, color=self.clist[i]) + f.write( + f"\nUpper - {self.pred_dict[k]['name']} Trend Slope : " + f"{self.pred_dict[k]['coeffs_OHC_pred_trend_upper'][0]}" + ) + self.pred_dict[k]["upper_trend_pred"] = ( + self.pred_dict[k]["coeffs_OHC_pred_trend_upper"][0] * 73 + ) + print( + self.pred_dict[k]["name"], + " upper_trend_pred: ", + self.pred_dict[k]["upper_trend_pred"], + ) + self.pred_dict[k]["total_trend_pred"] = self.pred_dict[k][ + "upper_trend_pred" + ] + + # ax[0].legend() + ax[0].set_title("OHC Anomaly") + + # Middle - GT + OHC_truth_mid = ( + (self.data["thetao"].sel(lev=slice(700, 2000)) * c_p * rho_0) + * self.data["areacello"] + * self.data["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + OHC_truth_mid = remove_climatology(OHC_truth_mid) + OHC_truth_mid.plot(ax=ax[1], label=self.dataset_name, c="k") # type: ignore + coeffs_OHC_ground_trend = np.polyfit( + np.arange(OHC_truth_mid[Days_to_Eq:].size), OHC_truth_mid[Days_to_Eq:], 1 + ) + (pos,) = ax[1].plot( + OHC_truth_mid[Days_to_Eq:].time.data, + np.arange(OHC_truth_mid[Days_to_Eq:].size) * coeffs_OHC_ground_trend[0] + + coeffs_OHC_ground_trend[1], + c="k", + ls="--", + ) + # ax[1].annotate(f'{coeffs_OHC_ground_trend[0]:.2e}', + # xy=(pos.get_xdata()[0], pos.get_ydata()[0]), + # xytext=(pos.get_xdata()[1], pos.get_ydata()[1]), + # fontsize=9, color='k') + f.write(f"\nMiddle - GT Trend Slope : {coeffs_OHC_ground_trend[0]}") + GT_mid = coeffs_OHC_ground_trend[0] + mid_trend_truth = coeffs_OHC_ground_trend[0] * 73 + print("mid_trend_truth: ", mid_trend_truth) + + # Middle - Pred + for i, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["OHC_pred_mid"] = ( + ( + self.pred_dict[k]["ds_prediction"]["thetao"].sel( + lev=slice(700, 2000) + ) + * c_p + * rho_0 + ) + * self.pred_dict[k]["ds_prediction"]["areacello"] + * self.pred_dict[k]["ds_prediction"]["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + self.pred_dict[k]["OHC_pred_mid"] = remove_climatology( + self.pred_dict[k]["OHC_pred_mid"] + ) + self.pred_dict[k]["OHC_pred_mid"] = self.pred_dict[k][ + "OHC_pred_mid" + ].rename("0.7-2.0km") + self.pred_dict[k]["OHC_pred_mid"] = self.pred_dict[k][ + "OHC_pred_mid" + ].assign_attrs(units="ZJ") + self.pred_dict[k]["coeffs_OHC_pred_trend_mid"] = np.polyfit( + np.arange(self.pred_dict[k]["OHC_pred_mid"][Days_to_Eq:].size), + self.pred_dict[k]["OHC_pred_mid"][Days_to_Eq:], + 1, + ) + self.pred_dict[k]["OHC_pred_mid"].plot( + ax=ax[1], label=self.pred_dict[k]["name"], c=self.clist[i] + ) + (pos,) = ax[1].plot( + self.pred_dict[k]["OHC_pred_mid"][Days_to_Eq:].time.data, + np.arange(self.pred_dict[k]["OHC_pred_mid"][Days_to_Eq:].size) + * self.pred_dict[k]["coeffs_OHC_pred_trend_mid"][0] + + self.pred_dict[k]["coeffs_OHC_pred_trend_mid"][1], + c=self.clist[i], + ls="--", + ) + # ax[1].annotate(f'{self.pred_dict[k]["coeffs_OHC_pred_trend_mid"][0]:.2e}', + # xy=(pos.get_xdata()[-1], pos.get_ydata()[-1]), + # xytext=(pos.get_xdata()[-2], pos.get_ydata()[-2]), + # fontsize=9, color=self.clist[i]) + f.write( + f"\nMiddle - {self.pred_dict[k]['name']} Trend Slope : " + f"{self.pred_dict[k]['coeffs_OHC_pred_trend_mid'][0]}" + ) + self.pred_dict[k]["mid_trend_pred"] = ( + self.pred_dict[k]["coeffs_OHC_pred_trend_mid"][0] * 73 + ) + print( + self.pred_dict[k]["name"], + " mid_trend_pred: ", + self.pred_dict[k]["mid_trend_pred"], + ) + self.pred_dict[k]["total_trend_pred"] += self.pred_dict[k]["mid_trend_pred"] + + # ax[1].legend() + ax[1].set_title("") + + # Deep - GT + OHC_truth_deep = ( + (self.data["thetao"].sel(lev=slice(2000, None)) * c_p * rho_0) + * self.data["areacello"] + * self.data["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + OHC_truth_deep = remove_climatology(OHC_truth_deep) + OHC_truth_deep.plot(ax=ax[2], label=self.dataset_name, c="k") # type: ignore + coeffs_OHC_ground_trend = np.polyfit( + np.arange(OHC_truth_deep[Days_to_Eq:].size), OHC_truth_deep[Days_to_Eq:], 1 + ) + (pos,) = ax[2].plot( + OHC_truth_deep[Days_to_Eq:].time.data, + np.arange(OHC_truth_deep[Days_to_Eq:].size) * coeffs_OHC_ground_trend[0] + + coeffs_OHC_ground_trend[1], + c="k", + ls="--", + ) + # ax[2].annotate(f'{coeffs_OHC_ground_trend[0]:.2e}', + # xy=(pos.get_xdata()[0], pos.get_ydata()[0]), + # xytext=(pos.get_xdata()[1], pos.get_ydata()[1]), + # fontsize=9, color='k') + f.write(f"\nDeep - GT Trend Slope : {coeffs_OHC_ground_trend[0]}") + GT_deep = coeffs_OHC_ground_trend[0] + deep_trend_truth = coeffs_OHC_ground_trend[0] * 73 + print("deep_trend_truth: ", deep_trend_truth) + + # Deep - Pred + for i, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["OHC_pred_deep"] = ( + ( + self.pred_dict[k]["ds_prediction"]["thetao"].sel( + lev=slice(2000, None) + ) + * c_p + * rho_0 + ) + * self.pred_dict[k]["ds_prediction"]["areacello"] + * self.pred_dict[k]["ds_prediction"]["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + self.pred_dict[k]["OHC_pred_deep"] = remove_climatology( + self.pred_dict[k]["OHC_pred_deep"] + ) + self.pred_dict[k]["OHC_pred_deep"] = self.pred_dict[k][ + "OHC_pred_deep" + ].rename("2.0-6.0km") + self.pred_dict[k]["OHC_pred_deep"] = self.pred_dict[k][ + "OHC_pred_deep" + ].assign_attrs(units="ZJ") + self.pred_dict[k]["coeffs_OHC_pred_trend_deep"] = np.polyfit( + np.arange(self.pred_dict[k]["OHC_pred_deep"][Days_to_Eq:].size), + self.pred_dict[k]["OHC_pred_deep"][Days_to_Eq:], + 1, + ) + self.pred_dict[k]["OHC_pred_deep"].plot( + ax=ax[2], label=self.pred_dict[k]["name"], c=self.clist[i] + ) + (pos,) = ax[2].plot( + self.pred_dict[k]["OHC_pred_deep"][Days_to_Eq:].time.data, + np.arange(self.pred_dict[k]["OHC_pred_deep"][Days_to_Eq:].size) + * self.pred_dict[k]["coeffs_OHC_pred_trend_deep"][0] + + self.pred_dict[k]["coeffs_OHC_pred_trend_deep"][1], + c=self.clist[i], + ls="--", + ) + # ax[2].annotate(f'{self.pred_dict[k]["coeffs_OHC_pred_trend_deep"][0]:.2e}', + # xy=(pos.get_xdata()[-1], pos.get_ydata()[-1]), + # xytext=(pos.get_xdata()[-2], pos.get_ydata()[-2]), + # fontsize=9, color=self.clist[i]) + f.write( + f"\nDeep - {self.pred_dict[k]['name']} Trend Slope : " + f"{self.pred_dict[k]['coeffs_OHC_pred_trend_deep'][0]}" + ) + + self.pred_dict[k]["deep_trend_pred"] = ( + self.pred_dict[k]["coeffs_OHC_pred_trend_deep"][0] * 73 + ) + print( + self.pred_dict[k]["name"], + " deep_trend_pred: ", + self.pred_dict[k]["deep_trend_pred"], + ) + self.pred_dict[k]["total_trend_pred"] += self.pred_dict[k][ + "deep_trend_pred" + ] + + # ax[2].legend() + ax[2].set_title("") + total_trend_truth = upper_trend_truth + mid_trend_truth + deep_trend_truth + + handles, labels = ax[0].get_legend_handles_labels() + fig.legend( + handles, labels, loc="lower center", bbox_to_anchor=(0.5, 0.91), ncols=3 + ) + + f.write( + f"\nGT Trend Ratio (Upper, Mid, Deep): " + f"{upper_trend_truth / total_trend_truth:.2f}, " + f"{mid_trend_truth / total_trend_truth:.2f}, " + f"{deep_trend_truth / total_trend_truth:.2f}" + ) + for k in self.pred_dict.keys(): + total_trend = self.pred_dict[k]["total_trend_pred"] + f.write( + f"\n{self.pred_dict[k]['name']} Trend Ratio (Upper, Mid, Deep): " + f"{self.pred_dict[k]['upper_trend_pred'] / total_trend:.2f}, " + f"{self.pred_dict[k]['mid_trend_pred'] / total_trend:.2f}, " + f"{self.pred_dict[k]['deep_trend_pred'] / total_trend:.2f}" + ) + + f.write("\n") + f.close() + + plt.savefig( + os.path.join(self.ohc_path, "OHC_Timeseries_depths"), + bbox_inches="tight", + dpi=600, + ) + + pd_data = [] + pd_data.append( + { + "Model": self.dataset_name, + "Upper": GT_upper, + "Middle": GT_mid, + "Deep": GT_deep, + } + ) + + for k in self.pred_dict.keys(): + pd_data.append( + { + "Model": self.pred_dict[k]["name"], + "Upper": self.pred_dict[k]["coeffs_OHC_pred_trend_upper"][0], + "Upper Slope Ratio": self.pred_dict[k][ + "coeffs_OHC_pred_trend_upper" + ][0] + / GT_upper, + "Middle": self.pred_dict[k]["coeffs_OHC_pred_trend_mid"][0], + "Middle Slope Ratio": self.pred_dict[k][ + "coeffs_OHC_pred_trend_mid" + ][0] + / GT_mid, + "Deep": self.pred_dict[k]["coeffs_OHC_pred_trend_deep"][0], + "Deep Slope Ratio": self.pred_dict[k]["coeffs_OHC_pred_trend_deep"][ + 0 + ] + / GT_deep, + } + ) + + # Create a DataFrame + df = pd.DataFrame(pd_data) + + # Define the file path + file_path = os.path.join(self.ohc_path, "depthwise_ohc_slopes_table.csv") + + # Save the DataFrame to a CSV file + df.to_csv(file_path, index=False) + + def step_basin_ohc_plots(self): + f = open(os.path.join(self.output_path, "compare_info.txt"), "a") + + c_p = 3850 # J/(kg C) + rho_0 = 1025 # kg/m^3 + + plt.rcParams.update({"font.size": 10}) + fig, ax = plt.subplots( + 2, + 3, + figsize=(16, 8), + gridspec_kw={ + "width_ratios": [1, 1, 1], + "height_ratios": [1, 1], + "wspace": 0.25, + "hspace": 0.5, + }, + ) + + ax_flat = ax.flatten() + + GT_regionwise_ohc = {} + GT_regionwise_ohc["Model"] = self.dataset_name + for j, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["regionwise_ohc"] = {} + + for i, var in enumerate(list(self.basin_masks.keys())): + var = str(var) + OHC = ( + (self.data["thetao"] * c_p * rho_0 * self.basin_masks[var]) + * self.data["areacello"] + * self.data["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + OHC = remove_climatology(OHC) + OHC = OHC.rename("OHC Anomaly") + OHC = OHC.assign_attrs(units="ZJ") + coeffs_OHC_trend = np.polyfit(np.arange(OHC.size), OHC, 1) + OHC.plot(ax=ax_flat[i], label=self.dataset_name, c="k") # type: ignore + (pos,) = ax_flat[i].plot( + OHC.time.data, + np.arange(OHC.size) * coeffs_OHC_trend[0] + coeffs_OHC_trend[1], + c="k", + ls="--", + ) + # ax_flat[i].annotate(f'{coeffs_OHC_trend[0]:.2e}', + # xy=(pos.get_xdata()[0], pos.get_ydata()[0]), + # xytext=(pos.get_xdata()[1], pos.get_ydata()[1]), + # fontsize=9, color='k') + f.write(f"\nOHC {var} GT Trend Slope : {coeffs_OHC_trend[0]}") + GT_regionwise_ohc[var] = coeffs_OHC_trend[0] + for j, k in enumerate(self.pred_dict.keys()): + OHC_pred = ( + ( + self.pred_dict[k]["ds_prediction"]["thetao"] + * c_p + * rho_0 + * self.basin_masks[var] + ) + * self.pred_dict[k]["ds_prediction"]["areacello"] + * self.pred_dict[k]["ds_prediction"]["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + OHC_pred = remove_climatology(OHC_pred) + OHC_pred = OHC_pred.rename("OHC Anomaly") + OHC_pred = OHC_pred.assign_attrs(units="ZJ") + coeffs_OHC_pred_trend = np.polyfit( + np.arange(OHC_pred.size), OHC_pred, 1 + ) + OHC_pred.plot( + ax=ax_flat[i], label=self.pred_dict[k]["name"], c=self.clist[j] + ) + (pos,) = ax_flat[i].plot( + OHC_pred.time.data, + np.arange(OHC_pred.size) * coeffs_OHC_pred_trend[0] + + coeffs_OHC_pred_trend[1], + c=self.clist[j], + ls="--", + ) + # ax_flat[i].annotate(f'{coeffs_OHC_pred_trend[0]:.2e}', + # xy=(pos.get_xdata()[-1], pos.get_ydata()[-1]), + # xytext=(pos.get_xdata()[-2], pos.get_ydata()[-2]), + # fontsize=9, color=self.clist[j]) + f.write( + f"\nOHC {var} {self.pred_dict[k]['name']} Trend Slope : " + f"{coeffs_OHC_pred_trend[0]}" + ) + self.pred_dict[k]["regionwise_ohc"][var] = coeffs_OHC_pred_trend[0] + + ax_flat[i].set_title(var + " Ocean") + + fig.delaxes(ax_flat[5]) + handles, labels = ax_flat[0].get_legend_handles_labels() + fig.legend( + handles, labels, loc="lower center", bbox_to_anchor=(0.5, 0.93), ncols=3 + ) + + f.write("\n") + f.close() + # plt.show() + plt.savefig( + os.path.join(self.ohc_path, "OHC_Basin"), bbox_inches="tight", dpi=600 + ) + + pd_data = [] + pd_data.append(GT_regionwise_ohc) + + for k in self.pred_dict.keys(): + k = str(k) + d = {} + d["Model"] = self.pred_dict[k]["name"] + for var in self.basin_masks.keys(): + var = str(var) + d[var] = self.pred_dict[k]["regionwise_ohc"][var] + d[var + " Slope Ratio"] = ( + self.pred_dict[k]["regionwise_ohc"][var] / GT_regionwise_ohc[var] + ) + pd_data.append(d) + + # Create a DataFrame + df = pd.DataFrame(pd_data) + + # Define the file path + file_path = os.path.join(self.ohc_path, "regionwise_ohc_slopes_table.csv") + + # Save the DataFrame to a CSV file + df.to_csv(file_path, index=False) + + def step_basin_ohc_upto_700_plots(self): + # TODO(jder): this is a copy-paste of the above with a limit + + f = open(os.path.join(self.output_path, "compare_info.txt"), "a") + + c_p = 3850 # J/(kg C) + rho_0 = 1025 # kg/m^3 + max_level = 700 + + plt.rcParams.update({"font.size": 10}) + fig, ax = plt.subplots( + 2, + 3, + figsize=(16, 8), + gridspec_kw={ + "width_ratios": [1, 1, 1], + "height_ratios": [1, 1], + "wspace": 0.25, + "hspace": 0.5, + }, + ) + + ax_flat = ax.flatten() + + GT_regionwise_ohc = {} + GT_regionwise_ohc["Model"] = self.dataset_name + for j, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["regionwise_ohc"] = {} + + for i, var in enumerate(list(self.basin_masks.keys())): + var = str(var) + OHC = ( + ( + self.data["thetao"].sel(lev=slice(None, max_level)) + * c_p + * rho_0 + * self.basin_masks[var] + ) + * self.data["areacello"] + * self.data["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + OHC = remove_climatology(OHC) + OHC = OHC.rename(f"OHC Anomaly (Upto {max_level}m)") + OHC = OHC.assign_attrs(units="ZJ") + coeffs_OHC_trend = np.polyfit(np.arange(OHC.size), OHC, 1) + OHC.plot(ax=ax_flat[i], label=self.dataset_name, c="k") # type: ignore + (pos,) = ax_flat[i].plot( + OHC.time.data, + np.arange(OHC.size) * coeffs_OHC_trend[0] + coeffs_OHC_trend[1], + c="k", + ls="--", + ) + f.write(f"\nOHC {var} GT Trend Slope : {coeffs_OHC_trend[0]}") + GT_regionwise_ohc[var] = coeffs_OHC_trend[0] + for j, k in enumerate(self.pred_dict.keys()): + OHC_pred = ( + ( + self.pred_dict[k]["ds_prediction"]["thetao"].sel( + lev=slice(None, max_level) + ) + * c_p + * rho_0 + * self.basin_masks[var] + ) + * self.pred_dict[k]["ds_prediction"]["areacello"] + * self.pred_dict[k]["ds_prediction"]["dz"] + ).sum(["x", "y", "lev"]) / 1e21 + + OHC_pred = remove_climatology(OHC_pred) + OHC_pred = OHC_pred.rename(f"OHC Anomaly (Upto {max_level}m)") + OHC_pred = OHC_pred.assign_attrs(units="ZJ") + coeffs_OHC_pred_trend = np.polyfit( + np.arange(OHC_pred.size), OHC_pred, 1 + ) + OHC_pred.plot( + ax=ax_flat[i], label=self.pred_dict[k]["name"], c=self.clist[j] + ) + (pos,) = ax_flat[i].plot( + OHC_pred.time.data, + np.arange(OHC_pred.size) * coeffs_OHC_pred_trend[0] + + coeffs_OHC_pred_trend[1], + c=self.clist[j], + ls="--", + ) + # ax_flat[i].annotate(f'{coeffs_OHC_pred_trend[0]:.2e}', + # xy=(pos.get_xdata()[-1], pos.get_ydata()[-1]), + # xytext=(pos.get_xdata()[-2], pos.get_ydata()[-2]), + # fontsize=9, color=self.clist[j]) + f.write( + f"\nOHC {var} {self.pred_dict[k]['name']} Trend Slope : " + f"{coeffs_OHC_pred_trend[0]}" + ) + self.pred_dict[k]["regionwise_ohc"][var] = coeffs_OHC_pred_trend[0] + + ax_flat[i].set_title(var + " Ocean") + + fig.delaxes(ax_flat[5]) + handles, labels = ax_flat[0].get_legend_handles_labels() + fig.legend( + handles, labels, loc="lower center", bbox_to_anchor=(0.5, 0.93), ncol=3 + ) + + f.write("\n") + f.close() + # plt.show() + plt.savefig( + os.path.join(self.ohc_path, f"OHC_Basin_upto_{max_level}m"), + bbox_inches="tight", + dpi=600, + ) + + def step_ocean_temperature_profile_plots(self): + def ocean_temperature_profile(datasets, titles, plot_title): + # We import here to avoid importing tk unconditionally above + from xarrayutils.plotting import linear_piecewise_scale # type: ignore + + fig, axs = plt.subplots( + 2, + 3, + figsize=(16, 6), + gridspec_kw={ + "width_ratios": [0.05] * 3, + "height_ratios": [0.05] * 2, + "wspace": 0.01, + "hspace": 0.5, + }, + dpi=300, + ) + fig.suptitle(plot_title, fontsize=15, fontweight="bold", y=0.95) + + ax = axs.flatten() + norm = symmetric_percentile_norm(datasets) + + for i, (data, title) in enumerate(zip(datasets, titles)): + data = data.rename(r"$\theta_O$").assign_attrs(units=r"$\degree C$") + data["y"] = data.y.assign_attrs( + long_name="Latitude", units=r"$\degree$" + ) + data["lev"] = data.lev.assign_attrs(long_name="depth", units="m") + + im = data.plot(ax=ax[i], cmap="bwr", norm=norm, add_colorbar=False) + ax[i].invert_yaxis() + + linear_piecewise_scale(1000, 5, ax=ax[i]) + ax[i].axhline(1000, color="0.5", ls="--") + ax[i].set_yticks([0, 250, 500, 750, 1000, 3000, 5000]) + ax[i].set_xticks([-60, -30, 0, 30, 60]) + ax[i].set_xticklabels( + [ + r"$60^\circ S$", + r"$30^\circ S$", + "0", + r"$30^\circ N$", + r"$60^\circ N$", + ] + ) + ax[i].set_title(title, fontsize=14) + ax[i].set_box_aspect(0.7) + + cbar = fig.colorbar( + im, ax=ax[:], orientation="vertical", fraction=0.02, pad=0.02 + ) + cbar.set_label(r"$\theta_O$ [$\degree C$]") + # plt.show() + plt.savefig( + os.path.join(self.temp_path, plot_title), bbox_inches="tight", dpi=600 + ) + + data_thetao_ano = remove_climatology(self.data.thetao) + pred_thetao_ano = remove_climatology( + self.pred_dict[self.key1]["ds_prediction"]["thetao"] + ) + + # Full time mean and bias [Last year - First year] + GT_lastyear_change = data_thetao_ano.isel(time=slice(-73, None)).mean( + dim="time" + ) - data_thetao_ano.isel(time=slice(0, 73)).mean(dim="time") + datasets, titles = get_basin_datasets( + GT_lastyear_change, self.basin_masks, self.data + ) + plot_title = f"{self.dataset_name} (Last Year - First Year)" + ocean_temperature_profile(datasets, titles, plot_title) + + pred_lastyear_change = pred_thetao_ano.isel(time=slice(-73, None)).mean( + dim="time" + ) - pred_thetao_ano.isel(time=slice(0, 73)).mean(dim="time") + datasets, titles = get_basin_datasets( + pred_lastyear_change, self.basin_masks, self.data + ) + plot_title = f"{self.pred_dict[self.key1]['name']} (Last Year - First Year)" + ocean_temperature_profile(datasets, titles, plot_title) + + bias_lastyear_change = pred_lastyear_change - GT_lastyear_change + datasets, titles = get_basin_datasets( + bias_lastyear_change, self.basin_masks, self.data + ) + plot_title = f"(Last Year - First Year) Bias" + ocean_temperature_profile(datasets, titles, plot_title) + + def salinity_global(self, data: xr.Dataset) -> xr.DataArray: + rho_0 = 1025 # kg/m^3 + salinity = ((data["so"] * rho_0) * data["areacello"] * data["dz"]).sum( + ["x", "y", "lev"] + ) + salinity = salinity.rename("Salinity") + salinity = salinity.assign_attrs(units="g") + return salinity + + def step_ocean_salinity_profile_plots(self): + f = open(os.path.join(self.output_path, "compare_info.txt"), "a") + + plt.rcdefaults() + fig, ax = plt.subplots(1, 1, figsize=(10, 5)) # Single axis for salinity plot + plt.rcParams.update({"font.size": 9}) + + salinity = self.salinity_global(self.data) + + for i, k in enumerate(self.pred_dict.keys()): + if "so" in self.pred_dict[k]["ls"]: + salinity_pred = self.salinity_global(self.pred_dict[k]["ds_prediction"]) + salinity_pred.plot( # type: ignore + ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i] + ) + coeffs_salinity_pred_trend = np.polyfit( + np.arange(salinity_pred.size), salinity_pred, 1 + ) + (pos,) = ax.plot( + salinity_pred.time.data, + np.arange(salinity_pred.size) * coeffs_salinity_pred_trend[0] + + coeffs_salinity_pred_trend[1], + c=self.clist[i], + ls="--", + ) + f.write( + f"\nSalinity {self.pred_dict[k]['name']} Trend Slope : " + f"{coeffs_salinity_pred_trend[0]}" + ) + self.pred_dict[k]["salinity_slope"] = coeffs_salinity_pred_trend[0] + + coeffs_salinity_trend = self.linear_fit(salinity) + salinity.plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + (pos,) = ax.plot( + salinity.time.data, + np.arange(salinity.size) * coeffs_salinity_trend[0] + + coeffs_salinity_trend[1], + c="k", + ls="--", + ) + f.write(f"\nSalinity GT Trend Slope : {coeffs_salinity_trend[0]}") + ax.set_ylim((5.861e22, 5.8632e22)) + ax.legend(ncol=3) + ax.set_title("") + f.write("\n") + f.close() + + print(coeffs_salinity_trend[0] * 73) + plt.savefig( + os.path.join(self.salinity_path, "Salinity"), bbox_inches="tight", dpi=600 + ) + + return coeffs_salinity_trend[0] + + def step_create_ohc_salinity_slopes_table(self): + """Create a CSV table with OHC and salinity slopes.""" + GT_ohc_slope = self.linear_fit(self.ohc_anomaly_global(self.data))[0] + GT_salinity_slope = self.linear_fit(self.salinity_global(self.data))[0] + + pd_data = [] + pd_data.append( + { + "Model": self.dataset_name, + "OHC": GT_ohc_slope, + "Salinity": GT_salinity_slope, + } + ) + + for k in self.pred_dict.keys(): + pd_data.append( + { + "Model": self.pred_dict[k]["name"], + "OHC": self.pred_dict[k]["OHC_slope"], + "OHC Slope Ratio": self.pred_dict[k]["OHC_slope"] / GT_ohc_slope, + "Salinity": self.pred_dict[k]["salinity_slope"], + "Salinity Slope Ratio": ( + self.pred_dict[k]["salinity_slope"] / GT_salinity_slope + ), + } + ) + + # Create a DataFrame + df = pd.DataFrame(pd_data) + + # Define the file path + file_path = os.path.join(self.output_path, "ohc_salinity_slopes_table.csv") + + # Save the DataFrame to a CSV file + df.to_csv(file_path, index=False) + + def step_salinity_deseasonalized_plots(self): + f = open( + os.path.join(self.metrics_path, "salinity_deseasonalized_info.txt"), "w" + ) + + plt.rcdefaults() + fig, ax = plt.subplots( + 1, 1, figsize=(10, 3), gridspec_kw={"wspace": 0.25, "hspace": 0.5} + ) + plt.rcParams.update({"font.size": 9}) + + salinity = ( + self.data["so"] + .weighted(self.data["areacello"] * self.data["dz"]) + .mean(["x", "y", "lev"]) + ) + + salinity = remove_climatology(salinity) + salinity = salinity.rename("Salinity") + salinity = salinity.assign_attrs(units="psu") + + for i, k in enumerate(self.pred_dict.keys()): + if "so" in self.pred_dict[k]["ls"]: + salinity_pred = ( + self.pred_dict[k]["ds_prediction"]["so"] + .weighted( + self.pred_dict[k]["ds_prediction"]["areacello"] + * self.pred_dict[k]["ds_prediction"]["dz"] + ) + .mean(["x", "y", "lev"]) + ) + salinity_pred = remove_climatology(salinity_pred) + salinity_pred = salinity_pred.rename("Salinity") + salinity_pred = salinity_pred.assign_attrs(units="psu") + salinity_pred.plot( + ax=ax, label=self.pred_dict[k]["name"], c=self.clist[i] + ) + coeffs_salinity_pred_trend = np.polyfit( + np.arange(salinity_pred.size), salinity_pred, 1 + ) + (pos,) = ax.plot( + salinity_pred.time.data, + np.arange(salinity_pred.size) * coeffs_salinity_pred_trend[0] + + coeffs_salinity_pred_trend[1], + c=self.clist[i], + ls="--", + ) + # ax[1].annotate(f'{coeffs_salinity_pred_trend[0]:.2e}', + # xy=(pos.get_xdata()[-1], pos.get_ydata()[-1]), + # xytext=(pos.get_xdata()[-2], pos.get_ydata()[-2]), + # fontsize=9, color=self.clist[i]) + f.write( + f"\nSalinity {self.pred_dict[k]['name']} Trend Slope : " + f"{coeffs_salinity_pred_trend[0]}" + ) + self.pred_dict[k]["salinity_slope_des"] = coeffs_salinity_pred_trend[0] + + coeffs_salinity_trend = np.polyfit(np.arange(salinity.size), salinity, 1) + salinity.plot(ax=ax, label=self.dataset_name, c="k") # type: ignore + (pos,) = ax.plot( + salinity.time.data, + np.arange(salinity.size) * coeffs_salinity_trend[0] + + coeffs_salinity_trend[1], + c="k", + ls="--", + ) + ax.set_title("") + # ax[1].annotate(f'{coeffs_salinity_trend[0]:.2e}', + # xy=(pos.get_xdata()[0], pos.get_ydata()[0]), + # xytext=(pos.get_xdata()[1], pos.get_ydata()[1]), + # fontsize=9, color='k') + f.write(f"\nSalinity GT Trend Slope : {coeffs_salinity_trend[0]}") + GT_salinity_slope = coeffs_salinity_trend[0] + f.write("\n") + f.close() + + print(coeffs_salinity_trend[0] * 73) + plt.savefig( + os.path.join(self.salinity_path, "salinity_deseasonalized"), + bbox_inches="tight", + dpi=600, + ) + return GT_salinity_slope + + def step_thetao_mae_metrics(self): + da_temp = self.data["thetao"] # Directly use temperature variable + section_mask = isnan(da_temp).all("x").isel(time=0) + da_temp_int_x = da_temp.weighted(self.data["areacello"]).mean(["x", "time"]) + temp_pred = da_temp_int_x.where(~section_mask) + GT_temp_pred = temp_pred + + for j, (model_key, model_data) in enumerate(self.pred_dict.items(), start=1): + da_temp = model_data["ds_prediction"]["thetao"] # Use temperature variable + section_mask = isnan(da_temp).all("x").isel(time=0) + da_temp_int_x = da_temp.weighted(self.data["areacello"]).mean(["x", "time"]) + temp_pred = da_temp_int_x.where(~section_mask) + self.pred_dict[model_key]["temp_profile"] = temp_pred + + f = open(os.path.join(self.metrics_path, "thetao_mae_info.txt"), "w") + + for i, key in enumerate(self.pred_dict): + mae_key = ( + np.abs(self.pred_dict[key]["temp_profile"] - GT_temp_pred) + .mean() + .compute() + ) + f.write(f"\n Thetao {self.pred_dict[key]['name']} MAE : {mae_key.item()}") + + f.close() + + def step_sst_mae_metrics(self): + section_mask = isnan(self.data["thetao"]).isel(lev=0).isel(time=5) + SST_gt = self.data["thetao"].isel(lev=0).mean("time") + SST_gt = SST_gt.where(~section_mask) + + for j, (model_key, model_data) in enumerate(self.pred_dict.items(), start=1): + section_mask = ( + isnan(model_data["ds_prediction"]["thetao"]).isel(lev=0).isel(time=5) + ) + SST_pred = model_data["ds_prediction"]["thetao"].isel(lev=0).mean("time") + SST_pred = SST_pred.where(~section_mask) + self.pred_dict[model_key]["sst"] = SST_pred + + f = open(os.path.join(self.metrics_path, "sst_mae_info.txt"), "w") + + for i, key in enumerate(self.pred_dict): + mae_key = np.abs(self.pred_dict[key]["sst"] - SST_gt).mean().compute() + f.write(f"\n SST {self.pred_dict[key]['name']} MAE : {mae_key.item()}") + + f.close() + + def step_pdf_plots_short(self): + plt.rcParams.update({"font.size": 9}) + # Create a figure + fig = plt.figure(figsize=(24, 15)) + plt.rc("axes", titlesize=30) # fontsize of the axes title + plt.rc("axes", labelsize=30) # fontsize of the x and y labels + plt.rc("xtick", labelsize=30) # fontsize of the tick labels + plt.rc("ytick", labelsize=30) # fontsize of the tick labels + plt.rc("legend", fontsize=20) # legend fontsize + plt.rc("figure", titlesize=30) + # Manual positioning using add_axes with uniform width and height + width = 0.22 + height = 0.3 + + # Top row: 3 plots, evenly spaced horizontally + axs = [ + fig.add_axes((0.05, 0.55, width, height)), # First plot in top row + fig.add_axes((0.38, 0.55, width, height)), # Second plot in top row + fig.add_axes((0.71, 0.55, width, height)), # Third plot in top row + ] + + # Bottom row: 2 plots centered, manually positioned + axs += [ + fig.add_axes((0.22, 0.1, width, height)), # First plot in bottom row + fig.add_axes((0.54, 0.1, width, height)), # Second plot in bottom row + ] + + # Plot PDFs + for i, v in enumerate(["thetao", "so", "zos", "uo", "vo"]): + min_val, max_val = ( + float(self.data[v].min().values), + float(self.data[v].max().values), + ) + true_pdf, bins_true = np.histogram( + self.data[v], bins=150, density=True, range=(min_val, max_val) + ) + axs[i].semilogy( + bins_true[:-1], true_pdf, label=self.dataset_name, color="k", lw=8 + ) + + for j, k in enumerate(self.pred_dict.keys()): + if v in self.pred_dict[k]["ls"]: + pdf_net, bins_net = np.histogram( + self.pred_dict[k]["ds_prediction"][v], + bins=150, + density=True, + range=(min_val, max_val), + ) + axs[i].semilogy( + bins_net[:-1], + pdf_net, + label=self.pred_dict[k]["name"], + color=self.clist[j], + lw=2, + ) + + axs[i].xaxis.set_major_locator(MaxNLocator(5, prune="both")) + if i == 0: + axs[i].legend() + axs[i].set_xlabel( + r"" + self.data[v].long_name + "[" + self.data[v].units + "]" + ) + axs[i].set_ylabel(r"${p(}$" + self.data[v].long_name + " " + "${)}$") + + if v not in ["thetao", "SSH"]: + axs[i].set_ylim( + (min(true_pdf.min(), pdf_net.min()) + 1e-5, true_pdf.max()) + ) + else: + axs[i].set_ylim((1e-3, true_pdf.max())) + matplotlib.style.use("default") + + # Save or show the figure + # plt.show() + plt.savefig( + os.path.join(self.pdfs_path, "PDF_Plots_Short"), + bbox_inches="tight", + dpi=600, + ) + + def step_enso_plots(self): + clim = ( + self.data["thetao"] + .sel(lev=slice(0, 500)) + .groupby("time.dayofyear") + .mean() + .compute() + ) + data_surface = self.data.sel(lev=slice(0, 500)) + for k in self.pred_dict.keys(): + self.pred_dict[k]["ds_prediction_surface"] = self.pred_dict[k][ + "ds_prediction" + ].sel(lev=slice(0, 500)) + self.pred_dict[k]["clim_pred"] = ( + self.pred_dict[k]["ds_prediction_surface"]["thetao"] + .groupby("time.dayofyear") + .mean() + .compute() + ) + + def NinoIndexComputeClim(T, area, dt=5, window=150): + T = T.load() + T_clim = T.copy() + T_clim = T_clim.sel(x=slice(190, 240), y=slice(-5, 5)) + area = area.sel(x=slice(190, 240), y=slice(-5, 5)).load() + clim = T_clim.groupby("time.dayofyear").mean("time").compute() + window = int(window / dt) + for i, t in enumerate(T_clim.time.values): + day = int(t.dayofyr) + T_clim[i] = (T[i] - clim.sel(dayofyear=day)).data + + T_clim = T_clim.rolling(time=window).mean() + # T_clim = (T_clim*area).sum(["x","y"])/area.sum(["x","y"]) + T_clim = T_clim.weighted(area).mean(["x", "y"]) + + return T_clim[window:] + + nino_true_compute_clim = NinoIndexComputeClim( + data_surface["thetao"][:, 0], self.data["areacello"] + ) + nino_true_compute_clim = nino_true_compute_clim.rename("Nino 3.4") + nino_true_compute_clim = nino_true_compute_clim.assign_attrs( + units=r"$\degree C$" + ) + + for k in self.pred_dict.keys(): + self.pred_dict[k]["nino_pred_compute_clim"] = NinoIndexComputeClim( + self.pred_dict[k]["ds_prediction_surface"]["thetao"][:, 0], + self.pred_dict[k]["ds_prediction"]["areacello"], + ) + self.pred_dict[k]["nino_pred_compute_clim"] = self.pred_dict[k][ + "nino_pred_compute_clim" + ].rename("Nino 3.4") + self.pred_dict[k]["nino_pred_compute_clim"] = self.pred_dict[k][ + "nino_pred_compute_clim" + ].assign_attrs(units=r"$\degree C$") + + day_max = int( + ( + np.argwhere( + nino_true_compute_clim.values + == np.nanmax(nino_true_compute_clim.values[5:-5]) + ) + + 30 + ).squeeze() + ) + day_min = int( + ( + np.argwhere( + nino_true_compute_clim.values + == np.nanmin(nino_true_compute_clim.values[5:-5]) + ) + + 30 + ).squeeze() + ) + + plt.rcParams.update({"font.size": 14}) + plt.figure(figsize=(10, 5)) + nino_true_compute_clim.plot(label=self.dataset_name, c="k") + for i, k in enumerate(self.pred_dict.keys()): + self.pred_dict[k]["nino_pred_compute_clim"].plot( + label=self.pred_dict[k]["name"], c=self.clist[i] + ) + + ax = plt.gca() + ax.legend() + ax.set_title("Nino 3.4 Index") + + plt.savefig( + os.path.join(self.enso_path, "Climatology"), bbox_inches="tight", dpi=600 + ) + + for k in self.pred_dict.keys(): + mae = np.abs( + ( + self.pred_dict[k]["nino_pred_compute_clim"] - nino_true_compute_clim + ).mean(["time"]) + ) + cor = ( + (self.pred_dict[k]["nino_pred_compute_clim"] * nino_true_compute_clim) + / np.sqrt( + (self.pred_dict[k]["nino_pred_compute_clim"] ** 2) + * (nino_true_compute_clim**2) + ) + ).mean() + print(self.pred_dict[k]["name"], "mae: ", mae.item(), "cor: ", cor.item()) + + # ### Profiles and Maps + + keys = list(self.pred_dict.keys()) + # assert len(keys) >= 2, "Maps supported by atleast two keys" + self.key1 = keys[0] + + # #### ENSO Maps + + plt.rcParams.update({"font.size": 14}) + fig, axs = plt.subplot_mosaic( + [ + ["time series", "map"], + ["nino_true", "nina_true"], + ["nino_pred", "nina_pred"], + ["colorbar", "colorbar"], + ], + figsize=(16, 15), + per_subplot_kw={"map": dict(projection=ccrs.Robinson(190))}, + gridspec_kw={ + "width_ratios": [1, 1], + "height_ratios": [0.5, 0.3, 0.3, 0.05], + "wspace": 0.25, + "hspace": 0.5, + }, + ) + + ############################################ + # Time Series + ############################################ + self.pred_dict[self.key1]["nino_pred_compute_clim"].plot( + label=self.pred_dict[self.key1]["name"], + c=self.clist[0], + ax=axs["time series"], + ) + + nino_true_compute_clim.plot( + label=self.dataset_name, c="k", ax=axs["time series"] + ) + nino_true_compute_clim.isel( + time=slice(day_max - 30, day_max - 30 + 1) + ).drop_vars(["lev"]).plot.scatter(s=80, c="k", ax=axs["time series"]) + nino_true_compute_clim.isel( + time=slice(day_min - 30, day_min - 30 + 1) + ).drop_vars(["lev"]).plot.scatter(s=80, c="k", ax=axs["time series"]) + + axs["time series"].set_title("") + axs["time series"].set_xlabel("") + axs["time series"].legend(ncols=3, bbox_to_anchor=(0.5, 0.89)) + # axs['time series'].set_title('Nino 3.4 Index') + + ############################################ + # El Nino Tropics Profiles + ############################################ + + day_start = day_max + window = 3 + + time_slice = slice( + self.pred_dict[self.key1]["ds_prediction_surface"]["time"][day_start], + self.pred_dict[self.key1]["ds_prediction_surface"]["time"][ + day_start + window + ], + ) + times = self.pred_dict[self.key1]["ds_prediction"]["time"][ + day_start : day_start + window + ].data + days_of_year = [i.dayofyr for i in times] + true_clim_to_remove = clim.sel(dayofyear=days_of_year).rename( + {"dayofyear": "time"} + ) + true_clim_to_remove["time"] = times + pred_clim_to_remove = ( + self.pred_dict[self.key1]["clim_pred"] + .sel(dayofyear=days_of_year) + .rename({"dayofyear": "time"}) + ) + pred_clim_to_remove["time"] = times + + tropics_profile = ( + data_surface["thetao"].sel( + time=time_slice, x=slice(118, 260), y=slice(-5, 5) + ) + - true_clim_to_remove.sel(x=slice(118, 260), y=slice(-5, 5)) + ).mean(["time", "y"]) + tropics_profile_pred = ( + self.pred_dict[self.key1]["ds_prediction_surface"]["thetao"][ + day_start : day_start + window + ].sel(lev=slice(0, 500), x=slice(118, 260), y=slice(-5, 5)) + - pred_clim_to_remove.sel(x=slice(118, 260), y=slice(-5, 5)) + ).mean(["time", "y"]) + tropics_profile = tropics_profile.rename("Anomaly") + tropics_profile_pred = tropics_profile_pred.rename("Anomaly") + tropics_profile = tropics_profile.assign_attrs(units=r"${^oC}$") + tropics_profile_pred = tropics_profile_pred.assign_attrs(units=r"${^oC}$") + tropics_profile_pred["x"] = tropics_profile_pred["x"].assign_attrs( + units=r"${^o}$" + ) + + tropics_profile.plot.pcolormesh( + ax=axs["nino_true"], + y="lev", + cmap=cm.cm.diff, + vmin=-2, + vmax=2, + add_colorbar=False, + ) + axs["nino_true"].set_title("Nino Conditions GT") + axs["nino_true"].set_xlabel("") + axs["nino_true"].invert_yaxis() + tropics_profile_pred.plot.pcolormesh( + ax=axs["nino_pred"], + y="lev", + cmap=cm.cm.diff, + vmin=-2, + vmax=2, + add_colorbar=False, + ) + axs["nino_pred"].set_title( + f"Nino Conditions {self.pred_dict[self.key1]['name']}" + ) + axs["nino_pred"].invert_yaxis() + + ## Stats + for name, pred in zip( + [self.pred_dict[self.key1]["name"]], + [tropics_profile_pred], + ): + mae = np.abs( + (pred - tropics_profile).weighted(self.data["dz"]).mean(["x", "lev"]) + ) + cor = (pred * tropics_profile).weighted(self.data["dz"]).mean( + ["x", "lev"] + ) / np.sqrt( + (pred**2).weighted(self.data["dz"]).mean(["x", "lev"]) + * (tropics_profile**2).weighted(self.data["dz"]).mean(["x", "lev"]) + ) + + print( + name, + "El Nino mae: ", + mae.compute().item(), + "El Nino cor: ", + cor.compute().item(), + ) + + ############################################ + # La Nina Tropics Profiles + ############################################ + + day_start = day_min + window = 3 + + time_slice = slice( + self.pred_dict[self.key1]["ds_prediction"]["time"][day_start], + self.pred_dict[self.key1]["ds_prediction"]["time"][day_start + window], + ) + times = self.pred_dict[self.key1]["ds_prediction"]["time"][ + day_start : day_start + window + ].data + days_of_year = [i.dayofyr for i in times] + true_clim_to_remove = clim.sel(dayofyear=days_of_year).rename( + {"dayofyear": "time"} + ) + true_clim_to_remove["time"] = times + pred_clim_to_remove = ( + self.pred_dict[self.key1]["clim_pred"] + .sel(dayofyear=days_of_year) + .rename({"dayofyear": "time"}) + ) + pred_clim_to_remove["time"] = times + + tropics_profile = ( + data_surface["thetao"].sel( + time=time_slice, x=slice(118, 260), y=slice(-5, 5) + ) + - true_clim_to_remove.sel(x=slice(118, 260), y=slice(-5, 5)) + ).mean(["time", "y"]) + tropics_profile_pred = ( + self.pred_dict[self.key1]["ds_prediction_surface"]["thetao"][ + day_start : day_start + window + ].sel(lev=slice(0, 500), x=slice(118, 260), y=slice(-5, 5)) + - pred_clim_to_remove.sel(x=slice(118, 260), y=slice(-5, 5)) + ).mean(["time", "y"]) + + tropics_profile = tropics_profile.rename("Anomaly") + tropics_profile_pred = tropics_profile_pred.rename("Anomaly") + tropics_profile = tropics_profile.assign_attrs(units=r"${^oC}$") + tropics_profile_pred = tropics_profile_pred.assign_attrs(units=r"${^oC}$") + tropics_profile_pred["x"] = tropics_profile_pred["x"].assign_attrs( + units=r"${^o}$" + ) + + tropics_profile.plot.pcolormesh( + ax=axs["nina_true"], + y="lev", + cmap=cm.cm.diff, + vmin=-2, + vmax=2, + add_colorbar=False, + ) + axs["nina_true"].set_title("Nina Conditions GT") + axs["nina_true"].set_xlabel("") + axs["nina_true"].invert_yaxis() + tropics_profile_pred.plot.pcolormesh( + ax=axs["nina_pred"], + y="lev", + cmap=cm.cm.diff, + vmin=-2, + vmax=2, + cbar_ax=axs["colorbar"], + cbar_kwargs={ + "orientation": "horizontal", + "shrink": 0.3, + "extend": "both", + }, + ) + axs["nina_pred"].set_title( + f"Nina Conditions {self.pred_dict[self.key1]['name']}" + ) + axs["nina_pred"].invert_yaxis() + + ## Stats + for name, pred in zip( + [self.pred_dict[self.key1]["name"]], + [tropics_profile_pred], + ): + mae = np.abs( + (pred - tropics_profile).weighted(self.data["dz"]).mean(["x", "lev"]) + ) + cor = (pred * tropics_profile).weighted(self.data["dz"]).mean( + ["x", "lev"] + ) / np.sqrt( + (pred**2).weighted(self.data["dz"]).mean(["x", "lev"]) + * (tropics_profile**2).weighted(self.data["dz"]).mean(["x", "lev"]) + ) + + print( + name, + "La Nina mae: ", + mae.compute().item(), + "La Nina cor: ", + cor.compute().item(), + ) + + ############################################ + # Maps + ############################################ + # maybe define this centrally and use for all plots from this variable? + bound_east = 118 # i think these are outdated? + bound_west = 260 + bound_north = 5 + bound_south = -5 + # nino 3.4 box + # nino 3.4 box + nino_east = 190 + nino_west = 240 + + axs["map"].set_extent([70, 320, -25, 25], crs=ccrs.PlateCarree()) + axs["map"].stock_img() + axs["map"].coastlines(color="0.3", lw=0.5) + axs["map"].gridlines(draw_labels=True, color="0.4") + # We import here to avoid importing tk unconditionally above + from xarrayutils.plotting import box_plot # type: ignore + + box_plot( + [bound_east, bound_west, bound_south, bound_north], + ax=axs["map"], + color="orange", + transform=ccrs.PlateCarree(), + label="Full Profile", + ) + box_plot( + [nino_east, nino_west, bound_south, bound_north], + ax=axs["map"], + color="red", + ls="--", + transform=ccrs.PlateCarree(), + label="Niño 3.4", + ) + axs["map"].legend(bbox_to_anchor=[0.5, -0.5], loc="center", ncol=2) + + plt.savefig( + os.path.join(self.enso_path, "Nino_Figure_Short_with_map_single.png"), + bbox_inches="tight", + dpi=600, + ) + + def step_ohc_maps(self): + # OHC Map + Bias + Days_to_Eq = 0 + c_p = 3850 # J/(kg C) + rho_0 = 1025 # kg/m^3 + zeta_joules_factor = 1e21 # Conversion factor to ZJ + + plt.rcParams.update({"font.size": 14}) + fig, axs = plt.subplots( + 2, + 2, + figsize=(24, 9), + subplot_kw={"projection": ccrs.PlateCarree()}, + gridspec_kw={"wspace": 0.02, "hspace": 0.27}, + ) + axs = axs.flatten() + + def ohc_map(ohc_intz): + ohc_intz = remove_climatology(ohc_intz) + # return last 1 year - first 1 year + return ohc_intz.isel(time=slice(-73, None)).mean("time") - ohc_intz.isel( + time=slice(0, 73) + ).mean("time") + + # Define a common plotting function for Cartesian lat-lon grids + + def plot_ohc(ax, ohc_data, title, i): + # Configure colormap and set color for NaN values (land) + colormap = ( + cm.cm.balance + ) # cm.cm.thermal # Using thermal colormap from cmocean + colormap.set_bad(color=(0.7, 0.7, 0.7, 0)) + norm = symmetric_percentile_norm(ohc_data) + im = ax.pcolormesh( + ohc_data["x"], + ohc_data["y"], + ohc_data, + shading="auto", + cmap=colormap, + transform=ccrs.PlateCarree(), + norm=norm, + ) + ax.add_feature(cfeature.COASTLINE, edgecolor="black") + ax.set_title(title, fontsize=14) + # Set longitude and latitude labels + gl = ax.gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + + if i > 0: + gl.left_labels = False + return im + + def plot_diff_ohc(ax, ohc_data, gt_ohc_data, title, i): + # Configure colormap and set color for NaN values (land) + colormap = cm.cm.balance # Using thermal colormap from cmocean + colormap.set_bad(color=(0.7, 0.7, 0.7, 0)) + bias_ohc = ohc_data - gt_ohc_data + norm = symmetric_percentile_norm(bias_ohc) + im = ax.pcolormesh( + bias_ohc["x"], + bias_ohc["y"], + bias_ohc, + shading="auto", + cmap=colormap, + transform=ccrs.PlateCarree(), + norm=norm, + ) + ax.add_feature(cfeature.COASTLINE, edgecolor="black") + ax.set_title(title, fontsize=14) + # Set longitude and latitude labels + gl = ax.gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + + # if i > 4: + # gl.left_labels = False + return im + + # Calculate Ocean Heat Content for different scenarios and convert to Zeta Joules + titles = [self.dataset_name, self.pred_dict[self.key1]["name"]] + bias_titles = [self.pred_dict[self.key1]["name"] + " Bias"] + datasets = [self.data, self.pred_dict[self.key1]["ds_prediction"]] + + ohc_maps = [] + for ds in datasets: + section_mask = isnan(ds["thetao"]).all("lev").isel(time=5) + OHC_pred = ( + (ds["thetao"][Days_to_Eq:] * c_p * rho_0 / zeta_joules_factor) + .weighted(ds["areacello"] * ds["dz"]) + .sum(["lev"]) + .compute() + ) + OHC_pred = ohc_map(OHC_pred) + OHC_pred = OHC_pred.where(~section_mask) + OHC_pred = OHC_pred.rename("Ocean Heat Content") + OHC_pred["y"] = OHC_pred.y.assign_attrs( + long_name="latitude", units=r"${^o}$" + ) + OHC_pred["x"] = OHC_pred.x.assign_attrs( + long_name="longitude", units=r"${^o}$" + ) + OHC_pred = OHC_pred.assign_attrs(units="ZJ") + ohc_maps.append(OHC_pred) + + gt_ohc, pred1_ohc = ohc_maps + for i, (ax, title, ohc_data) in enumerate(zip(axs[:2], titles, ohc_maps)): + # Plot using the Cartesian lat-lon grid + im = plot_ohc(ax, ohc_data, title, i) + + # Add colorbar + cbar = fig.colorbar( + im, ax=axs[:2], orientation="vertical", fraction=0.01, pad=0.02 + ) + cbar.set_label("Ocean Heat Content [ZJ]", fontsize=14) + + im = plot_diff_ohc(axs[3], pred1_ohc, gt_ohc, bias_titles[0], 4) + + # Add colorbar + cbar = fig.colorbar( + im, ax=axs[3:], orientation="vertical", fraction=0.1, pad=0.02 + ) + cbar.set_label("Ocean Heat Content [ZJ]", fontsize=14) + + fig.delaxes(axs[2]) + + fig.suptitle( + "OHC Anomaly (Last Year - First Year) and Bias", fontsize=14, y=0.95 + ) + # Save or display the plot + plt.savefig( + os.path.join(self.ohc_path, "OHC_Global_map.png"), + bbox_inches="tight", + dpi=600, + ) + # plt.show() + + def map_bias_avg(data_pred1, fig, title="", **kwargs): + var_name = kwargs["var_name"] + + plt.clf() + plt.rcParams.update({"font.size": 14}) + + # Define colormap + new_cmap = cm.cm.balance + new_cmap.set_bad(color="grey", alpha=0.0) + + # Create figure with one subplot + ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree()) + + # Plot Prediction + norm = symmetric_percentile_norm(data_pred1) + im = data_pred1.plot( + ax=ax, + cmap=new_cmap, + norm=norm, + add_colorbar=False, + ) + ax.add_feature(cfeature.COASTLINE, edgecolor="black", linewidth=0.1) + ax.set_title(self.pred_dict[self.key1]["name"] + " Bias", fontsize=14) + + gl = ax.gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + + # Add colorbar + cbar = fig.colorbar( + im, ax=ax, orientation="vertical", fraction=0.02, pad=0.02 + ) + cbar.set_label(self.var_list[var_name]) + + # Add title + plt.text( + 0.5, + 1.1, + title, + ha="center", + va="bottom", + transform=ax.transAxes, + fontsize=14, + fontweight="bold", + ) + + return ax, im + + first_year = slice(None, 73) + last_year = slice(-73, -1) + second_last_year = slice(-146, -74) + third_last_year = slice(-219, -147) + + datasets = [self.data, self.pred_dict[self.key1]["ds_prediction"]] + + for i, (ax, title, ds) in enumerate(zip(axs, titles, datasets)): + section_mask = isnan(ds["thetao"]).all("lev") + OHC_pred = ( + (ds["thetao"][Days_to_Eq:] * c_p * rho_0 / zeta_joules_factor) + .weighted(ds["areacello"] * ds["dz"]) + .sum(["lev"]) + .compute() + ) + OHC_pred = remove_climatology(OHC_pred) + OHC_pred = OHC_pred.where(~section_mask) + OHC_pred = OHC_pred.rename("OHC Anomaly") + OHC_pred["y"] = OHC_pred.y.assign_attrs( + long_name="latitude", units=r"${^o}$" + ) + OHC_pred["x"] = OHC_pred.x.assign_attrs( + long_name="longitude", units=r"${^o}$" + ) + OHC_pred = OHC_pred.assign_attrs(units="ZJ") + + if i == 0: + gt_ohc = OHC_pred + elif i == 1: + pred1_ohc = OHC_pred + + da = ( + pred1_ohc.isel(time=last_year).mean("time") + - pred1_ohc.isel(time=second_last_year).mean("time") + ).compute() + fig, ax = plt.subplots(figsize=(10, 10)) + # increase size of plot + ax, im = map_bias_avg( + da, + fig, + var_name="OHC", + title="OHC Bias (Last Year - Second Last Year)", + ) + + # plot + fig.tight_layout() + plt.savefig( + os.path.join(self.ohc_path, "OHC_Bias_Map_Diff1_2.png"), + bbox_inches="tight", + dpi=600, + ) + + da = ( + pred1_ohc.isel(time=last_year).mean("time") + - pred1_ohc.isel(time=third_last_year).mean("time") + ).compute() + # increase size of plot + fig, ax = plt.subplots(figsize=(10, 10)) + ax, im = map_bias_avg( + da, + fig, + var_name="OHC", + title="OHC Bias (Last Year - Third Last Year)", + ) + + # plot + fig.tight_layout() + plt.savefig( + os.path.join(self.ohc_path, "OHC_Bias_Map_Diff1_3.png"), + bbox_inches="tight", + dpi=600, + ) + + da = ( + pred1_ohc.isel(time=last_year).mean("time") + - pred1_ohc.isel(time=first_year).mean("time") + ).compute() + # increase size of plot + fig, ax = plt.subplots(figsize=(10, 10)) + ax, im = map_bias_avg( + da, + fig, + var_name="OHC", + title="OHC Bias (Last Year - First Year)", + ) + + # plot + fig.tight_layout() + plt.savefig( + os.path.join(self.ohc_path, "OHC_Bias_Map_Diff_Last_First.png"), + bbox_inches="tight", + dpi=600, + ) + + def plot_surface_map(self, ax, data, title, i): + colormap = cm.cm.thermal + colormap.set_bad(color=(0.7, 0.7, 0.7, 0)) + norm = percentile_norm(data) + im = ax.pcolormesh( + data["x"], + data["y"], + data, + shading="auto", + cmap=colormap, + transform=ccrs.PlateCarree(), + norm=norm, + ) + ax.add_feature(cfeature.COASTLINE, edgecolor="black") + ax.set_title(title, fontsize=14) + gl = ax.gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + + if i > 0: + gl.left_labels = False + return im + + def plot_bias(self, ax, data, gt_data, title): + colormap = cm.cm.balance + colormap.set_bad(color=(0.7, 0.7, 0.7, 0)) + bias = data - gt_data + norm = symmetric_percentile_norm(bias) + im = ax.pcolormesh( + bias["x"], + bias["y"], + bias, + shading="auto", + cmap=colormap, + transform=ccrs.PlateCarree(), + norm=norm, + ) + ax.add_feature(cfeature.COASTLINE, edgecolor="black") + ax.set_title(title, fontsize=14) + gl = ax.gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + + return im + + def step_sst_mean_maps(self): + plt.rcParams.update({"font.size": 14}) + fig, axs = plt.subplots( + 2, + 2, + figsize=(16, 6), + subplot_kw={"projection": ccrs.PlateCarree()}, + gridspec_kw={"wspace": 0.02, "hspace": 0.23}, + ) + axs = axs.flatten() + + # Calculate Sea Surface Temperature (SST) for different scenarios + titles = [self.dataset_name, self.pred_dict[self.key1]["name"]] + bias_titles = [self.pred_dict[self.key1]["name"] + " Bias"] + datasets = [self.data, self.pred_dict[self.key1]["ds_prediction"]] + + for i, (ax, title, ds) in enumerate(zip(axs, titles, datasets)): + section_mask = isnan(ds["thetao"]).isel(lev=0).isel(time=5) + SST_pred = ds["thetao"].isel(lev=0).mean("time") + SST_pred = SST_pred.where(~section_mask) + SST_pred = SST_pred.rename("2.5m " + r"$\theta_O$") + SST_pred["y"] = SST_pred.y.assign_attrs( + long_name="latitude", units=r"${^o}$" + ) + SST_pred["x"] = SST_pred.x.assign_attrs( + long_name="longitude", units=r"${^o}$" + ) + SST_pred = SST_pred.assign_attrs(units=r"$\degree C$") + + if i == 0: + gt_sst = SST_pred + elif i == 1: + pred1_sst = SST_pred + + # Plot using the Cartesian lat-lon grid + im = self.plot_surface_map(ax, SST_pred, title, i) + + # Add colorbar for SST plots + cbar = fig.colorbar( + im, ax=axs[:2], orientation="vertical", fraction=0.01, pad=0.02 + ) + cbar.set_label(r"$\theta_O$ [$\degree C$]", fontsize=14) + + # Plot biases for SST + im = self.plot_bias(axs[3], pred1_sst, gt_sst, bias_titles[0]) + + # Add colorbar for bias plots + cbar = fig.colorbar( + im, ax=axs[3:], orientation="vertical", fraction=0.1, pad=0.02 + ) + cbar.set_label(r"$\theta_O$ [$\degree C$]", fontsize=14) + + # Remove the empty axis + fig.delaxes(axs[2]) + + fig.suptitle("2.5m Potential Temperature", fontsize=14, y=1.0) + + # Save or display the plot + plt.savefig( + os.path.join(self.temp_path, "SST_Global_map.png"), + bbox_inches="tight", + dpi=600, + ) + # plt.show() + + def step_sst_time_snapshot_maps(self): + for t_index in self.time_indices: + plt.rcParams.update({"font.size": 14}) + fig, axs = plt.subplots( + 2, + 2, + figsize=(16, 6), + subplot_kw={"projection": ccrs.PlateCarree()}, + gridspec_kw={"wspace": 0.02, "hspace": 0.23}, + ) + axs = axs.flatten() + + # Define a common plotting function for Cartesian lat-lon grids + + titles = [ + self.dataset_name + f" t={t_index}", + self.pred_dict[self.key1]["name"] + f" t={t_index}", + ] + bias_titles = [self.pred_dict[self.key1]["name"] + " Bias"] + datasets = [ + self.data, + self.pred_dict[self.key1]["ds_prediction"], + ] + + for i, (ax, title, ds) in enumerate(zip(axs, titles, datasets)): + section_mask = isnan(ds["thetao"]).isel(lev=0).isel(time=5) + SST_pred = ds["thetao"].isel(lev=0).isel(time=t_index) + SST_pred = SST_pred.where(~section_mask) + SST_pred = SST_pred.rename("2.5m " + r"$\theta_O$") + SST_pred["y"] = SST_pred.y.assign_attrs( + long_name="latitude", units=r"${^o}$" + ) + SST_pred["x"] = SST_pred.x.assign_attrs( + long_name="longitude", units=r"${^o}$" + ) + SST_pred = SST_pred.assign_attrs(units=r"$\degree C$") + + if i == 0: + gt_sst = SST_pred + elif i == 1: + pred1_sst = SST_pred + + # Plot using the Cartesian lat-lon grid + im = self.plot_surface_map(ax, SST_pred, title, i) + + # Add colorbar for SST plots + cbar = fig.colorbar( + im, ax=axs[:2], orientation="vertical", fraction=0.01, pad=0.02 + ) + cbar.set_label(r"$\theta_O$ [$\degree C$]", fontsize=14) + + # Plot biases for SST + im = self.plot_bias(axs[3], pred1_sst, gt_sst, bias_titles[0]) + + # Add colorbar for bias plots + cbar = fig.colorbar( + im, ax=axs[3:], orientation="vertical", fraction=0.1, pad=0.02 + ) + cbar.set_label(r"$\theta_O$ [$\degree C$]", fontsize=14) + + # Remove the empty axis + fig.delaxes(axs[2]) + + # Save or display the plot + plt.savefig( + os.path.join(self.temp_path, f"SST_map_snapshot_t_{t_index}.png"), + bbox_inches="tight", + dpi=600, + ) + # plt.show() + plt.close() + + # #### Salinity Map + + def step_salinity_mean_map(self): + plt.rcParams.update({"font.size": 14}) + fig, axs = plt.subplots( + 2, + 2, + figsize=(16, 6), + subplot_kw={"projection": ccrs.PlateCarree()}, + gridspec_kw={"wspace": 0.02, "hspace": 0.23}, + ) + axs = axs.flatten() + + # Define a common plotting function for Cartesian lat-lon grids + + # Calculate Sea Surface Salinity (SSS) for different scenarios + titles = [self.dataset_name, self.pred_dict[self.key1]["name"]] + bias_titles = [self.pred_dict[self.key1]["name"] + " Bias"] + datasets = [self.data, self.pred_dict[self.key1]["ds_prediction"]] + + for i, (ax, title, ds) in enumerate(zip(axs, titles, datasets)): + section_mask = isnan(ds["so"]).isel(lev=0).isel(time=5) + SSS_pred = ds["so"].isel(lev=0).mean("time") + SSS_pred = SSS_pred.where(~section_mask) + SSS_pred = SSS_pred.rename("2.5m " + r"$so$") + SSS_pred["y"] = SSS_pred.y.assign_attrs( + long_name="latitude", units=r"${^o}$" + ) + SSS_pred["x"] = SSS_pred.x.assign_attrs( + long_name="longitude", units=r"${^o}$" + ) + SSS_pred = SSS_pred.assign_attrs(units=r"$psu$") + + if i == 0: + gt_sss = SSS_pred + elif i == 1: + pred1_sss = SSS_pred + + # Plot using the Cartesian lat-lon grid + im = self.plot_surface_map(ax, SSS_pred, title, i) + + # Add colorbar for SSS plots + cbar = fig.colorbar( + im, ax=axs[:2], orientation="vertical", fraction=0.01, pad=0.02 + ) + cbar.set_label(r"$so$ [$psu$]", fontsize=14) + + # Plot biases for SSS + im = self.plot_bias(axs[3], pred1_sss, gt_sss, bias_titles[0]) + + # Add colorbar for bias plots + cbar = fig.colorbar( + im, ax=axs[3:], orientation="vertical", fraction=0.1, pad=0.02 + ) + cbar.set_label(r"$so$ [$psu$]", fontsize=14) + + # Remove the empty axis + fig.delaxes(axs[2]) + + # Save or display the plot + plt.savefig( + os.path.join(self.salinity_path, "SeaSurfaceSalinity_Global_map.png"), + bbox_inches="tight", + dpi=600, + ) + # plt.show() + + def step_salinity_snapshot_maps(self): + for t_index in self.time_indices: + plt.rcParams.update({"font.size": 14}) + fig, axs = plt.subplots( + 2, + 2, + figsize=(16, 6), + subplot_kw={"projection": ccrs.PlateCarree()}, + gridspec_kw={"wspace": 0.02, "hspace": 0.23}, + ) + axs = axs.flatten() + + # Calculate Sea Surface Salinity (SSS) for different scenarios + titles = [ + self.dataset_name + f" t={t_index}", + self.pred_dict[self.key1]["name"] + f" t={t_index}", + ] + bias_titles = [self.pred_dict[self.key1]["name"] + " Bias"] + datasets = [ + self.data, + self.pred_dict[self.key1]["ds_prediction"], + ] + + for i, (ax, title, ds) in enumerate(zip(axs, titles, datasets)): + section_mask = isnan(ds["so"]).isel(lev=0).isel(time=5) + SSS_pred = ds["so"].isel(lev=0).isel(time=t_index) + SSS_pred = SSS_pred.where(~section_mask) + SSS_pred = SSS_pred.rename("2.5m " + r"$so$") + SSS_pred["y"] = SSS_pred.y.assign_attrs( + long_name="latitude", units=r"${^o}$" + ) + SSS_pred["x"] = SSS_pred.x.assign_attrs( + long_name="longitude", units=r"${^o}$" + ) + SSS_pred = SSS_pred.assign_attrs(units=r"$psu$") + + if i == 0: + gt_sss = SSS_pred + elif i == 1: + pred1_sss = SSS_pred + + # Plot using the Cartesian lat-lon grid + im = self.plot_surface_map(ax, SSS_pred, title, i) + + # Add colorbar for SSS plots + cbar = fig.colorbar( + im, ax=axs[:2], orientation="vertical", fraction=0.01, pad=0.02 + ) + cbar.set_label(r"$so$ [$psu$]", fontsize=14) + + # Plot biases for SSS + im = self.plot_bias(axs[3], pred1_sss, gt_sss, bias_titles[0]) + + # Add colorbar for bias plots + cbar = fig.colorbar( + im, ax=axs[3:], orientation="vertical", fraction=0.1, pad=0.02 + ) + cbar.set_label(r"$so$ [$psu$]", fontsize=14) + + # Remove the empty axis + fig.delaxes(axs[2]) + + # Save or display the plot + plt.savefig( + os.path.join(self.salinity_path, f"SSS_map_snapshot_t_{t_index}.png"), + bbox_inches="tight", + dpi=600, + ) + # plt.show() + + # ### Movies + + # Need atleast two keys otherwise duplicate maps + + def step_movies(self): + keys = list(self.pred_dict.keys()) + # assert len(keys) >= 2, "Maps supported by atleast two keys" + self.key1 = keys[0] + if len(keys) == 1: + print( + "Maps only support two models for now!!! Using the same key for both maps" + ) + key2 = keys[0] + elif len(keys) > 2: + print("Maps only support two models for now!!! Using the first two keys") + assert False + else: + key2 = keys[1] + + # #### Core + + def _core_plot(ax, data, plotmethod=None, **kwargs): + """Core plotting functionality.""" + # Deactivate cbar for contours (not sure this should be hardcoded...) + if plotmethod == "contour": + kwargs.pop("cbar_kwargs", None) + + # I am probably recoding something from matplotlib...is there a way to get + # the plot.something functionslity with a keyword? + # For now do it the hard way + if plotmethod is None: + p = data.plot(ax=ax, **kwargs) + # doesnt work,...i want this for smoother images + elif plotmethod == "imshow": + # p = data.plot.imshow(ax=ax, **kwargs) + # testing interpolation + p = data.plot.imshow(ax=ax, interpolation="gaussian", **kwargs) + # print(p.get_interpolation()) + elif plotmethod == "pcolormesh": + p = data.plot.pcolormesh(ax=ax, **kwargs) + elif plotmethod == "contour": + p = data.plot.contour(ax=ax, **kwargs) + elif plotmethod == "contourf": + p = data.plot.contourf(ax=ax, **kwargs) + else: + raise RuntimeError( + f"Input '{plotmethod}' not recognized \ + as plotmode" + ) + return p + + def _base_plot( + ax, + base_data, + timestamp, + timestep_value, + framedim, + plotmethod=None, + **kwargs, + ): + data = base_data.isel({framedim: timestamp}) + p = _core_plot(ax, data, plotmethod=plotmethod, **kwargs) + return p + + ### Presets (should proabably put all others into a submodule) + def basic( + da, + fig, + timestamp, + timestep_value, + framedim="time", + plotmethod=None, + subplot_kw=None, + **kwargs, + ): + # create axis + ax = fig.subplots(subplot_kw=subplot_kw) + pp = _base_plot( + ax, + da, + timestamp, + timestep_value, + framedim, + plotmethod=plotmethod, + **kwargs, + ) + return ax, pp + + # data treatment + + def _parse_plot_defaults(da, kwargs): + if isinstance(da, xr.DataArray): + data = da + else: + raise RuntimeError(f"input of type ({type(da)}) not supported yet.") + + # check these explicitly to avoid any computation if these are set. + if "vmin" not in kwargs.keys(): + warnings.warn( + "No `vmin` provided. data limits are calculated from input. " + "Depending on the input this can take long. Pass `vmin` to avoid " + "this step", + UserWarning, + ) + kwargs["vmin"] = data.min().data + + if "vmax" not in kwargs.keys(): + warnings.warn( + "No `vmax` provided. data limits are calculated from input. " + "Depending on the input this can take long. Pass `vmax` to avoid " + "this step", + UserWarning, + ) + kwargs["vmax"] = data.max().data + + # There is a bug that prevents this from working... + # I'll have to fix that upstream. + # defaults["cbar_kwargs"] = dict(extend="neither") + # This works for now + kwargs.setdefault("extend", "neither") + + # if any value is dask.array compute them here. + for k in ["vmin", "vmax"]: + if isinstance(kwargs[k], DaskArray): + kwargs[k] = kwargs[k].compute() + + return kwargs + + def _check_plotfunc_output(func, da, framedim="time", **kwargs): + timestep = 0 + timestep_value = da[framedim].data[timestep] + fig = plt.figure() + oargs = func(da, fig, timestep, timestep_value, framedim, **kwargs) + # I just want the number of output args, delete plot + plt.close(fig) + if oargs is None: + return 0 + else: + return len(oargs) + + def _check_ffmpeg_version(): + p = Popen("ffmpeg -version", stdout=PIPE, shell=True) + (output, err) = p.communicate() + p_status = p.wait() + # Parse version + if p_status != 0: + print("No ffmpeg found") + return None + else: + # parse version number + if copyright := re.search( + "ffmpeg version (.+?) Copyright", str(output) + ): + return copyright.group(1).replace(" ", "") + return None + + def _execute_command( + command, verbose=False, error=True, log_file="output.log", max_lines=10 + ): + with open(log_file, "w") as f: + p = Popen( + command, + stdout=PIPE, + stderr=STDOUT, + shell=True, + bufsize=1, + universal_newlines=True, + ) + line_count = 0 + assert p.stdout is not None + + for line in iter(p.stdout.readline, ""): + f.write(line) # Write to log file + + if verbose and line_count < max_lines: + sys.stdout.write(line) # Display line in console + sys.stdout.flush() + line_count += 1 + + p.stdout.close() + p.wait() + + # Inform if output was truncated + if verbose and line_count >= max_lines: + print(f"\n...Output truncated. Full log saved to {log_file}") + + if error and p.returncode != 0: + raise RuntimeError( + f"Command '{command}' failed with return code {p.returncode}" + ) + + return p + + def _check_ffmpeg_execute(command, verbose=False): + if _check_ffmpeg_version() is None: + raise RuntimeError( + "Could not find an ffmpeg version on the system. \ + Please install ffmpeg with e.g. `conda install -c conda-forge ffmpeg`" + ) + else: + try: + p = _execute_command(command, verbose=verbose) + return p + except RuntimeError: + raise RuntimeError( + "Something has gone wrong. Use `verbose=True` to check " + "if ffmpeg displays a problem" + ) + + def convert_gif( + mpath, + gpath="movie.gif", + gif_palette=False, + resolution=(480, 320), + verbose=False, + remove_movie=True, + gif_framerate=5, + ): + if gif_palette: + palette_filter = ( + '-filter_complex "[0:v] split [a][b];[a] palettegen [p];' + '[b][p] paletteuse"' + ) + else: + palette_filter = "" + + command = ( + f'ffmpeg -y -i "{mpath}" {palette_filter} -r {gif_framerate}' + f' -s {resolution[0]}x{resolution[1]} "{gpath}"' + ) + p = _check_ffmpeg_execute(command, verbose=verbose) + + print(f"GIF created at {gpath}") + if remove_movie: + if os.path.exists(mpath): + os.remove(mpath) + return p + + def _combine_ffmpeg_command( + sourcefolder, moviename, framerate, frame_pattern, ffmpeg_options + ): + # we need `-y` because i can not properly diagnose the errors here... + command = ( + f'ffmpeg -r {framerate} -i "{os.path.join(sourcefolder, frame_pattern)}"' + f" -y {ffmpeg_options} -r {framerate}" + f' "{os.path.join(sourcefolder, moviename)}"' + ) + return command + + def write_movie( + sourcefolder, + moviename, + frame_pattern="frame_%05d.png", + remove_frames=True, + verbose=False, + ffmpeg_options="-c:v libvpx-vp9 -b:v 2M -f mp4", + framerate=20, + ): + command = _combine_ffmpeg_command( + sourcefolder, moviename, framerate, frame_pattern, ffmpeg_options + ) + p = _check_ffmpeg_execute(command, verbose=verbose) + + print(f"Movie created at {moviename}") + if remove_frames: + rem_name = frame_pattern.replace("%05d", "*") + for f in glob.glob(os.path.join(sourcefolder, rem_name)): + if os.path.exists(f): + os.remove(f) + return p + + def frame_save(fig, frame, odir=None, frame_pattern="frame_%05d.png", dpi=100): + fig.savefig( + os.path.join(odir, frame_pattern % (frame)), + dpi=dpi, + facecolor=fig.get_facecolor(), + transparent=True, + ) + # I am trying everything to *wipe* this figure, hoping that it could + # help with the dask glitches I experienced earlier. + # TBD if this is all needed...how this might affect performance. + plt.close(fig) + del fig + gc.collect(2) + + class Movie: + def __init__( + self, + da, + plotfunc=None, + framedim="time", + pixelwidth=1920, + pixelheight=1080, + dpi=200, + frame_pattern="frame_%05d.png", + fieldname=None, + input_check=True, + **kwargs, + ): + self.pixelwidth = pixelwidth + self.pixelheight = pixelheight + self.dpi = dpi + self.width = self.pixelwidth / self.dpi + self.height = self.pixelheight / self.dpi + self.frame_pattern = frame_pattern + self.data = da + self.framedim = framedim + if plotfunc is None: + self.plotfunc = basic + else: + self.plotfunc = plotfunc + # set sensible defaults + self.raw_kwargs = kwargs + + # Check input + + # optional checks (these might need to be deactivated when using custom + # plot functions.) + if input_check: + if isinstance(self.data, xr.Dataset): + raise ValueError( + "xmovie presets do not yet support the input of xr.Datasets. \ + In order to use datasets as inputs, set `input_check` to False. \ + Note that this requires you to manually set colorlimits etc." + ) + + # Set defaults + self.kwargs = _parse_plot_defaults(self.data, self.raw_kwargs) + else: + self.kwargs = self.raw_kwargs + + # Mandatory checks + # Check if `framedim` exists. + if self.framedim not in list(self.data.dims): + raise ValueError( + f"Framedim ({self.framedim}) not found in input data" + ) + # Check the output of plotfunc + self.plotfunc_n_outargs = _check_plotfunc_output( + self.plotfunc, self.data, self.framedim, **self.kwargs + ) + + def render_frame(self, timestep, timestep_value): + """Renders complete figure (frame) for given timestep. + + Parameters + ---------- + timestep : type + Description of parameter `timestep`. + + Returns: + ------- + type + Description of returned object. + + """ + fig = plt.figure(figsize=(self.width, self.height)) + # create_frame(self.pixelwidth, self.pixelheight, self.dpi) + # produce dummy output for ax and pp if the plotfunc does not provide them + if self.plotfunc_n_outargs == 2: + # this should be the case for all presets provided by xmovie + ax, pp = self.plotfunc( + self.data, + fig, + timestep, + timestep_value, + self.framedim, + **self.kwargs, + ) + else: + warnings.warn( + "The provided `plotfunc` does not provide the expected number of " + "output arguments. Expected a function `ax,pp =plotfunc(...)` but " + f"got {self.plotfunc_n_outargs} output arguments. Inserting dummy " + "values. This should not affect output.", + UserWarning, + ) + _ = self.plotfunc( + self.data, + fig, + timestep, + timestep_value, + self.framedim, + **self.kwargs, + ) + ax, pp = None, None + return fig, ax, pp + + def save_frames(self, odir, progress=False): + """Save movie frames as picture files. + + Parameters + ---------- + odir : path + path to output directory + progress : type + Show progress bar. Requires + + """ + # create range of frames + timesteps = self.data[self.framedim].data + frame_range: Iterable[int] = range(len(timesteps)) + if progress: + frame_range = tqdm(frame_range) + + for fi in frame_range: + fig, ax, pp = self.render_frame(fi, timesteps[fi]) + frame_save( + fig, + fi, + odir=odir, + frame_pattern=self.frame_pattern, + dpi=self.dpi, + ) + + # Needs more testing! Slower than the sequential counterpart + def save_frames_parallel(self, odir, batch_size=10, progress=False): + """Save frames in parallel batches to reduce I/O overhead.""" + timesteps = self.data[self.framedim].data + frame_range = range(len(timesteps)) + + for i in tqdm(range(0, len(frame_range), batch_size)): + batch = frame_range[i : i + batch_size] + tasks = [] + for fi in batch: + delayed_frame = delayed(self.render_frame)(fi, timesteps[fi]) + tasks.append( + delayed(frame_save)( + delayed_frame[0], + fi, + odir=odir, + frame_pattern=self.frame_pattern, + dpi=self.dpi, + ) + ) + compute(*tasks) + + def save( + self, + filename, + remove_frames=True, + remove_movie=True, + progress=False, + verbose=False, + overwrite_existing=False, + framerate=15, + ffmpeg_options="-c:v mjpeg -q:v 2 -pix_fmt yuvj420p", + gif_palette=False, + gif_resolution_factor=0.5, + gif_framerate=10, + parallel=False, + batch_size=10, + ): + """Save out animation from Movie object. + + Parameters + ---------- + filename : str + Pathname to final movie/animation. Output is dependent on filetype: + Creates movie for `*.mp4` and gif for `*.gif` + remove_frames : Bool + Optional removal of frame pictures (the default is True; False will + leave all picture files in folder). + remove_movie : Bool + As `remove_frames` but for movie file. Only applies when filename + is given as `.gif` (the default is True). + progress : Bool + Experimental switch to show progress output. This will be refined + in future version (the default is False). + verbose : Bool + Experimental switch to show output of ffmpeg commands. Useful for + debugging but can quickly flood your notebook + (the default is False). + overwrite_existing : Bool + Set to overwrite existing files with `filename` + (the default is False). + framerate : int + Frames per second for the output movie file. Only relevant for + '.mp4' files. + (the default is 15) + ffmpeg_options: str + Encoding options to pass to ffmpeg call. + Defaults to : `"-c:v libx264 -preset veryslow -crf 10 -pix_fmt yuv420p"` + gif_palette : Bool + Use a gif colorpalette to improve quality. Can lead to artifacts + in very contrasty situations (the default is False). + gif_resolution_factor : float + Factor used to reduce gif resolution compared to movie. + Use 1.0 to put out the same resolutions for both products. + (the default is 0.5). + gif_framerate : int + As `framerate` but for the gif output file. Only relevant to + `.gif` files. + (the default is 10) + """ + # parse out directory and filename + dirname = os.path.dirname(filename) + filename = os.path.basename(filename) + + # detect gif filename + + isgif = ".gif" in filename + if isgif: + giffile = filename + moviefile = filename.replace("gif", "mp4") + gpath = os.path.join(dirname, giffile) + else: + moviefile = filename + + mpath = os.path.join(dirname, moviefile) + + # check existing files + if os.path.exists(mpath): + if not overwrite_existing: + raise RuntimeError( + f"File `{mpath}` already exists. Set `overwrite_existing` " + "to True to overwrite." + ) + if isgif: + if os.path.exists(gpath): + if not overwrite_existing: + raise RuntimeError( + f"File `{gpath}` already exists. Set " + "`overwrite_existing` to True to overwrite." + ) + + # print frames + if parallel: + self.save_frames_parallel( + dirname, batch_size=batch_size, progress=progress + ) + else: + self.save_frames(dirname, progress=progress) + + # Create movie + write_movie( + dirname, + moviefile, + frame_pattern=self.frame_pattern, + remove_frames=remove_frames, + verbose=verbose, + framerate=framerate, + ffmpeg_options=ffmpeg_options, + ) + + # Create gif + if isgif: + # if ppath: + # create_gif_palette(mpath, ppath=ppath, verbose=verbose) + convert_gif( + mpath, + gpath=gpath, + resolution=[480, 320], + gif_palette=gif_palette, + verbose=verbose, + remove_movie=remove_movie, + gif_framerate=gif_framerate, + ) + + # #### Global Map + + def global_surface_map( + da, fig, timestamp, timestamp_val, framedim="time", **kwargs + ): + var_name = kwargs["var_name"] + + da_gt = da[0] + da_pred1 = da[1] + da_pred2 = da[2] + + data_gt = da_gt.isel({framedim: timestamp}) + data_pred1 = da_pred1.isel({framedim: timestamp}) + data_pred2 = da_pred2.isel({framedim: timestamp}) + + plt.clf() + plt.rcParams.update({"font.size": 14}) + + # Define colormap + new_cmap = ( + cm.cm.balance + if var_name == "thetao" or var_name == "OHC" + else cm.cm.haline + ) + new_cmap.set_bad(color="grey", alpha=0.0) + # Set common color range for the colorbar + vmin, vmax = { + "thetao": (0, 30), + "so": (30, 40), + "uo": (-2, 2), + "vo": (-2, 2), + "zos": (-1, 1), + "OHC": (-0.05, 0.05), + }[var_name] + + # Create figure with appropriate layout + ax = fig.subplots( + 1, + 3, + subplot_kw={"projection": ccrs.PlateCarree()}, + gridspec_kw={"wspace": 0.02, "hspace": 0.05}, + ) + ax = np.array(ax) # Ensure ax is an array for easy indexing + + # Plot Ground Truth (GT) + im = data_gt.plot( + ax=ax[0], + cmap=new_cmap, + vmin=vmin, + vmax=vmax, + add_colorbar=False, + ) + ax[0].add_feature(cfeature.COASTLINE, edgecolor="black", linewidth=0.1) + ax[0].set_title(f"{self.dataset_name}", fontsize=14) + gl = ax[0].gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + + # Plot Predictions + im = data_pred1.plot( + ax=ax[1], + cmap=new_cmap, + vmin=vmin, + vmax=vmax, + add_colorbar=False, + ) + ax[1].add_feature(cfeature.COASTLINE, edgecolor="black", linewidth=0.1) + ax[1].set_title(self.pred_dict[self.key1]["name"], fontsize=14) + gl = ax[1].gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.left_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + ax[1].set_yticks([]) + ax[1].set_ylabel("") + + im = data_pred2.plot( + ax=ax[2], + cmap=new_cmap, + vmin=vmin, + vmax=vmax, + add_colorbar=False, + ) + ax[2].add_feature(cfeature.COASTLINE, edgecolor="black", linewidth=0.1) + ax[2].set_title(self.pred_dict[key2]["name"], fontsize=14) + gl = ax[2].gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.left_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + ax[2].set_yticks([]) + ax[2].set_ylabel("") + + # Add colorbar for plots + cbar = fig.colorbar( + im, ax=ax[:], orientation="vertical", fraction=0.02, pad=0.02 + ) + cbar.set_label(self.var_list[var_name]) + + # Add timestamp text + ax[0].text( + 1.0, + 1.4, + f"{timestamp_val.year}-{timestamp_val.month:02d}-{timestamp_val.day:02d}", + horizontalalignment="center", + verticalalignment="center", + transform=ax[0].transAxes, + fontsize=14, + fontweight="bold", + ) + + return ax, im + + def global_surface_map_bias( + da, fig, timestamp, timestamp_val, framedim="time", **kwargs + ): + var_name = kwargs["var_name"] + + da_pred1 = da[0] + da_pred2 = da[1] + + data_pred1 = da_pred1.isel({framedim: timestamp}) + data_pred2 = da_pred2.isel({framedim: timestamp}) + + plt.clf() + plt.rcParams.update({"font.size": 14}) + + # Define colormap + new_cmap = cm.cm.balance + new_cmap.set_bad(color="grey", alpha=0.0) + + # Set common color range for the colorbar + vmin, vmax = { + "thetao": (-5, 5), + "so": (-1, 1), + "uo": (-0.01, 0.01), + "vo": (-0.01, 0.01), + "zos": (-1, 1), + "OHC": (-0.1, 0.1), + }[var_name] + + # Create figure with appropriate layout + ax = fig.subplots( + 1, + 2, + subplot_kw={"projection": ccrs.PlateCarree()}, + gridspec_kw={"wspace": 0.02, "hspace": 0.05}, + ) + ax = np.array(ax) # Ensure ax is an array for easy indexing + + # Plot Predictions + im = data_pred1.plot( + ax=ax[0], + cmap=new_cmap, + vmin=vmin, + vmax=vmax, + add_colorbar=False, + ) + ax[0].add_feature(cfeature.COASTLINE, edgecolor="black", linewidth=0.1) + ax[0].set_title(self.pred_dict[self.key1]["name"] + " Bias", fontsize=14) + gl = ax[0].gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.left_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + + im = data_pred2.plot( + ax=ax[1], + cmap=new_cmap, + vmin=vmin, + vmax=vmax, + add_colorbar=False, + ) + ax[1].add_feature(cfeature.COASTLINE, edgecolor="black", linewidth=0.1) + ax[1].set_title(self.pred_dict[key2]["name"] + " Bias", fontsize=14) + gl = ax[1].gridlines(draw_labels=True, color="0.4", linestyle="--", alpha=0) + gl.top_labels = False + gl.right_labels = False + gl.left_labels = False + gl.xlabel_style = {"size": 14} + gl.ylabel_style = {"size": 14} + gl.xlocator = FixedLocator([-120, -60, 0, 60, 120]) + ax[1].set_yticks([]) + ax[1].set_ylabel("") + + # Add colorbar for plots + cbar = fig.colorbar( + im, ax=ax[:], orientation="vertical", fraction=0.02, pad=0.02 + ) + cbar.set_label(self.var_list[var_name]) + + # Add timestamp text + ax[0].text( + 1.0, + 1.4, + f"{timestamp_val.year}-{timestamp_val.month:02d}-{timestamp_val.day:02d}", + horizontalalignment="center", + verticalalignment="center", + transform=ax[0].transAxes, + fontsize=14, + fontweight="bold", + ) + + return ax, im + + if "mask" in self.data.data_vars: + mask = self.data.mask + surface_mask = mask.isel( + time=0, lev=0, missing_dims="ignore" + ) # TODO: we don't have time anymore + else: + mask = self.data.wetmask + surface_mask = mask.isel(lev=0) + + # combine the two datasets into a single xarray new dimension + movie_var_list = ["thetao", "so"] + + c_p = 3850 # J/(kg C) + rho_0 = 1025 # kg/m^3 + zeta_joules_factor = 1e21 + + for var in movie_var_list: + if var == "OHC": + ohc_gt = ( + (self.data["thetao"] * c_p * rho_0 / zeta_joules_factor) + .weighted(self.data["areacello"] * self.data["dz"]) + .sum(["lev"]) + .compute() + ) + data_gt = remove_climatology(ohc_gt) + + ohc_pred1 = ( + ( + self.pred_dict[self.key1]["ds_prediction"]["thetao"] + * c_p + * rho_0 + / zeta_joules_factor + ) + .weighted(self.data["areacello"] * self.data["dz"]) + .sum(["lev"]) + .compute() + ) + data_pred1 = remove_climatology(ohc_pred1) + + ohc_pred2 = ( + ( + self.pred_dict[key2]["ds_prediction"]["thetao"] + * c_p + * rho_0 + / zeta_joules_factor + ) + .weighted(self.data["areacello"] * self.data["dz"]) + .sum(["lev"]) + .compute() + ) + data_pred2 = remove_climatology(ohc_pred2) + + else: + data_gt = self.data[var].isel(lev=0).compute() + data_pred1 = ( + self.pred_dict[self.key1]["ds_prediction"][var] + .isel(lev=0) + .compute() + ) + data_pred2 = ( + self.pred_dict[key2]["ds_prediction"][var].isel(lev=0).compute() + ) + + data_gt = data_gt.where(surface_mask).compute() + data_pred1 = data_pred1.where(surface_mask).compute() + data_pred2 = data_pred2.where(surface_mask).compute() + + da = xr.concat([data_gt, data_pred1, data_pred2], dim="dummy") + mov = Movie( + da, plotfunc=global_surface_map, var_name=var, input_check=False + ) + mov.save( + os.path.join(self.movie_path, f"{var}_surface_map_movie.mp4"), + progress=True, + overwrite_existing=True, + ) + + # Bias + da = xr.concat( + [(data_pred1 - data_gt).compute(), (data_pred2 - data_gt).compute()], + dim="dummy", + ) + mov = Movie( + da, plotfunc=global_surface_map_bias, var_name=var, input_check=False + ) + mov.save( + os.path.join(self.movie_path, f"{var}_surface_map_bias_movie.mp4"), + progress=True, + overwrite_existing=True, + ) + + # Observation comparison + # + # These recompute through `samudra.metrics` rather than reading anything + # `samudra.eval` wrote, so viz runs against any rollout on its own. Sharing + # the kernels is what keeps the numbers on the figures equal to the numbers + # in W&B. + # ------------------------------------------------------------------ + + def _obs_inputs(self): + """Observation products, rollouts on the observation grid, and the frame. + + The same pass `samudra.eval` runs, so the CSV beside these figures is the + one the eval job would have written and the numbers annotated on them are + the ones it would have logged. + + Computed once and cached: it costs a full reduction over the record, and + every observation step wants the same one. + """ + if getattr(self, "_obs_cache", None) is not None: + return self._obs_cache + if self.observations is None or self.obs_data_root is None: + raise ValueError( + "This viz config has no `observations` block, so there is " + "nothing to compare against. Add one (see " + "src/samudra/configs/data/obs.yaml)." + ) + + # Copied because `process_data` writes stacked variables back into these + # same objects during construction; scoring what was read keeps this + # independent of whatever the other steps have done to them. + runs = {run.name: run.data.copy() for run in self.raw_runs} + if self.observations.baselines: + # `baselines` adds the ground-truth comparison to an eval job, which + # scores whatever it was handed. viz draws the runs its config + # lists, so the key has no effect here; saying so beats a figure + # quietly missing a panel the CSV beside it would have had. + logger.info( + "Ignoring observations.baselines=%s: viz draws the runs in its " + "own `runs:` list. Add the baseline there to see it.", + self.observations.baselines, + ) + scored = score_rollouts( + self.observations, + rollouts=runs, + dataset_spec=self.dataset_spec, + data_root=self.obs_data_root, + primary_label=next(iter(runs)), + output_dir=Path(self.obs_path), + ) + self._obs_cache = (scored.rollouts, scored.products, scored.frame) + return self._obs_cache + + def _observations_configured(self, step: str) -> bool: + """Whether there is anything for the observation steps to draw against. + + The observation steps are part of the default run, so a preset without + an `observations` block has to skip them rather than fail: most presets + do not configure one, and `VizConfig.observations` says omitting it + skips these steps. + """ + if self.observations is None or self.obs_data_root is None: + logger.info( + "Skipping %s: no `observations` block in this viz config " + "(see src/samudra/configs/data/obs.yaml).", + step, + ) + return False + return True + + def step_obs_rmse_maps(self): + """Per-cell RMSE maps for velocity, EKE, SST and both OHC layers.""" + if not self._observations_configured("obs_rmse_maps"): + return + + rollouts, products, frame = self._obs_inputs() + assert self.observations is not None # established by _obs_inputs + obs_figures.rmse_map_figures( + rollouts, + products, + frame, + self.observations.window, + self.obs_path, + self.observations.velocity_kind, + ) + + def step_obs_annual_rmse(self): + """Per-year totals behind each headline score, with its interval.""" + if not self._observations_configured("obs_annual_rmse"): + return + + _, _, frame = self._obs_inputs() + obs_figures.save( + obs_figures.annual_rmse_panel(frame, "Annual total RMSE vs observations"), + self.obs_path, + "annual_total_rmse_vs_observations", + ) + + def step_obs_variance_maps(self): + """Residual-anomaly variance maps for SST and upper-700 m OHC.""" + if not self._observations_configured("obs_variance_maps"): + return + + rollouts, products, frame = self._obs_inputs() + obs_figures.variance_map_figures(rollouts, products, frame, self.obs_path) + + def step_obs_timeseries(self): + """Global-mean SST, EKE and OHC series, with trends and residuals.""" + if not self._observations_configured("obs_timeseries"): + return + + rollouts, products, _ = self._obs_inputs() + assert self.observations is not None # established by _obs_inputs + obs_figures.timeseries_figures( + rollouts, products, self.obs_path, self.observations.velocity_kind + ) + + def step_obs_spectra(self): + """Spatial and temporal spectra, and their interannual bands.""" + if not self._observations_configured("obs_spectra"): + return + + rollouts, products, _ = self._obs_inputs() + assert self.observations is not None # established by _obs_inputs + obs_figures.spectra_figures( + rollouts, products, self.obs_path, self.observations.velocity_kind + ) + + +def isnan(x: xr.DataArray) -> xr.DataArray: + """Wrapped around np.isnan which correctly reflects the type we use it on.""" + return np.isnan(x) # type: ignore + + +def _combine_variables_by_level( + ds: xr.Dataset, combine_vars: list[str], dataset_spec: DatasetSpec +) -> xr.Dataset: + """ + Combine variables in the dataset along a new 'lev' dimension based on their suffix. + + Parameters: + ds (xarray.Dataset): The input dataset containing variables with suffixes + (e.g., thetao_0, so_1). + combine_vars (list): List of variable prefixes to combine. + + Returns: + xarray.Dataset: The dataset with combined variables and a new 'lev' dimension. + """ + for v in combine_vars: + level_numbers = [i for i in range(len(dataset_spec.depth_levels))] + sorted_vars = [v + "_" + str(lev) for lev in level_numbers] + if sorted_vars[0] not in ds.data_vars: + continue + combined = xr.concat([ds[var] for var in sorted_vars], dim="lev") + combined = combined.assign_coords(lev=list(dataset_spec.depth_levels)) + ds[v] = combined + ds = ds.drop_vars(sorted_vars) + return ds + + +def combine_variables_by_level( + ds_groundtruth: xr.Dataset, + pred_dict: dict[str, dict[str, Any]], + dataset_spec: DatasetSpec, +) -> tuple[xr.Dataset, dict[str, dict[str, Any]]]: + """ + Combine variables by level for ground truth and predictions. + + Parameters: + ds_groundtruth (xarray.Dataset): The ground truth dataset. + pred_dict (dict): Dictionary containing prediction datasets. + + Returns: + xarray.Dataset, dict: Updated ground truth and prediction datasets. + """ + ds_groundtruth = _combine_variables_by_level( + ds_groundtruth, ["thetao", "so", "uo", "vo", "mask"], dataset_spec + ) + for key in pred_dict.keys(): + pred_dict[key]["ds_prediction"] = _combine_variables_by_level( + pred_dict[key]["ds_prediction"], pred_dict[key]["ls"], dataset_spec + ) + return ds_groundtruth, pred_dict + + +def _postprocess_for_plot( + ds, + areacello: np.ndarray, + areacello_spherical: np.ndarray, + dz, + times, + wetmask, + coords=None, +): + """ + Postprocess the dataset to make it compatible with plotting functions. + """ + ds = ds.transpose("time", "lev", ...) + ds["time"] = times + if coords is not None: + ds = ds.assign_coords(coords) + if "thetao" in ds.data_vars: + ds["thetao"] = ds["thetao"].assign_attrs( + long_name=r"${\theta_O}$", units=r"$\degree C$" + ) + if "so" in ds.data_vars: + ds["so"] = ds["so"].assign_attrs(long_name=r"${s}$", units=r"psu") + if "zos" in ds.data_vars: + ds["zos"] = ds["zos"].assign_attrs(long_name=r"SSH", units=r"m") + if "vo" in ds.data_vars: + ds["vo"] = ds["vo"].assign_attrs(long_name=r"${v}$", units=r"m/s") + if "uo" in ds.data_vars: + ds["uo"] = ds["uo"].assign_attrs(long_name=r"${u}$", units=r"m/s") + + ds["lev"] = ds["lev"].assign_attrs(long_name="depth", units="m") + if "init_time" in ds.coords: + ds = ds.drop(["init_time", "valid_time"]) + + for var in ds.data_vars: + if "lev" in ds[var].dims: + ds[var] = ds[var].where(wetmask) + else: + ds[var] = ds[var].where(wetmask.isel(lev=0)) + + ds["areacello"] = (["lat", "lon"], areacello) + ds["areacello_spherical"] = (["lat", "lon"], areacello_spherical) + ds["dz"] = ("lev", dz) + return ds + + +def postprocess_for_plot( + ds_groundtruth, areacello: xr.DataArray, dz: np.ndarray, pred_dict +): + """ + Postprocess for plotting. + + Parameters: + ds_groundtruth (xarray.Dataset): The ground truth dataset. + areacello (xarray.DataArray): areacello dataarray. + pred_dict (dict): Dictionary containing prediction datasets. + + Returns: + xarray.Dataset, dict: Postprocessed ground truth and prediction datasets. + """ + areacello_values = areacello.values + times = ds_groundtruth.time + areacello_spherical_values = ds_groundtruth["areacello_spherical"].values + + # Masking land with NaNs + if "mask" in ds_groundtruth.data_vars: + wetmask = ds_groundtruth["mask"].isel( + time=0, missing_dims="ignore" + ) # our data does not always have time for a mask + else: + wetmask = ds_groundtruth.wetmask + + ds_groundtruth = _postprocess_for_plot( + ds_groundtruth, areacello_values, areacello_spherical_values, dz, times, wetmask + ) + + coords = ds_groundtruth.coords + + for key in pred_dict.keys(): + pred_dict[key]["ds_prediction"] = _postprocess_for_plot( + pred_dict[key]["ds_prediction"], + areacello_values, + areacello_spherical_values, + dz, + times, + wetmask, + coords=coords, + ) + + # Rename lat and lon to y and x + pred_dict[key]["ds_prediction"] = pred_dict[key]["ds_prediction"].rename( + {"lat": "y", "lon": "x"} + ) + + # Rename lat and lon to y and x (This needs to be done in the end!) + ds_groundtruth = ds_groundtruth.rename({"lat": "y", "lon": "x"}) + + return ds_groundtruth, pred_dict + + +def process_data( + data: xr.Dataset, + pred_dict: dict[str, dict[str, Any]], + dataset_spec: DatasetSpec, +) -> tuple[xr.Dataset, dict[str, dict[str, Any]]]: + """ + Get plot ready OM4 data. + """ + ds_groundtruth = with_level_index_vars(data, dataset_spec=dataset_spec) + + # Store ds_prediction + copy_dict = deepcopy(pred_dict) + + for key in pred_dict.keys(): + ds_prediction = pred_dict[key]["data"] + + assert ds_prediction.time.size == ds_groundtruth.time.size, ( + f"Sizes different for {key}: {ds_prediction.time.size}!=" + f"{ds_groundtruth.time.size}; prediction range is " + f"{ds_prediction.time.values[0]} to " + f"{ds_prediction.time.values[-1]}\n" + f"groundtruth range is {ds_groundtruth.time.values[0]} to " + f"{ds_groundtruth.time.values[-1]}" + ) + if "model_path" in ds_prediction.attrs: + copy_dict[key]["model_path"] = ds_prediction.attrs["model_path"] + + pred_dict[key]["ds_prediction"] = ds_prediction + + ### Combine Variables by level + ds_groundtruth, pred_dict = combine_variables_by_level( + ds_groundtruth, pred_dict, dataset_spec + ) + + ### Postprocess predictions for plotting + ds_groundtruth, pred_dict = postprocess_for_plot( + ds_groundtruth, + ds_groundtruth.areacello, + np.array(dataset_spec.depth_thickness), + pred_dict, + ) + + return ds_groundtruth, pred_dict + + +def remove_climatology(ds): + # Compute the climatology on the detrended data + climatology = ds.groupby("time.dayofyear").mean("time").compute() + + # Remove the seasonal cycle (climatology) from the detrended data + day_of_year = ds["time"].dt.dayofyear + res = (ds - climatology.sel(dayofyear=day_of_year)).compute() + + return res + + +def process_mask(data, mask): + mask = mask.where(mask != 0, np.nan) + mask = mask.transpose("lat", "lon") + mask = mask.assign_coords(lat=data.y.values, lon=data.x.values) + mask = mask.rename({"lat": "y", "lon": "x"}) + return mask + + +def profile_mean(ds: xr.Dataset) -> xr.Dataset: + """ + Compute the mean of each variable for each time step. + """ + return ds.weighted(ds.areacello).mean(["y", "x"]) + + +def get_basin_datasets(ds, basin_masks, data): + da_temp = ds * basin_masks["Atlantic"] + section_mask = isnan(da_temp).all("x") + da_temp_int_x = da_temp.weighted(data["areacello"]).mean(["x"]) + At = da_temp_int_x.where(~section_mask) + + da_temp = ds * basin_masks["Indian"] + section_mask = isnan(da_temp).all("x") + da_temp_int_x = da_temp.weighted(data["areacello"]).mean(["x"]) + In = da_temp_int_x.where(~section_mask) + + da_temp = ds * basin_masks["Pacific"] + section_mask = np.isnan(da_temp).all("x") + da_temp_int_x = da_temp.weighted(data["areacello"]).mean(["x"]) + Pa = da_temp_int_x.where(~section_mask) + + da_temp = ds * basin_masks["Southern"] + section_mask = isnan(da_temp).all("x") + da_temp_int_x = da_temp.weighted(data["areacello"]).mean(["x"]) + So = da_temp_int_x.where(~section_mask) + + da_temp = ds * basin_masks["Arctic"] + section_mask = isnan(da_temp).all("x") + da_temp_int_x = da_temp.weighted(data["areacello"]).mean(["x"]) + Ar = da_temp_int_x.where(~section_mask) + + da_temp = ds + section_mask = isnan(da_temp).all("x") + da_temp_int_x = da_temp.weighted(data["areacello"]).mean(["x"]) + Gl = da_temp_int_x.where(~section_mask) + + return [At, In, Pa, So, Ar, Gl], [ + "Atlantic", + "Indian", + "Pacific", + "Southern", + "Arctic", + "Global", + ] diff --git a/samudra/viz/norms.py b/samudra/viz/norms.py new file mode 100644 index 0000000000000000000000000000000000000000..9c0887c5e26c736ee87e331d7d93088c4c0ccf10 --- /dev/null +++ b/samudra/viz/norms.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Colour limits for map panels. + +Percentile rather than min/max, so a single outlier cell cannot flatten a whole +figure, with the degenerate cases -- an all-NaN field, a constant one, an +inverted range -- resolved to something matplotlib can draw. + +Its own module rather than part of `viz.core` so that `viz.observations` can use +it without the two importing each other. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import numpy as np +import xarray as xr +from dask.array.core import Array as DaskArray +from matplotlib import colors + + +def _flatten_for_norm(data: Any) -> np.ndarray: + if isinstance(data, xr.DataArray): + arr = data.data + elif isinstance(data, np.ndarray): + arr = data + elif isinstance(data, DaskArray): + arr = data + elif isinstance(data, Iterable) and not isinstance(data, (str, bytes)): + arrays = [_flatten_for_norm(item) for item in data] + if arrays: + return np.concatenate(arrays) + return np.array([], dtype=float) + else: + arr = data + + if isinstance(arr, DaskArray): + arr = arr.compute() + return np.asarray(arr).ravel() + + +def symmetric_percentile_norm( + data: Any, percentile: float = 98.0, fallback: float = 1.0 +) -> colors.Normalize: + flat = _flatten_for_norm(data) + flat = flat[~np.isnan(flat)] + if flat.size == 0: + max_abs = fallback + else: + max_abs = np.percentile(np.abs(flat), percentile) + if not np.isfinite(max_abs) or max_abs == 0: + max_abs = fallback + return colors.Normalize(vmin=-max_abs, vmax=max_abs) + + +def percentile_norm( + data: Any, lower: float = 2.0, upper: float = 98.0, fallback: float = 1.0 +) -> colors.Normalize: + flat = _flatten_for_norm(data) + flat = flat[~np.isnan(flat)] + if flat.size == 0: + vmin = 0.0 + vmax = fallback + else: + vmin = np.percentile(flat, lower) + vmax = np.percentile(flat, upper) + if not np.isfinite(vmin) or not np.isfinite(vmax): + vmin = 0.0 + vmax = fallback + elif vmin == vmax: + vmin = vmin - fallback + vmax = vmax + fallback + elif vmin > vmax: + vmin, vmax = vmax, vmin + return colors.Normalize(vmin=vmin, vmax=vmax) diff --git a/samudra/viz/observations.py b/samudra/viz/observations.py new file mode 100644 index 0000000000000000000000000000000000000000..f0c47fc4c3cd27802f78db6e48d6febce393bb96 --- /dev/null +++ b/samudra/viz/observations.py @@ -0,0 +1,974 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +"""Figures comparing a rollout against observations. + +Every number drawn here comes from `samudra.metrics`, the same kernels +`samudra.eval` reduces with, so a value printed on a figure is the value logged +to W&B rather than a second implementation that happens to agree. + +Figures are written to disk rather than logged as images: a full suite runs to +tens of megabytes per eval, which is worth a bucket and not worth W&B's +storage. +""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Callable + +import cartopy.crs as ccrs # type: ignore +import cartopy.feature as cfeature # type: ignore +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr +from cartopy.mpl.geoaxes import GeoAxes # type: ignore + +from samudra.constants import build_om4_spec +from samudra.metrics import comparisons, kernels, spectra +from samudra.viz.norms import percentile_norm, symmetric_percentile_norm + +logger = logging.getLogger(__name__) + +# viz assumes OM4 throughout; the observation figures inherit that. +_OM4_SPEC = build_om4_spec(prognostic_vars_key="thermo_dynamic_all") + +PROJECTION = ccrs.Robinson(central_longitude=-150) +PLATE = ccrs.PlateCarree() + +# Colour per run, with the reference first. Distinguishable in greyscale and to +# the most common colour-vision deficiencies. +SERIES_COLOURS = ("#000000", "#0072b2", "#d55e00", "#009e73", "#cc79a7") + + +# What an RMSE-map builder hands back: squared error over time, and the cell +# areas to weight it by. +_SquaredError = tuple[xr.DataArray, xr.DataArray] + + +def save(figure: plt.Figure, directory: str, name: str) -> str: + """Write a figure as PDF and close it. Returns the path.""" + os.makedirs(directory, exist_ok=True) + path = os.path.join(directory, f"{name}.pdf") + figure.savefig(path, bbox_inches="tight", dpi=200) + plt.close(figure) + logger.info(" wrote %s", path) + return path + + +def _annotate(value: float, units: str, row: pd.Series | None = None) -> str: + """Format a metric for a panel title, with its interval when there is one.""" + if not np.isfinite(value): + return "unavailable" + text = f"{value:.4g} {units}".strip() + if row is None: + return text + low, high = row.get("ci_low"), row.get("ci_high") + if low is not None and high is not None and np.isfinite(low) and np.isfinite(high): + text += f" [{low:.4g}, {high:.4g}]" + return text + + +def _map_axes(axis: GeoAxes) -> None: + axis.add_feature(cfeature.LAND, facecolor="0.85", zorder=2) + axis.coastlines(linewidth=0.3, color="0.4", zorder=3) + axis.set_global() + + +def map_panels( + fields: dict[str, xr.DataArray], + title: str, + units: str, + subtitles: dict[str, str] | None = None, + cmap: str = "viridis", + diverging: bool = False, +) -> plt.Figure: + """One map per run, on a shared colour scale so panels are comparable.""" + names = list(fields) + figure, axes = plt.subplots( + 1, + len(names), + figsize=(6.5 * len(names), 4.2), + subplot_kw={"projection": PROJECTION}, + squeeze=False, + ) + # Shared percentile limits, so one outlier cell cannot flatten every panel. + # `viz.core` already computes these, including the degenerate cases a + # constant or all-NaN field produces. + pooled = list(fields.values()) + norm = symmetric_percentile_norm(pooled) if diverging else percentile_norm(pooled) + + mesh = None + for axis, name in zip(axes.flat, names, strict=True): + field = fields[name] + mesh = axis.pcolormesh( + field["lon"], + field["lat"], + field.values, + transform=PLATE, + cmap=cmap, + norm=norm, + shading="auto", + ) + _map_axes(axis) + label = name if subtitles is None else f"{name}\n{subtitles.get(name, '')}" + axis.set_title(label, fontsize=10) + + if mesh is not None: + bar = figure.colorbar( + mesh, ax=list(axes.flat), orientation="horizontal", fraction=0.05, pad=0.05 + ) + bar.set_label(units) + figure.suptitle(title, fontsize=12) + return figure + + +def series_panel( + series: dict[str, pd.Series], + title: str, + ylabel: str, + annotations: dict[str, str] | None = None, +) -> plt.Figure: + """One line per run against a shared time axis.""" + figure, axis = plt.subplots(figsize=(11, 4)) + for (name, values), colour in zip(series.items(), SERIES_COLOURS, strict=False): + label = name if annotations is None else f"{name} — {annotations.get(name, '')}" + axis.plot(values.index, values.to_numpy(), label=label, color=colour, lw=1.2) + axis.set_ylabel(ylabel) + axis.set_title(title, fontsize=12) + axis.grid(alpha=0.3) + axis.legend(fontsize=8) + return figure + + +def spectra_panel( + curves: dict[str, dict[str, tuple[np.ndarray, np.ndarray]]], + title: str, + xlabel: str, + ylabel: str, + scores: dict[str, dict[str, float]] | None = None, +) -> plt.Figure: + """One panel per region, each with every run's spectrum on log-log axes.""" + regions = list(next(iter(curves.values()), {})) + if not regions: + regions = ["(no data)"] + figure, axes = plt.subplots( + 1, len(regions), figsize=(5.2 * len(regions), 4.2), squeeze=False + ) + + for axis, region in zip(axes.flat, regions, strict=True): + for (name, per_region), colour in zip( + curves.items(), SERIES_COLOURS, strict=False + ): + x, y = per_region.get(region, (np.array([]), np.array([]))) + if x.size == 0: + continue + label = name + if scores is not None: + score = scores.get(name, {}).get(region) + if score is not None and np.isfinite(score): + label = f"{name} ({score:.3f} dex)" + axis.loglog(x, y, label=label, color=colour, lw=1.4) + axis.set_title(region, fontsize=10) + axis.set_xlabel(xlabel) + axis.grid(alpha=0.3, which="both") + if axis.get_legend_handles_labels()[0]: + axis.legend(fontsize=8) + else: + axis.text( + 0.5, + 0.5, + "no spectrum:\nbox too small or\nnot fully observed", + ha="center", + va="center", + transform=axis.transAxes, + fontsize=9, + color="0.5", + ) + axes.flat[0].set_ylabel(ylabel) + figure.suptitle(title, fontsize=12) + return figure + + +def interannual_spectra_panel( + bands: dict[str, dict[str, tuple]], + title: str, + xlabel: str, + ylabel: str, +) -> plt.Figure: + """Per-region geometric-mean spectra with their year-to-year spread. + + The shaded band is +/- one log standard deviation across years, so a model + whose spectrum sits inside the observed band is within natural variability + rather than merely close on average. + """ + regions = list(next(iter(bands.values()), {})) or ["(no data)"] + figure, axes = plt.subplots( + 1, len(regions), figsize=(5.2 * len(regions), 4.2), squeeze=False + ) + for axis, region in zip(axes.flat, regions, strict=True): + for (name, per_region), colour in zip( + bands.items(), SERIES_COLOURS, strict=False + ): + band = per_region.get(region) + if band is None: + continue + x, mean, lower, upper, years = band + if getattr(x, "size", 0) == 0: + continue + axis.fill_between(x, lower, upper, color=colour, alpha=0.15, linewidth=0) + axis.loglog(x, mean, color=colour, lw=1.6, label=f"{name} ({years} yr)") + axis.set_title(region, fontsize=10) + axis.set_xlabel(xlabel) + axis.grid(alpha=0.3, which="both") + if axis.get_legend_handles_labels()[0]: + axis.legend(fontsize=8) + else: + axis.text( + 0.5, + 0.5, + "no spectrum:\nbox too small or\nnot fully observed", + ha="center", + va="center", + transform=axis.transAxes, + fontsize=9, + color="0.5", + ) + axes.flat[0].set_ylabel(ylabel) + figure.suptitle(title, fontsize=12) + return figure + + +def annual_rmse_panel(frame: pd.DataFrame, title: str) -> plt.Figure: + """Per-year totals behind each headline score, with its interval. + + A single number hides whether a model is steadily wrong or wrong in one + year, which is the difference between a bias and an event. + """ + annual = frame[frame["period_kind"] == "annual"] + primary = frame[frame["period_kind"] == "primary_complete_years"] + metrics = list(dict.fromkeys(annual["metric"] + annual["depth"].fillna(""))) + if not metrics: + metrics = ["(none)"] + + figure, axes = plt.subplots( + 1, len(metrics), figsize=(4.6 * len(metrics), 3.8), squeeze=False + ) + for axis, key in zip(axes.flat, metrics, strict=True): + rows = annual[(annual["metric"] + annual["depth"].fillna("")) == key] + for (name, group), colour in zip( + rows.groupby("model"), SERIES_COLOURS, strict=False + ): + group = group.sort_values("year") + axis.plot( + group["year"], + group["value"], + marker="o", + color=colour, + label=name, + lw=1.2, + ) + head = primary[ + (primary["model"] == name) + & ((primary["metric"] + primary["depth"].fillna("")) == key) + ] + if not head.empty: + row = head.iloc[0] + axis.axhline(row["value"], color=colour, ls="--", lw=0.9, alpha=0.7) + if np.isfinite(row.get("ci_low", np.nan)): + axis.axhspan( + row["ci_low"], row["ci_high"], color=colour, alpha=0.10, lw=0 + ) + axis.set_title(key, fontsize=9) + axis.set_xlabel("year") + axis.grid(alpha=0.3) + axis.legend(fontsize=8) + axes.flat[0].set_ylabel("total RMSE") + figure.suptitle(title, fontsize=12) + return figure + + +def rmse_map_figures( + rollouts: dict[str, xr.Dataset], + products: dict[str, xr.Dataset], + frame: pd.DataFrame, + window: tuple[pd.Timestamp, pd.Timestamp], + directory: str, + velocity_kind: str = "absolute", +) -> list[str]: + """Per-cell RMSE maps for every metric that produces one. + + Recomputes the maps rather than reading them back: `eval` emits scalars and + rows, not fields, and using the same kernels keeps the annotated totals + identical to the logged ones. + """ + written: list[str] = [] + duacs, oisst, argo = products["duacs"], products["oisst"], products["argo"] + + def scalar(metric: str, model: str, depth: str | None = None) -> pd.Series | None: + rows = frame[ + (frame["metric"] == metric) + & (frame["model"] == model) + & (frame["period_kind"] == "primary_complete_years") + & (frame["depth"].isna() if depth is None else frame["depth"] == depth) + ] + return None if rows.empty else rows.iloc[0] + + def emit( + name: str, + metric: str, + title: str, + units: str, + build: Callable[[str, xr.Dataset], tuple[xr.DataArray, xr.DataArray]], + depth: str | None = None, + cmap: str = "magma", + ) -> None: + maps, notes = {}, {} + for model, rollout in rollouts.items(): + row = scalar(metric, model, depth) + try: + error_squared, area = build(model, rollout) + # Reduce through the same kernel `eval` scores with, so the map + # and the total annotated on it are one calculation. A plain + # time mean here would weight years by their sample count while + # the annotated total weights them equally. + maps[model] = kernels.rmse_map_with_uncertainty( + error_squared, + area, + ("lat", "lon"), + f"{model} {metric}", + bootstrap_samples=0, + )[0] + except (ValueError, KeyError) as error: + # A metric that was scored must get the map that explains it: + # dropping the figure while its scalar goes to W&B leaves a + # number nobody can check. A metric the frame never scored has + # nothing to draw, and saying so is enough. + if row is not None: + raise ValueError( + f"{model}: {metric} has a scalar in the metric frame " + f"but its RMSE map could not be built" + ) from error + logger.info(" %s: %s was not scored, so no map", model, metric) + continue + notes[model] = ( + _annotate(float(row["value"]), units, row) if row is not None else "" + ) + if maps: + written.append( + save(map_panels(maps, title, units, notes, cmap=cmap), directory, name) + ) + + # Derived once per model: the velocity pair costs a geostrophic derivation + # and an alignment over the whole window, and both maps below want it. + velocities: dict[str, comparisons.VelocityComparison] = {} + + def velocity_of(model: str, rollout: xr.Dataset) -> comparisons.VelocityComparison: + if model not in velocities: + velocities[model] = comparisons.surface_velocity( + rollout, duacs, window, model, velocity_kind + ) + return velocities[model] + + def velocity_map(model: str, rollout: xr.Dataset) -> _SquaredError: + velocity = velocity_of(model, rollout) + return velocity.vector_error_squared, velocity.area + + def eke_map(model: str, rollout: xr.Dataset) -> _SquaredError: + eke = velocity_of(model, rollout).eddy_kinetic_energy() + return eke.error_squared, eke.area + + def sst_map(model: str, rollout: xr.Dataset) -> _SquaredError: + sst = comparisons.sea_surface_temperature(rollout, oisst, window, model) + return sst.error_squared, sst.area + + emit( + "surface_geostrophic_velocity_rmse_vs_duacs", + "surface_geostrophic_velocity_vector_total_rmse", + "Surface geostrophic velocity vector RMSE vs DUACS", + "m s-1", + velocity_map, + ) + emit( + "surface_eke_rmse_vs_duacs", + "instantaneous_surface_eke_total_rmse", + "Instantaneous surface EKE RMSE vs DUACS", + "m2 s-2", + eke_map, + ) + emit( + "sst_rmse_vs_oisst", + "surface_sst_total_rmse", + "SST RMSE vs OISST", + "degC", + sst_map, + ) + + for layer in kernels.OHC_LAYERS: + + def ohc_map(model: str, rollout: xr.Dataset, layer=layer) -> _SquaredError: + ohc = comparisons.ohc_layer( + rollout, + argo, + layer, + window, + model, + comparisons.model_depth_thickness(rollout, _OM4_SPEC), + ) + return ohc.error_squared, ohc.area + + slug = layer.label.split("(")[-1].strip(") ").replace("-", "_") + emit( + f"ohc_per_area_rmse_vs_argo_iap_{slug}", + "ohc_per_area_total_rmse", + f"OHC per-area RMSE vs ARGO-IAP, {layer.label}", + "J m-2", + ohc_map, + depth=layer.label, + ) + + return written + + +def _on_common_cells( + reference: xr.DataArray, runs: dict[str, xr.DataArray] +) -> tuple[xr.DataArray, dict[str, xr.DataArray]]: + """Restrict the reference and every run to the cells they all carry. + + A run covers only the observation cells that survived pairing: its own land, + plus the band regridding erodes around it, which is coastal and so is where + the energy and the error both are. Reducing the reference over every cell + the product has and each run over its own subset makes the gap between the + curves a property of the masks, which a mean-state plot then presents as + bias. + """ + mask = None + for field in runs.values(): + ever = comparisons.ever_finite(field) + mask = ever if mask is None else (mask & ever) + if mask is None: + return reference, runs + return reference.where(mask), { + name: field.where(mask) for name, field in runs.items() + } + + +def _over_one_window(items): + """Trim every comparison to the span they all cover; return the reference too. + + A panel draws several runs beside one observation curve on one colour scale + and one legend, so the reference has to cover the runs' period and no more, + and the runs have to cover each other's. Built with a `None` window, each + comparison spans only what that run and the product share; this intersects + those. + """ + starts = [pd.Timestamp(c.time.values.min()) for c in items.values()] + ends = [pd.Timestamp(c.time.values.max()) for c in items.values()] + window = slice(max(starts), min(ends)) + trimmed = {name: c.slice(window) for name, c in items.items()} + first = next(iter(trimmed.values())) + reference = first.obs if hasattr(first, "obs") else first.eastward.obs + return reference, trimmed + + +def variance_map_figures( + rollouts: dict[str, xr.Dataset], + products: dict[str, xr.Dataset], + frame: pd.DataFrame, + directory: str, +) -> list[str]: + """Residual-variance maps: how much unforced variability, and where. + + Variance is computed on each model's native grid and only the resulting map + regridded. Interpolating the field first would smooth it and damp the + variance by an amount set by the resolution ratio. + + Every panel covers the period all runs and the observations share, so the + reference is not a longer record drawn beside shorter ones on a common + colour scale -- and it matches the span behind the annotated numbers. + """ + written: list[str] = [] + oisst, argo = products["oisst"], products["argo"] + + def note(metric: str, model: str, units: str, drawn: tuple[str, str]) -> str: + rows = frame[ + (frame["metric"] == metric) + & (frame["model"] == model) + & (frame["period_kind"] == "full_overlap") + ] + if rows.empty: + return "" + row = rows.iloc[0] + text = _annotate(float(row["value"]), units) + # `report` scores each run over its own overlap with the product, while + # the panels share one window so they can share a colour scale. With + # runs of equal coverage those agree; when they do not, say so rather + # than print a number for a period the map does not cover. + scored = (str(row["period_start"]), str(row["period_end"])) + if scored != drawn: + text += f" (scored {scored[0]}..{scored[1]})" + return text + + def draw( + items: dict[str, comparisons.Comparison], + reference_name: str, + slug: str, + title: str, + units: str, + rmse_metric: str, + corr_metric: str, + ) -> None: + reference, items = _over_one_window(items) + stamps = pd.DatetimeIndex(reference["time"].values) + drawn = (f"{stamps.min():%Y-%m-%d}", f"{stamps.max():%Y-%m-%d}") + obs_variance = kernels.residual_variance_map(reference) + maps = {reference_name: obs_variance} + notes = {reference_name: ""} + for model, item in items.items(): + maps[model] = kernels.model_field_on_obs_grid( + kernels.residual_variance_map(item.native), obs_variance + ) + notes[model] = ( + f"map RMSE {note(rmse_metric, model, units, drawn)}; " + f"corr {note(corr_metric, model, '', drawn)}" + ) + written.append( + save( + map_panels( + maps, + f"{title}, {drawn[0]} to {drawn[1]}", + units, + notes, + cmap="cividis", + ), + directory, + slug, + ) + ) + + draw( + { + model: comparisons.sea_surface_temperature(rollout, oisst, None, model) + for model, rollout in rollouts.items() + }, + "OISST", + "sst_residual_variance_vs_oisst", + "SST residual-anomaly variance vs OISST", + "degC2", + "surface_sst_residual_variance_map_rmse", + "surface_sst_residual_variance_pattern_corr", + ) + + upper = kernels.OHC_LAYERS[0] + draw( + { + model: comparisons.ohc_layer( + rollout, + argo, + upper, + None, + model, + comparisons.model_depth_thickness(rollout, _OM4_SPEC), + complete_years_only=False, + ) + for model, rollout in rollouts.items() + }, + "ARGO-IAP", + "ohc_upper700_residual_variance_vs_argo_iap", + "Upper-700 m OHC residual-anomaly variance vs ARGO-IAP", + "(J m-2)2", + "ohc_upper700_per_area_residual_variance_map_rmse", + "ohc_upper700_per_area_residual_variance_pattern_corr", + ) + return written + + +def timeseries_figures( + rollouts: dict[str, xr.Dataset], + products: dict[str, xr.Dataset], + directory: str, + velocity_kind: str = "absolute", +) -> list[str]: + """Global-mean series: the trend, the seasonal cycle, and what is left. + + The residual is the quantity of interest. Trend and seasonality are the easy + parts to reproduce; what remains is the variability a model either captures + or smooths away. + """ + written: list[str] = [] + duacs, oisst, argo = products["duacs"], products["oisst"], products["argo"] + + def draw( + items, + reference_name: str, + slug: str, + title: str, + units: str, + derive=None, + ) -> None: + reference, items = _over_one_window(items) + if derive is not None: + # Anything reducing over time has to be derived after the window is + # settled: an eddy anomaly is taken about the mean of whatever + # record it is handed, and trimming afterwards cannot undo that. + items = {name: derive(item) for name, item in items.items()} + reference = next(iter(items.values())).obs + area = next(iter(items.values())).area + reference, runs = _on_common_cells( + reference, {name: item.on_obs_grid for name, item in items.items()} + ) + series = {reference_name: _global_mean(reference, area)} + series.update({name: _global_mean(f, area) for name, f in runs.items()}) + if len(series) < 2: + return + + trends = { + name: ( + f"trend {kernels.series_linear_trend_per_year(values):+.3g} {units}/yr; " + f"residual var {kernels.series_residual_variance(values):.3g}" + ) + for name, values in series.items() + } + written.append( + save(series_panel(series, title, units, trends), directory, slug) + ) + + residuals = { + name: kernels.series_without_linear_trend( + kernels.series_without_seasonal_cycle(values) + ) + for name, values in series.items() + } + written.append( + save( + series_panel( + residuals, f"{title} — detrended and deseasonalized", units + ), + directory, + f"{slug}_residual", + ) + ) + + draw( + { + model: comparisons.sea_surface_temperature(rollout, oisst, None, model) + for model, rollout in rollouts.items() + }, + "OISST", + "global_sst", + "Global mean SST", + "degC", + ) + draw( + { + model: comparisons.surface_velocity( + rollout, duacs, None, model, velocity_kind + ) + for model, rollout in rollouts.items() + }, + "DUACS", + "global_surface_eke", + "Global mean surface EKE", + "m2 s-2", + derive=lambda velocity: velocity.eddy_kinetic_energy(), + ) + upper = kernels.OHC_LAYERS[0] + draw( + { + model: comparisons.ohc_layer( + rollout, + argo, + upper, + None, + model, + comparisons.model_depth_thickness(rollout, _OM4_SPEC), + complete_years_only=False, + ) + for model, rollout in rollouts.items() + }, + "ARGO-IAP", + "global_ohc_upper700", + "Global mean upper-700 m OHC per area", + "J m-2", + ) + return written + + +def _global_mean(field: xr.DataArray, area: xr.DataArray) -> pd.Series: + """Area-weighted global mean as a pandas series.""" + values = kernels.area_weighted_timeseries(field, area, ("lat", "lon")) + return pd.Series( + np.asarray(values.values, dtype=float), + index=pd.DatetimeIndex(values["time"].values), + ) + + +def spectra_figures( + rollouts: dict[str, xr.Dataset], + products: dict[str, xr.Dataset], + directory: str, + velocity_kind: str = "absolute", +) -> list[str]: + """Spatial and temporal spectra: at which scales the model is wrong. + + A model can carry the right total variance and still distribute it over the + wrong wavelengths. These are the only figures that show that directly, and + the log10 scores beside each curve say how far off it is in dex. + + Every observation curve is reduced over the span the runs share. That is not + only labelling: an eddy anomaly is taken about the mean of whatever record + it is handed, so a reference built from the product's full archive is not + the same quantity as a run built from eight years of it. + """ + written: list[str] = [] + duacs, oisst = products["duacs"], products["oisst"] + + # Trimmed once, up front: the energies below reduce over time, so a mean + # taken before the span is settled is a mean of the wrong record. + _, velocity = _over_one_window( + { + model: comparisons.surface_velocity( + rollout, duacs, None, model, velocity_kind + ) + for model, rollout in rollouts.items() + } + ) + obs_sst, temperature = _over_one_window( + { + model: comparisons.sea_surface_temperature(rollout, oisst, None, model) + for model, rollout in rollouts.items() + } + ) + + # --- surface EKE, spatial ------------------------------------------------ + mean_eke = {m: v.mean_eddy_kinetic_energy() for m, v in velocity.items()} + obs_eke = next(iter(mean_eke.values())).obs + eke_curves = {"DUACS": _region_curves(obs_eke, spectra.SPATIAL_REGIONS)} + for model, item in mean_eke.items(): + # The model's own grid, like every other spectrum here. Interpolating + # onto the finer observation grid is a low-pass filter, so it damps the + # high wavenumbers this figure exists to compare -- and the wavenumbers + # come back in rad/km either way, so the curves stay comparable. + eke_curves[model] = _region_curves(item.native, spectra.SPATIAL_REGIONS) + written.append( + save( + spectra_panel( + eke_curves, + "Surface EKE spatial spectra vs DUACS", + "wavenumber (rad km-1)", + "k x PSD", + _curve_scores(eke_curves, "DUACS"), + ), + directory, + "eke_spatial_spectra_vs_duacs", + ) + ) + + # --- SST anomaly, spatial ------------------------------------------------ + sst = temperature + obs_ssta = kernels.without_seasonal_cycle(obs_sst) + ssta_curves = {"OISST": _region_curves(obs_ssta, spectra.SPATIAL_REGIONS)} + for model, item in sst.items(): + ssta_curves[model] = _region_curves( + kernels.without_seasonal_cycle(item.native), spectra.SPATIAL_REGIONS + ) + written.append( + save( + spectra_panel( + ssta_curves, + "SST anomaly spatial spectra vs OISST", + "wavenumber (rad km-1)", + "k x PSD", + _curve_scores(ssta_curves, "OISST"), + ), + directory, + "ssta_spatial_spectra_vs_oisst", + ) + ) + + # --- surface KE, temporal ------------------------------------------------ + kinetic = {m: v.kinetic_energy() for m, v in velocity.items()} + obs_ke = next(iter(kinetic.values())).obs + obs_ke, ke_runs = _on_common_cells( + obs_ke, {m: item.on_obs_grid for m, item in kinetic.items()} + ) + ke_curves = {"DUACS": _temporal_curves(obs_ke, spectra.TEMPORAL_REGIONS)} + ke_curves.update( + {m: _temporal_curves(f, spectra.TEMPORAL_REGIONS) for m, f in ke_runs.items()} + ) + written.append( + save( + spectra_panel( + ke_curves, + "Surface geostrophic KE temporal spectra vs DUACS", + "frequency (cycles yr-1)", + "PSD", + _curve_scores(ke_curves, "DUACS"), + ), + directory, + "ke_temporal_spectra_vs_duacs", + ) + ) + + # --- interannual bands --------------------------------------------------- + first = next(iter(velocity.values())) + eke_bands = { + "DUACS": _yearly_bands(obs_eke_by_year(first.eastward.obs, first.northward.obs)) + } + for model, item in velocity.items(): + eke_bands[model] = _yearly_bands( + obs_eke_by_year(item.eastward.native, item.northward.native) + ) + written.append( + save( + interannual_spectra_panel( + eke_bands, + "Interannual surface EKE spatial spectra", + "wavenumber (rad km-1)", + "k x PSD", + ), + directory, + "eke_spatial_spectra_interannual", + ) + ) + + ssta_bands = {"OISST": _yearly_bands(_anomaly_by_year(obs_ssta))} + for model, item in sst.items(): + ssta_bands[model] = _yearly_bands( + _anomaly_by_year(kernels.without_seasonal_cycle(item.native)) + ) + written.append( + save( + interannual_spectra_panel( + ssta_bands, + "Interannual SST anomaly spatial spectra", + "wavenumber (rad km-1)", + "k x PSD", + ), + directory, + "ssta_spatial_spectra_interannual", + ) + ) + + ke_bands = {"DUACS": _yearly_temporal_bands(obs_ke)} + for model, field in ke_runs.items(): + ke_bands[model] = _yearly_temporal_bands(field) + written.append( + save( + interannual_spectra_panel( + ke_bands, + "Interannual surface geostrophic KE temporal spectra", + "frequency (cycles yr-1)", + "PSD", + ), + directory, + "ke_temporal_spectra_interannual", + ) + ) + + return written + + +def _anomaly_by_year(anomaly: xr.DataArray) -> dict[int, xr.DataArray]: + """Split an anomaly field into complete calendar years, keeping the time axis. + + The anomaly must already be taken about the whole record's climatology, so + that every year shares one baseline and the band shows real year-to-year + spread rather than each year's deviation from itself. + """ + return { + year: anomaly.sel(time=slice(f"{year}-01-01", f"{year}-12-31")) + for year in kernels.complete_calendar_years(anomaly["time"]) + } + + +def _yearly_temporal_bands(field: xr.DataArray) -> dict[str, tuple]: + """Welch PSD of each complete calendar year, aggregated into a per-region band.""" + years = kernels.complete_calendar_years(field["time"]) + bands = {} + for name, lon, lat in spectra.TEMPORAL_REGIONS: + series = spectra.region_mean_series(field, lon, lat) + stamps = pd.DatetimeIndex(series.index) + curves = [] + for year in years: + frequencies, power = spectra.welch_psd(series[stamps.year == year]) + if np.asarray(frequencies).size: + curves.append((np.asarray(frequencies), np.asarray(power))) + bands[name] = spectra.interannual_band(curves) + return bands + + +def obs_eke_by_year(u: xr.DataArray, v: xr.DataArray) -> dict[int, xr.DataArray]: + """Annual-mean EKE maps for each complete calendar year. + + A year the record only partly covers is dropped rather than averaged: the + default rollout starts on 20 October, so its first year holds fifteen + samples of one season and would enter the band as if it were a year. + + The eddy anomaly is taken about the mean of the whole record, not of each + year: a per-year mean would absorb the year-to-year change in the mean flow, + which is the variability these bands exist to show. + """ + eke = kernels.instantaneous_surface_eke(u, v) + return { + year: eke.sel(time=slice(f"{year}-01-01", f"{year}-12-31")).mean( + "time", skipna=True + ) + for year in kernels.complete_calendar_years(u["time"]) + } + + +def _region_curves( + field: xr.DataArray, regions: tuple[tuple[str, slice, slice], ...] +) -> dict[str, tuple[np.ndarray, np.ndarray]]: + return { + name: spectra.region_spectrum(field, lon, lat, name=name) + for name, lon, lat in regions + } + + +def _temporal_curves( + field: xr.DataArray, regions: tuple[tuple[str, slice, slice], ...] +) -> dict[str, tuple[np.ndarray, np.ndarray]]: + return { + name: spectra.welch_psd(spectra.region_mean_series(field, lon, lat)) + for name, lon, lat in regions + } + + +def _yearly_bands( + per_year: dict[int, xr.DataArray], +) -> dict[str, tuple]: + bands = {} + for name, lon, lat in spectra.SPATIAL_REGIONS: + curves = [ + spectra.region_spectrum(field, lon, lat, name=f"{name} {year}") + for year, field in per_year.items() + ] + bands[name] = spectra.interannual_band( + [(x, y) for x, y in curves if getattr(x, "size", 0)] + ) + return bands + + +def _curve_scores( + curves: dict[str, dict[str, tuple[np.ndarray, np.ndarray]]], reference: str +) -> dict[str, dict[str, float]]: + """log10 RMSE of each run's spectrum against the reference, per region.""" + baseline = curves.get(reference, {}) + empty = (np.array([]), np.array([])) + + scores: dict[str, dict[str, float]] = {} + for name, per_region in curves.items(): + if name == reference: + continue + scores[name] = {} + for region, (wavenumber, power) in per_region.items(): + reference_wavenumber, reference_power = baseline.get(region, empty) + scores[name][region] = spectra.log10_rmse_between_curves( + reference_wavenumber, reference_power, wavenumber, power + ) + return scores