Add GFDL OM4/CM4 ocean dataset pipeline on xarray_beam + Dataflow by jpdunc23 · Pull Request #1340 · ai2cm/ace

jpdunc23 · 2026-07-04T04:02:26Z

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. Four production configs cover the 5-daily ocean stores for the full 200-year piControl and 140-year 1pctCO2 simulations at 1° and 4°, reading durable input artifacts (statics, weights, face masks) from a dated immutable GCS prefix; a 6-hourly sea-ice store is a follow-on config over the same machinery. Each config has a DirectRunner smoke-test Makefile target (a few timesteps; asserts the wetmask conform step is a no-op and the output opens with the expected variable set), the umbrella suite also verifies the two simulations derive identical wetmasks and that the run driver refuses to initialize into a pre-existing output store, and the README documents the smoke → launch → inspect production checklist.

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_dataflow smoke-tests imports in-container before any push; weight and face-mask generation refuse to overwrite an existing artifact unless --overwrite is passed; per-config smoke_test_*/dataflow_* Makefile targets, with post-run output checks (pipeline/check_output.py) and a cross-simulation wetmask identity check (pipeline/check_wetmask_equivalence.py); the Dataflow temp location is a dated scratch prefix so launches can't pick up a previous generation's staged artifacts.
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

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

…am-pipeline

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 · 2026-07-06T22:06:12Z

Delete prior to merge

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 · 2026-07-08T22:13:05Z

