Add GFDL OM4/CM4 ocean dataset pipeline on xarray_beam + Dataflow#1340
Draft
jpdunc23 wants to merge 11 commits into
Draft
Add GFDL OM4/CM4 ocean dataset pipeline on xarray_beam + Dataflow#1340jpdunc23 wants to merge 11 commits into
jpdunc23 wants to merge 11 commits into
Conversation
Producing training-reference datasets for ocean emulation currently runs through scripts/data_process/compute_ocean_dataset.py: one giant lazy dask graph over the full simulation executed with xpartition on a hand-managed cluster on LEAP's 2i2c hub. That approach is fragile (whole-graph memory blowups, manual cluster scaling, S3/OSN credential coupling) and doesn't consume the regenerated source data, which now provides true snapshots alongside interval means, a full surface flux/forcing suite, and vertical coarsening to the 19-level Samudra grid already applied upstream. This PR adds a self-contained pipeline in scripts/gfdl_om4/, modeled on scripts/era5/: an xarray_beam pipeline run on Google Cloud Dataflow (DirectRunner for local subset validation) that reads the native-grid 0.25° tripolar source zarrs, applies per-chunk transforms — vector rotation, wetmask-normalized conservative regridding to an analytic Gaussian grid (F90 1° or F22.5 4°), vertical-level splitting, derived variables, and the established masking/NaN conventions — and writes one templated sharded zarr v3 store per run. Each invocation is driven by a YAML config declaring source stores, streams (variables, renames, transforms), target grid, and output layout, so a new simulation is a config change rather than a code change. Prognostic ocean state comes from snapshot stores with raw timestamps; interval-mean fluxes and ice fields sampled at the snapshot instants share the same single time coordinate; every output variable carries provenance attrs naming its source store and variable. xESMF regridding weights are precomputed once by a setup step and stored as a versioned GCS artifact, and the needed ocean_emulators utilities are ported into the pipeline directory so there is no external dependency on that repo. The first production config is the 5-daily 1° ocean store from the piControl sample year; a 6-hourly sea-ice store and 4° outputs are follow-on configs over the same machinery. Changes: - New scripts/gfdl_om4/ directory (pipeline package, configs/, Makefile, Dockerfile, environment, run script); no changes to existing code. - scripts/data_process/compute_ocean_dataset.py and compute_sea_ice_dataset.py are unchanged. - [ ] Tests added - [ ] If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated
The 12h offset in the coarsened-daily flux store's time labels was an upstream configuration error, since fixed in the regenerated source store. Flux time labels are now consumed verbatim; the cross-stream time-alignment assertions remain as the guard against any mismatch.
…rtifact New scripts/gfdl_om4/ directory with the regridding library layer for the OM4/CM4 ocean dataset pipeline: - pipeline/grids.py: analytic Gaussian target grids (F90, F22.5) with exact quadrature-weight cell areas (mean radius 6371 km). The F90 grid matches the reference gaussian_grid_180_by_360.nc and the existing processed-store coordinates to machine precision. - pipeline/ocean_emulators_port.py: supergrid conversion, vector rotation, C-grid-to-tracer-center interpolation, and explicit wetmask-normalized conservative regridding, ported from the ai2cm fork of ocean_emulators (commit bb88b2586) so the pipeline has no dependency on that repo. Known defects of the ported code are fixed rather than reproduced: coastal fillna(0) velocity bias, polar-radius cell areas, implicit coastal threshold (na_thres) now an explicit named constant, regridded ocean fraction kept as an output. - pipeline/weights.py: setup entry point precomputing conservative xESMF weights per source x target grid pair into a versioned GCS artifact (weights + source geometry), and the per-process cached regridder loader used by workers. - environment.yaml, dataflow-requirements.txt, Makefile, README.md following the scripts/era5 operational pattern. Verified: F90 vs reference grid max coord diff 4e-14 deg; regridded wetmask footprint identical to the existing processed store's mask_0; vector rotation preserves speed and round-trips to machine epsilon; area-weighted global mean preserved through the normalized regrid.
jpdunc23
commented
Jul 6, 2026
A YAML-configured xarray_beam pipeline (DirectRunner locally, Dataflow-ready options passthrough) that reads the native-grid snapshot and static stores, applies per-chunk transforms (C-grid-to-center interpolation, vector rotation, wetmask-normalized conservative regridding, level splitting), and writes a templated zarr v3 store with per-level masks, interface depths, exact Gaussian cell areas, provenance attrs on every variable, and fail-fast assertions on variables, levels, time alignment, and valid-data footprints. Chunks are normalized by their instantaneous ocean footprint, which varies slightly in time in the source data.
Slicing the 3D wetmask to a level kept the scalar level coordinate, which rode through masking and regridding into the output store's coordinates. Slice with drop=True everywhere and assert the output template carries only the time/lat/lon coordinates so any future leak fails the run.
Stream options for dim normalization (ice xT/yT/xB/yB onto ocean conventions), time subsampling to the shared snapshot instants, full-cell regridding, and named postprocess transforms (Kelvin sst, hfds_total_area, sea-ice conventions incl. sea_ice_volume in m^3). UI/VI and HI are zeroed where there is no ice (no time-varying NaN pattern); rotated pairs zero-fill ocean cells with no valid staggered neighbor, keeping the shared-footprint contract exact for tauuo/tauvo.
…mask_k The snapshot sources' instantaneous ocean coverage drifts from the reference-time wetmask: the upstream z-level remap uses instantaneous layer thicknesses, so bottom sliver cells in columns whose depth (deptho + zos) sits near a level interface dry and re-wet with sea level (SSH-driven in 100% of observed events; sub-surface levels only; ~0.004% of wet cells at the worst timestep of the piControl sample year). Regridding each chunk by its own instantaneous footprint therefore produced output cells with mask_k == 1 but NaN, and mask_k == 0 but valid. Training requires NaN exactly where mask_k == 0: a finite target at a masked cell puts NaN into the loss (the loss NaN-mask keys on the target while the prediction is NaN-filled at mask == 0), and a NaN target at an unmasked cell NaNs per-variable metrics and can leak into training histograms. The old 140-yr processed store has zero such cells over its full record, so no trained model has ever seen NaN != mask. Each chunk is now conformed to the wetmask before regridding: values at cells wet at time t but outside the wetmask are dropped, and cells inside the wetmask that are instantaneously dry are filled from the level above (the water immediately overlying the vacated sliver; a wet cell directly above exists at all ~19k dry events in the sample year), via a targeted read of the level-(k-1) slice from the source store. The regrid is normalized by the static wetmask, every variable is asserted to match it exactly per chunk, and the conformed-cell count stays bounded by MAX_FOOTPRINT_DRIFT_FRACTION so an inconsistent source still fails loudly. Verified on a 6-step DirectRunner subset: zero mask/NaN mismatches in both directions across all level-split variables (the same steps previously showed up to 13 mask==1-NaN and 15 mask==0-valid cells), with filled values within 1.3 K of the overlying level. Analysis details and verification are summarized in the PR discussion.
jpdunc23
force-pushed
the
scripts/gfdl-om4-xrbeam-pipeline
branch
from
ea33cd8 to
0359218
Compare
jpdunc23
commented
Jul 8, 2026
| @@ -0,0 +1,145 @@ | |||
| """Regridding utilities ported from the ai2cm fork of ocean_emulators | |||
Member
Author
There was a problem hiding this comment.
Suggested change
| """Regridding utilities ported from the ai2cm fork of ocean_emulators | |
| """Regridding utilities ported from the ai2cm fork of m2lines/ocean_emulators |
Dockerfile, run script, and Makefile targets for running the pipeline on Google Cloud Dataflow, following the era5 operational pattern. Because the pipeline is a package rather than a single file, the worker image copies pipeline/ onto PYTHONPATH instead of relying on --save_main_session; the base image additionally needs libgomp1 for esmf. apache_beam gains the tfrecord extra for fast crc32, matching era5.
The previous Dockerfile created the conda env from the defaults channel and then installed esmf/xesmf with -c conda-forge; that channel mix resolved defaults' _openmp_mutex without an env-local libgomp, so conda-forge's libesmf could not find libgomp.so.1 and needed a system libgomp1 apt workaround. Solving everything from conda-forge with --override-channels puts libgomp in the env (matching the era5 image) and the apt workaround is removed. Also: build_dataflow now smoke-tests the image with an in-container 'import pipeline.run, xesmf' so missing shared libraries or a broken PYTHONPATH fail at build time instead of at Dataflow launch, and the ENV lines use the modern key=value form.
… target Weight artifacts are treated as immutable published versions: generating onto a URL where source_grid.nc or weights.nc already exists now fails fast with instructions to bump the version prefix, so a colleague experimenting with a new config cannot silently replace the weights a production config points at. An explicit --overwrite flag remains for deliberate replacement. Also adds a generate_weights Makefile target covering both target grids and a README setup section documenting the one-time environment and weight-generation steps.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Producing training-reference datasets for ocean emulation currently runs through scripts/data_process/compute_ocean_dataset.py: one giant lazy dask graph over the full simulation executed with xpartition on a hand-managed cluster on LEAP's 2i2c hub. That approach is fragile (whole-graph memory blowups, manual cluster scaling, S3/OSN credential coupling) and doesn't consume the regenerated source data, which now provides true snapshots alongside interval means, a full surface flux/forcing suite, and vertical coarsening to the 19-level Samudra grid already applied upstream.
This PR adds a self-contained pipeline in scripts/gfdl_om4/, modeled on scripts/era5/: an xarray_beam pipeline run on Google Cloud Dataflow (DirectRunner for local subset validation) that reads the native-grid 0.25° tripolar source zarrs, applies per-chunk transforms — vector rotation, wetmask-normalized conservative regridding to an analytic Gaussian grid (F90 1° or F22.5 4°), vertical-level splitting, derived variables, and the established masking/NaN conventions — and writes one templated sharded zarr v3 store per run. Each invocation is driven by a YAML config declaring source stores, streams (variables, renames, transforms), target grid, and output layout, so a new simulation is a config change rather than a code change. Prognostic ocean state comes from snapshot stores with raw timestamps; interval-mean fluxes and ice fields sampled at the snapshot instants share the same single time coordinate; every output variable carries provenance attrs naming its source store and variable. xESMF regridding weights are precomputed once by a setup step and stored as a versioned GCS artifact, and the needed ocean_emulators utilities are ported into the pipeline directory so there is no external dependency on that repo. The first production config is the 5-daily 1° ocean store from the piControl sample year; a 6-hourly sea-ice store and 4° outputs are follow-on configs over the same machinery.
Notes for reviewers
deptho+zos) sits within ~1 m of a level interface dry and re-wet with sea level (SSH-driven at 100% of the 2231 affected columns in the sample year; mostly Arctic shelves). Training requires NaN exactly wheremask_k == 0— a finite target at a masked cell NaNs the fme loss, a NaN target at an unmasked cell NaNs per-variable metrics — and the old 140-yr store satisfies that equality at literally every cell-timestep. Each chunk is therefore conformed to the reference-time wetmask before regridding: newly-wet cells outside the wetmask are dropped, dried cells inside it are filled from the level above (the water immediately overlying the vacated sliver; a wet cell directly above exists at all ~19k observed dry events), the regrid is normalized by the static wetmask, and output NaN pattern ==mask_kis asserted exactly per chunk, with the conformed-cell count still bounded at 0.1% of wet cells so an inconsistent source fails loudly. Analysis and verification:explore2/jamesd/2026-07-06-gfdl-om4-xrbeam-pipeline/02b-time-varying-wetmask.ipynbin the workspace repo.areacellofrom Gaussian quadrature at R=6371 km, ~0.5% above the old polar-radius areas; masked coastal velocity interpolation instead of fillna(0); explicit land-NaN exemption list;sea_ice_volumeinm^3with divisor and units attr agreeing, where the old store divided to km³ but labeled"1000 km^3").tauuo/tauvocarry a stricter staggered-point mask thanuo/vo: after center-interpolation, 1373 tracer-wet coastal cells (~0.14% of surface ocean) have no valid staggered neighbor. Rotated pairs zero-fill those cells (matching the legacy values there; a no-op foruo/vo), keeping the shared-footprint contract exact.UI/VIfollow theHIconvention (zero wheresea_ice_fractionis zero, NaN over land) instead of the old NaN-where-EXT==0 mask —EXTis a binary 15% concentration threshold, so that mask discarded velocities over thin ice and produced a time-varying NaN pattern; and fraction units use the GFDL-style0-1throughout.Changes:
New scripts/gfdl_om4/ directory (pipeline package, configs/, Makefile, environment/requirements, Dockerfile + run-dataflow.sh Dataflow ops, README); no changes to existing code.
Operational details: the worker image COPYs the pipeline package onto PYTHONPATH (the pipeline is a package, so era5's --save_main_session approach doesn't apply) and solves its conda env strictly from conda-forge;
make build_dataflowsmoke-tests imports in-container before any push; weight generation refuses to overwrite an existing versioned artifact unless --overwrite is passed.scripts/data_process/compute_ocean_dataset.py and compute_sea_ice_dataset.py are unchanged.
Tests added
If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated