Per-column z ingest for ERA5 pressure-level data by glwagner · Pull Request #241 · NumericalEarth/NumericalEarth.jl

glwagner · 2026-05-14T21:22:52Z

Summary

PressureLevelVerticalDiscretization + PressureLevelGrid (new src/Grids submodule). Stores raw geopotential Φ(λ,φ,p) (m²/s²) and gravitational_acceleration; rnode(i,j,k,grid) reads Φ[i,j,k] / g. _fractional_indices is specialized so Oceananigans.interpolate! bisects each column's actual heights inside the kernel, instead of the 1-D znodes(grid, loc) vector that the default assumes.
ERA5 pressure-level ingest is per-column by default. z_interfaces(::ERA5PressureMetadata) builds the discretization from the instantaneous Φ and clips at the local surface geopotential via clip_subsurface!. No new dataset keyword — old z_mode / mean_geopotential_height plumbing is gone.
znodes dispatch on AbstractField. Returns the column-mean Vector{Float64} when the field is horizontally absent (Flat x/y topology or Nothing horizontal locations from mean(...; dims=(1,2))); otherwise a 3-D per-cell Field. Same logic applies to FieldTimeSeries, so mean(znodes(T_series), dims=(1,2)) yields a Field{Nothing, Nothing, Center} as expected.
CDS request aggregation. Single-variable and multi-variable ERA5 downloads now batch by (year, month) (using CDS's Cartesian year/month/day/time semantics) and further cap each batch at 5000 fields per request to avoid HTTP 403 cost-limit errors on pressure-level multi-var pulls.
Example: examples/ERA5_hourly_data.jl. Adds a §4 that animates ERA5 specific humidity downscaled 4× from the 0.25° native grid onto a fixed-altitude target, illustrating the per-column-z interpolation end-to-end. Plotting routes through the Oceananigans Makie ext (Fields and AbstractOperations as plot args; Observable{<:Field} for animation). The §4 frame loop reuses the §3c q_series FieldTimeSeries via interpolate!, so no per-frame metadata lookup or grid reconstruction.
Example size tuned for CI. Date range trimmed from 7 days (Dec 27–Jan 2) to 4 days entirely in December, rico_region from 1.5°×1.5° to 1°×1°. Cold-cache local repro: 53 min → 30 min.

Closes #236.

Why

set!(field, era5_pressure_level_metadatum) previously placed the source ERA5 field on a LatitudeLongitudeGrid whose z-axis was a single 1-D vector of ⟨Φ⟩/g averaged over space and time. Over real terrain (sub-surface ERA5 fill leaks through that flat profile) and under synoptic Φ-surface displacement (~40 m for 5 hPa SLP gradients — comparable to LAM near-surface Δz), that profile is wrong by O(50–500 m) in cell placement. See #236 for details and the design write-up.

Design highlights

No advection-side changes. The discretization participates only in rnode and _fractional_indices; advection, halo, and model-side kernels never touch it because pressure-level grids aren't used as model grids.
Sub-surface clip at build time. clip_subsurface! rewrites Φ[i,j,k] ← max(Φ[i,j,k], Φ_sfc[i,j]) in place and refills halos. Columns are guaranteed monotone in k, so column_fractional_z_index needs no branching and the kernel stays clean.
Surface-Φ paired by dispatch. _matching_single_level_dataset(::ERA5HourlyPressureLevels) = ERA5HourlySingleLevel() and the monthly counterpart, so the surface-clip source matches the cadence of the data being clipped.
znodes(::Field) / znodes(::FieldTimeSeries) route via _znodes_for_plg. Two concrete-type methods (one on Field{…<:PressureLevelGrid}, one on FieldTimeSeries{…<:PressureLevelGrid}) — each more specific than Oceananigans' own znodes defaults on their respective wrappers, so dispatch is unambiguous.
CDS batching helper _batch_datetimes_for_cds. Groups by (year, month) first (preserving CDS's Cartesian product), then chunks each month into ≤ 5000 / (num_vars × num_levels) datetimes. Recovers most of the daily→monthly speedup for small payloads (§2/§3a: 2.5–3.3×) while staying safely under CDS's pressure-level cost limit.

What this defers

FTS/TimeSeriesInterpolation-backed Φ for forecast-time use. The discretization is already structurally compatible (its geopotential field accepts a TSI), but ERA5PrescribedAtmosphere and the geopotential field on PrescribedAtmosphere aren't here — separate PR.
Per-donor-column bisection in _interpolate. column_fractional_z_index bisects a single representative column — the cell at the integer indices clamp(unsafe_trunc(Int, ii), 1, Nx), clamp(unsafe_trunc(Int, jj), 1, Ny). The bilinear blend in Oceananigans.interpolate! then weights values at the neighbouring (i+1, j), (i, j+1), (i+1, j+1) columns with the same (k_low, k_high), even though those columns have slightly different z values at the same k. Near sharp orographic gradients this biases the result. Tractable as a kernel-level follow-up.
FieldTimeSeries support in OceananigansMakieExt — heatmap!(ax, fts) would let the §3 Hovmöller drop its manual (Nt × Nz) matrix assembly. Tracked upstream as CliMA/Oceananigans.jl#5599.

Test plan

CPU unit tests (test/test_pressure_level_grid.jl, 58 tests): PressureLevelVerticalDiscretization constructor, generate_coordinate dim/axis guards, clip_subsurface! behavior on a Field-backed Φ, grid-level znodes/rnodes agreement, znodes(::Field) per-cell vs column-mean dispatch on all four cases (horizontally resolved, Flat-Flat topology, Nothing-Nothing location, FieldTimeSeries parity).
CDS batching unit tests (test/test_cds_downloading.jl): _group_by_calendar_month, _batch_datetimes_for_cds (single-level monthly fits in one batch; pressure-level many-vars splits within a month; _max_dts_per_cds_request arithmetic; chronological ordering under shuffled input).
ExplicitImports QA after the Adapt import cleanup in src/Grids/pressure_level_vertical_discretization.jl.
End-to-end ERA5 example builds in CI Documenter (was hanging 7h+ before this PR's aggregations; now finishes in ~46 min).
Hurricane Helene demo (out-of-tree): 37-frame animation of ERA5 wind speed at z = 1.5 km over a 15°×18° SE-US bbox, downscaled 4×. Storm structure (eye, spiral, rapid translation) clearly resolved through landfall and inland decay.
Run examples/era5_breeze.jl and confirm field ranges match the previous two-stage workaround to within rectification bias — out of scope for this PR.

🤖 Generated with Claude Code

Adds `PressureLevelVerticalDiscretization` (new `src/Grids` submodule), a 3-D vertical coordinate backed by raw geopotential `Φ(λ,φ,p)` (m²/s²) and a `gravitational_acceleration` scalar. `rnode(i,j,k,grid)` reads `Φ[i,j,k] / g`, and `_fractional_indices` is specialized for the `PressureLevelGrid` alias to bisect each column's actual heights instead of the 1-D `znodes(grid, loc)` vector that the default kernel assumes. `ERA5HourlyPressureLevels` / `ERA5MonthlyPressureLevels` gain a tri-state `z_mode` keyword (`:per_column` default, `:mean`, `:standard_atmosphere`), replacing the boolean `mean_geopotential_height` with a deprecation shim. In `:per_column` mode, `z_interfaces` returns the new discretization built from the instantaneous Φ and clipped at the local surface geopotential (`ERA5HourlySingleLevel` / `ERA5MonthlySingleLevel` paired via dispatch on the pressure-level dataset type). `examples/ERA5_hourly_data.jl` gains a final §4 section that animates ERA5 specific humidity at z = 800 m over the RICO bbox, downscaled 4× from the 0.25° native grid — the entire ingest is one `set!(field_2d, metadatum)` per frame. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

codecov · 2026-05-14T21:52:10Z

Codecov Report

❌ Patch coverage is 53.94737% with 70 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/Grids/pressure_level_vertical_discretization.jl	25.30%	62 Missing ⚠️
src/DataWrangling/ERA5/ERA5_pressure_levels.jl	80.00%	5 Missing ⚠️
src/DataWrangling/metadata_field.jl	72.72%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

`z_mode` is removed entirely. The only honest mode is per-column, and `ERA5HourlyPressureLevels`/`ERA5MonthlyPressureLevels` always produce it when `z = nothing` (the default). Users who explicitly want a static profile pass `z = standard_atmosphere_z_interfaces(pressure_levels)` (or any other vector / discretization) directly. `PressureLevelVerticalDiscretization` no longer carries `cᵃᵃᶠ`/`cᵃᵃᶜ`/ `Δᵃᵃᶠ`/`Δᵃᵃᶜ` placeholders — only `gravitational_acceleration` and `geopotential`. `generate_coordinate` computes `grid.Lz` from `extrema(geopotential) / g` so it reports the actual physical extent (~47 km for ERA5's full level set) instead of a meaningless `Nz`. `znodes`/`rnodes` (the 1-D plural forms) error explicitly on a `PressureLevelGrid`: there is no representative 1-D z. Per-cell `znode(i, j, k, grid)` still works and is what `_fractional_indices` and downstream interpolation use. Also pulls in the `:geopotential` / `:geopotential_height` entries on `ERA5HourlySingleLevel`'s variable-name maps (previously only on the era5_breeze_land branch) so that the surface-clip path in `per_column_geopotential_discretization` actually resolves. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Match Oceananigans' existing `rnodes(::AUG, ...)` signatures one-to-one with strictly-more-specific dispatch on `PressureLevelGrid`. Previously the generic AUG fallback's vararg + my vararg fallback created ambiguity, and `znodes(grid, Center())` failed with `MethodError`; now it throws a clear `ArgumentError` pointing users at `znode(i, j, k, grid)`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…eanup - Parameterize `ERA5HourlyPressureLevels{Z}` and `ERA5MonthlyPressureLevels{Z}` by the `z` field's type instead of carrying `::Any`. - `clip_subsurface!(::Field, Φ_sfc)` now launches a KA kernel via `launch!` so it works on GPU fields. CPU loop kept for the TSI path (not exercised). - `rnodes(grid::PressureLevelGrid, ℓx, ℓy, ℓz)` returns a `KernelFunctionOperation` wrapping per-cell `rnode`, so callers can materialize the 3-D z-field via `Field(rnodes(...))`. The 1-loc form still throws an ArgumentError (no 1-D answer exists). - Reuse Oceananigans' `index_binary_search` via a tiny `_ColumnView` wrapper instead of reimplementing the bisection. - Fix docstring placement (was attached to the @kernel definition instead of the public `clip_subsurface!`). Drop stale references to removed `z_mode` and `column_index_binary_search`. - Remove alignment-style padding around `=` and `::`. Reorder local helpers (`_ColumnView`, `_znode_op`) to precede their use. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

glwagner · 2026-05-15T19:14:57Z

here's a movie that can be created after this PR, atmos state interpolated to a constant z-level

helene_animation_demo.mp4

The `mean_geopotential_height::Bool` field was removed when `z_mode` was dropped in favor of `z = nothing` triggering the per-column path. Update the constructor-sorting test accordingly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Constrain the 3-tuple `_fractional_indices` method to `NTuple{3, Any}` so it no longer matches the 1-tuple (z,) call that Oceananigans uses for Flat-Flat-Bounded grids (a Column-region source). Add a 1-tuple method that bisects the single (1,1) column directly. Fixes docs build, which evaluates §3 of the ERA5 tutorial — the RICO column path that loads `:specific_cloud_liquid_water_content` over a `Column` region. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

A `znodes(field)` call on a Column-region field (Field{Nothing, Nothing, Center} on a Flat-Flat-Bounded grid) routed to the 3-loc `rnodes` override and returned a `KernelFunctionOperation`, which Makie's heatmap recipe rejected — fails docs build for §3 of the ERA5 tutorial. Add an explicit `rnodes(grid, ::Nothing, ::Nothing, ℓz)` that materializes the single (1, 1) column's heights as a `Vector{Float64}`. The 3-D form (non- Nothing ℓx/ℓy) still returns a lazy `KernelFunctionOperation`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The previous KernelFunctionOperation return from `rnodes(grid, ℓx, ℓy, ℓz)` broke any consumer expecting a 1-D plot axis or `length`-able vector — including `znodes(T_series[1])` in §3 of the ERA5 tutorial, which then hits a DimensionMismatch when allocating `zeros(Nz, Nt)` (since `length` on the 3-D KFO is Nx·Ny·Nz, not Nz). Switch every `rnodes(grid::PressureLevelGrid, ...)` form to return the column-mean z profile as a plain `Vector{Float64}` of length `Nz`. This matches what the old `:mean` ingest mode produced and what plot recipes expect. Per-cell access is unchanged: `znode(i, j, k, grid)` still reads `geopotential[i, j, k] / g`. Users who want a lazy 3-D z-field can build one explicitly from `_znode_op` (no longer in the public surface — drop the import since nothing in the package needs it). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…umericalEarth/NumericalEarth.jl into eq/per-column-z-discretization

Addresses review feedback on PR #241: prefer `view` semantics over materializing into a fresh `Array`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

- Drop `pressure_level_heights` helper: `znodes(field)` already returns the column-mean `Vector{Float64}` for `PressureLevelGrid` (commit 53f1786). `Field(znodes(field))` doesn't exist as a constructor. - §1: pass the wind/Stokes `AbstractOperation`s directly to `heatmap!`; the Makie ext converts and supplies node coordinates. - §4: hold the `CenterField` itself in the `Observable` and let the ext handle the update path, instead of unpacking `interior(q2d)[:, :, 1]`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

A field on a PressureLevelGrid has per-cell heights — there's no single 1-D z axis at the Field level. Override `znodes(::Field)` to return a materialized 3-D `Field` of per-cell z (wrapping a `KernelFunctionOperation` over `rnode`), keeping the field's own location. The grid-level `znodes(grid, ℓ...)` / `rnodes(grid, ℓ...)` still return the column-mean `Vector{Float64}` so plot recipes, `Lz`, and length consumers that only have the grid in hand continue to work. Example: collapse to a 1-D axis via `vec(znodes(field))` for column (1×1×Nz) fields, or `vec(mean(znodes(field), dims=(1, 2)))` for horizontally extended fields. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Same numerical result for the (1×1×Nz) Column case (it's a no-op mean), but matches the profiles-plot pattern and states intent clearly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`znodes(grid, Center(), Center(), Center())` already returns the column-mean `Vector{Float64}` we want for the plot axis, so we don't need to round-trip through `vec(mean(znodes(field), dims=(1, 2)))`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`znodes(grid, nothing, nothing, Center())` is the semantically honest way to ask for the horizontally-averaged z profile — `Nothing` slots say "no horizontal dimension," which is what a column-mean axis is. The existing `rnodes(::PressureLevelGrid, ℓz::Center)` 1-arg override already handles this routing (Oceananigans' default 3-arg `rnodes` forwards to the 1-arg form). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Per the convention reaffirmed on PR #241: the leading `_` is reserved for KernelAbstractions `@kernel` functions and `func`/`_func` layered interfaces. Everything else is module-private by virtue of not being `export`ed. Renames in this PR's new code: ext/NumericalEarthCDSAPIExt.jl _CDS_MAX_FIELDS_PER_REQUEST → CDS_MAX_FIELDS_PER_REQUEST _group_by_calendar_month → group_by_calendar_month _max_dts_per_cds_request → max_dts_per_cds_request _batch_datetimes_for_cds → batch_datetimes_for_cds src/Grids/pressure_level_vertical_discretization.jl _geopotential_data_for_extrema → geopotential_data_for_extrema _mean_height_profile → mean_height_profile _znode_op → znode_op _znodes_for_plg → znodes_for_plg _ColumnView → ColumnView test/test_pressure_level_grid.jl _make_plvd → make_plvd _make_plg → make_plg Kept: _clip_subsurface_kernel! (a `@kernel` function) _fractional_indices, fractional_x_index, … (Oceananigans' names — imported, not introduced by this PR) All 58 PressureLevelVerticalDiscretization tests and 320 CDSAPIExt tests still pass locally. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Force-push of MetadataSet export fix didn't auto-trigger CI.

Co-authored-by: Eliot Quon <eliot@aeolus.earth>

Neither type is referenced unqualified from outside the `Grids` submodule — `src/DataWrangling/ERA5/…` already imports them via `using NumericalEarth.Grids: PressureLevelVerticalDiscretization`, and `PressureLevelGrid` only appears in doc cross-refs and a comment in the example. Resolves ewquon's question on PR #241. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`parent(fts)` returns a 4-D OffsetArray that includes halo cells. For default FieldBoundaryConditions those halos are zero, which collapsed both `extrema` (driving `Lz`) and the per-level `mean` in `mean_height_profile` toward zero. ewquon's repro on PR #241: znodes(grid, Center()) → [0.0, 0.0, 0.0, 480.0] expected → [3000.0, 6000.0, 9000.0, 12000.0] Use `interior(fts)` instead, which already excludes halos on every spatial dim while preserving the trailing time axis — matching the `Field` overload one line up. Adds a regression test asserting the time-mean column-mean and `Lz` on a TSI-backed PLVD. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

- Define `Base.show(io, grid::LatitudeLongitudeGrid{…<:PLVD}, withsummary=true)`. The default `LatitudeLongitudeGrid` show reads `grid.z.cᵃᵃᶠ`, which `PressureLevelVerticalDiscretization` doesn't carry — typing the grid name at the REPL raised `FieldError: PressureLevelVerticalDiscretization has no field cᵃᵃᶠ`. The override prints the horizontal axes, size/halo/Lz, and delegates the z line to PLVD's own `show`. Targeting `LLG{…<:PLVD}` (concrete wrapper + constrained `z` parameter) is strictly more specific than either the LLG default or `show(::PressureLevelGrid)`, so dispatch is unambiguous. Resolves ewquon on PR #241. - Drop the unused `locs` arg from `column_fractional_z_index`; callers in both `_fractional_indices` overloads stop passing it. Resolves ewquon on PR #241. Regression test asserts that both `sprint(show, grid)` and `sprint(show, MIME"text/plain"(), grid)` complete without `FieldError` and include the PLVD description. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`test_quality_assurance.jl` (CPU CI on Julia 1.12.5) flagged two issues in `NumericalEarth.Grids` introduced by this PR: ExplicitImportsFromNonOwnerException: `instantiated_location` has owner `Oceananigans` but it was imported from `Oceananigans.Fields`. StaleImportsException: `AbstractField` is imported but unused (the type was used in an earlier draft of `_znodes_for_plg`; switched to two concrete dispatch methods on `Field` / `FieldTimeSeries` in 8664a5d and never cleaned up). Move `instantiated_location` to the top-level `Oceananigans` import and drop the unused `AbstractField`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`ext/NumericalEarthCDSAPIExt.jl` does `using Dates: Dates` (module-only import per the ExplicitImports convention), so `DateTime` must be written as `Dates.DateTime`. The docs build picked up the unqualified reference in `batch_datetimes_for_cds`: ERROR: LoadError: LoadError: UndefVarError: `DateTime` not defined in `NumericalEarthCDSAPIExt` Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

glwagner · 2026-05-20T12:45:06Z

merging this because the CI errors are unrelated.

Resolves conflicts from PR #241 (per-column z ingest) landing on main: - src/DataWrangling/ERA5/ERA5.jl: take main's PressureLevelVerticalDiscretization import, keep our download_dataset → download rename. - src/DataWrangling/ERA5/ERA5_pressure_levels.jl: take main's wholesale rewrite (per_column_geopotential_discretization, parametric {Z} types), adapt the two download_dataset calls to use the renamed download. - examples/ERA5_hourly_data.jl: keep our lower-troposphere bullet, use download instead of download_dataset. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

glwagner added atmosphere ⛈️ where the weather happens data wrangling 🗃️ JRA55, ECCO, ERA5, and friends labels May 14, 2026

glwagner requested a review from ewquon May 14, 2026 21:32

glwagner and others added 4 commits May 14, 2026 22:06

Merge branch 'main' into eq/per-column-z-discretization

ed160ba

glwagner and others added 8 commits May 15, 2026 21:44

Merge branch 'main' into eq/per-column-z-discretization

885d969

Fix ERA5 docs pressure-level heights

18a2a8b

Merge branch 'eq/per-column-z-discretization' of https://github.com/N…

b98d35e

…umericalEarth/NumericalEarth.jl into eq/per-column-z-discretization

Merge branch 'main' into eq/per-column-z-discretization

8486c74

glwagner requested a review from navidcy May 16, 2026 03:04

Merge branch 'main' into eq/per-column-z-discretization

74bc0a3

navidcy reviewed May 16, 2026

View reviewed changes

Comment thread examples/ERA5_hourly_data.jl Outdated

Merge branch 'main' into eq/per-column-z-discretization

1b6aab6

Use a view of pressure-level heights instead of copying

7d44c3f

Addresses review feedback on PR #241: prefer `view` semantics over materializing into a fresh `Array`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

glwagner force-pushed the eq/per-column-z-discretization branch from 8e94880 to 7d44c3f Compare May 16, 2026 13:55

glwagner and others added 5 commits May 16, 2026 08:09

Make z_col use the explicit column-mean form

fbc3b5d

Same numerical result for the (1×1×Nz) Column case (it's a no-op mean), but matches the profiles-plot pattern and states intent clearly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>