@@ -0,0 +1,145 @@
+"""Regridding utilities ported from the ai2cm fork of ocean_emulators


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.

The 2026-06-24 z*-remapped sources carry coastal staggered velocity faces that are valid with value exactly 0.0 where the native vertical grid masks them as land; the valid-neighbor face->center interpolation averaged these into coastal centers, depressing coastal speeds (the entire 12% coastal mean-speed deficit vs the legacy store was this artifact). - pipeline/face_masks.py (new): one-time setup step scanning a stream's C-grid velocity pair over the config's time window and publishing per-level boolean face masks plus the resulting tracer-center footprint as a versioned GCS artifact. A face is flagged iff it is exactly 0.0 at every scanned step AND has >= 1 dry tracer neighbor (fold-aware on the tripolar seam row). Generation is self-verifying: expected flagged-face census counts, the coastal-geometry identity (flagged faces == faces with exactly one wet tracer neighbor at the surface), and time-static face validity. - Streams opt in via face_mask_url (set for snapshot_ocean). Flagged faces are invalidated before interpolation, with a per-chunk assert that no valid exactly-0.0 face with a dry tracer neighbor survives, so a stale artifact or changed source fails loudly instead of silently depressing coastal speeds again. - Face-masked pairs are restricted and regrid-normalized to the artifact's static center footprint instead of being zero-filled; 86 of 44,892 1-degree surface ocean cells whose contributing centers all lose their valid faces come out NaN. - Every rotated-pair output gets a per-variable mask static (mask_<name>_k for 3D pairs, mask_<name> for 2D) for training-side per-variable masking; face-masked pairs' masks mark their footprint, all others equal the tracer masks. Published artifact om4-picontrol-2026-06-24: 16,960 surface uo and 16,694 vo faces flagged, reconciling exactly with an independent census of the source; a 6-step DirectRunner store matched an independent by-hand corrected regrid to 6e-8 m/s. Intended next revision: return the velocity footprint to the tracer wetmask by restoring the zero-fill fallback for the orphaned centers. After face masking, that fallback fires only at centers whose faces on an axis are all native land (walls), where the model's resolved component is identically zero by no-normal-flow -- so zero is the model's own answer there, the NaN cells disappear, and every per-variable mask equals the tracer masks, without biasing any center that has resolved flow.

A wet tracer center whose staggered faces on an axis are all land is a wall for that axis, where the model's resolved normal component is identically zero by no-normal-flow. Fill that grid-relative component with 0 before rotation — rotation mixes the components, so filling afterwards zeroes a valid along-wall component too (at the surface, every such center is a one-cell-wide channel with real along-channel flow). Every rotated pair then keeps the tracer wetmask footprint, so the per-variable mask statics and the separate face-masked regrid footprint are removed; NaN-equals-mask_k holds for all outputs.

Add the four production configs ({piControl,1pctCO2} x {F90,F22.5}) reading only permanent input artifacts and writing flat dated zarrs in vcm-ml-intermediate. The run driver now refuses to initialize into a pre-existing output store, and a --max-conformed-cells bound lets smoke tests assert the wetmask conform step is a no-op. New check_output and check_wetmask_equivalence scripts back per-config DirectRunner smoke tests and a cross-simulation wetmask identity check, wired into Makefile smoke_test_*/dataflow_* targets; the face-mask generator gains scan-window overrides so per-simulation targets encode their census windows. README documents the smoke -> launch -> inspect checklist.

jpdunc23 · 2026-07-15T20:09:29Z

Notes for reviewers

Deliberate, documented legacy-defect fixes (areacello from 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_volume in m^3 with divisor and units attr agreeing, where the old store divided to km³ but labeled "1000 km^3").
The sources' staggered velocities carry coastal faces that are valid with value exactly 0.0 where MOM6's online z*-remap unmasked land columns; averaged into the C-grid→center interpolation they would bias coastal speeds low. A one-time setup step (pipeline/face_masks.py) scans a source and publishes a per-simulation face-mask artifact (a face is flagged iff exactly 0.0 at every scanned timestep AND it has a dry tracer neighbor, fold-aware on the tripolar seam), verified against an independent census of the source's zero faces (16,960 u / 16,694 v surface faces, identical across piControl and 1pctCO2). Streams opt in via face_mask_url; flagged faces are treated as invalid before interpolation, and a per-chunk assertion fails the run if any face matching the fingerprint survives the masks (a stale artifact or changed source).
Dropping invalid faces means some tracer-wet coastal cells have no valid face on an axis — a wall for that axis, where the model's resolved normal component is identically zero by no-normal-flow. That grid-relative component is set to 0.0 before rotation (rotation mixes the components, so filling after would also zero a valid along-wall component, e.g. the along-channel flow of a one-cell-wide channel). For tauuo/tauvo, whose staggered mask is stricter than uo/vo's, this reproduces the legacy zero values at the 1373 affected surface cells (~0.14% of surface ocean); either way every rotated pair keeps the shared tracer-wetmask footprint exactly.
Two legacy sea-ice conventions were deliberately changed: UI/VI follow the HI convention (zero where sea_ice_fraction is zero, NaN over land) instead of the old NaN-where-EXT==0 mask — EXT is 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-style 0-1 throughout.

spencerkclark

A few minor comments as I gave this an initial rough read. No red flags stood out so far. It's great to have this workflow so well organized / easy to run now.

spencerkclark · 2026-07-13T19:40:15Z

+
+
+def _interpolate_right_to_center(
+    da: xr.DataArray, dim: str, center_dim: str, periodic: bool


nit: maybe call the input right for maximum clarity

Suggested change

da: xr.DataArray, dim: str, center_dim: str, periodic: bool

right: xr.DataArray, dim: str, center_dim: str, periodic: bool

spencerkclark · 2026-07-13T19:44:26Z

+    """
+    if da.sizes[dim] != da[dim].size:
+        raise ValueError(f"Expected coordinate for staggered dim {dim}")
+    left = da.roll({dim: 1}, roll_coords=False) if periodic else da.shift({dim: 1})


It could be worth clarifying in a comment or docstring that the non-periodic case simply uses a row of NaNs to form the first "left" boundary, and that this is natural in the case of MOM6, since the southern boundary is Antarctica, which is land and therefore all NaN.

spencerkclark · 2026-07-15T20:35:55Z

+  The source footprint drifts slightly from the reference-time wetmask
+  (z*-remapped bottom slivers dry/re-wet with sea level at a handful of
+  sub-surface cells), so each chunk is 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). See _conform_to_wetmask.


Fine to keep _conform_to_wetmask—it's great that it facilitates this strict check even if we do not expect to use its primary funcationality—but reminder to update this section of the docstring.

spencerkclark · 2026-07-15T20:53:51Z

+                # Symmetric staggering carries both edges; drop the first
+                # point to match the right/north-edge convention of the
+                # shared interpolation and rotation machinery.


Indeed awkward that the sea ice and ocean outputs follow different conventions here...

spencerkclark · 2026-07-15T21:28:05Z

+# drift is ~0.01% of wet cells. Anything larger indicates an inconsistent
+# source (wrong grid, dropped mask) and fails the run rather than being
+# silently conformed.
+MAX_FOOTPRINT_DRIFT_FRACTION = 0.001


Oh so this is still giving some lenience to a changing mask when max_conformed_cells is None?

spencerkclark · 2026-07-15T22:11:04Z

+    for start in range(0, n_steps, TIME_BATCH):
+        batch = da.isel({time_dim: slice(start, start + TIME_BATCH)}).load().values
+        finite = np.isfinite(batch)
+        zero = finite & (batch == 0.0)
+        finite_all = (
+            finite.all(axis=0)
+            if finite_all is None
+            else (finite_all & finite.all(axis=0))
+        )
+        finite_any = (
+            finite.any(axis=0)
+            if finite_any is None
+            else (finite_any | finite.any(axis=0))
+        )
+        zero_all = (
+            zero.all(axis=0) if zero_all is None else (zero_all & zero.all(axis=0))
+        )
+        logger.info(
+            "  %s: scanned steps %d..%d of %d",
+            da.name,
+            start,
+            min(start + TIME_BATCH, n_steps) - 1,
+            n_steps,
+        )


nit: I think these operations should all be possible with xarray objects. I would find it easier to follow with labeled dimensions.

spencerkclark · 2026-07-15T22:20:59Z

Oof a lot of work to address this issue. We should find out why it's the case in MOM6 and whether they would be open to addressing it.

jpdunc23 added 4 commits July 3, 2026 20:35

Merge remote-tracking branch 'origin/main' into scripts/gfdl-om4-xrbe…

aede035

…am-pipeline

jpdunc23 commented Jul 6, 2026

View reviewed changes

Comment thread pr-plan.md

jpdunc23 Jul 6, 2026

Copy link
Copy Markdown

Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Delete prior to merge

jpdunc23 added 4 commits July 6, 2026 16:30

jpdunc23 force-pushed the scripts/gfdl-om4-xrbeam-pipeline branch from ea33cd8 to 0359218 Compare July 7, 2026 17:03

jpdunc23 commented Jul 8, 2026

View reviewed changes

jpdunc23 added 6 commits July 9, 2026 14:09

spencerkclark reviewed Jul 15, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add GFDL OM4/CM4 ocean dataset pipeline on xarray_beam + Dataflow#1340

Add GFDL OM4/CM4 ocean dataset pipeline on xarray_beam + Dataflow#1340
jpdunc23 wants to merge 14 commits into
mainfrom
scripts/gfdl-om4-xrbeam-pipeline

jpdunc23 commented Jul 4, 2026 •

edited

Loading

Uh oh!

jpdunc23 Jul 6, 2026

Uh oh!

jpdunc23 Jul 8, 2026

Uh oh!

jpdunc23 commented Jul 15, 2026

Uh oh!

spencerkclark left a comment

Uh oh!

spencerkclark Jul 13, 2026

Uh oh!

spencerkclark Jul 13, 2026

Uh oh!

spencerkclark Jul 15, 2026

Uh oh!

spencerkclark Jul 15, 2026

Uh oh!

spencerkclark Jul 15, 2026

Uh oh!

spencerkclark Jul 15, 2026

Uh oh!

spencerkclark Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -0,0 +1,145 @@
		"""Regridding utilities ported from the ai2cm fork of ocean_emulators



		def _interpolate_right_to_center(
		da: xr.DataArray, dim: str, center_dim: str, periodic: bool

	da: xr.DataArray, dim: str, center_dim: str, periodic: bool
	right: xr.DataArray, dim: str, center_dim: str, periodic: bool

Uh oh!

Conversation

jpdunc23 commented Jul 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jpdunc23 commented Jul 15, 2026

Notes for reviewers

Uh oh!

spencerkclark left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jpdunc23 commented Jul 4, 2026 •

edited

Loading