From 2084588680bfb422c5b5784333aa809ed3392b17 Mon Sep 17 00:00:00 2001 From: Eliot Quon Date: Fri, 28 Aug 2026 20:10:42 -0600 Subject: [PATCH] Materialize a time-varying pressure-level geopotential once per state update MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `column_fractional_z_index` bisects a column of the geopotential for every cell it interpolates. When the geopotential is a `TimeSeriesInterpolation` — the default for `ERA5HourlyPressureLevels()`, which passes `z = nothing` — each of those probes was a `FieldTimeSeries[i, j, k, Time(t)]`, carrying its own binary search over `times` and a two-slice blend. Within a step every probe asks for the same Φ(t), so that work was redundant. Materialize it instead. The discretization now always stores a plain `Field`; a `TimeSeriesInterpolation` is materialized at construction and recomputed by `materialize_geopotential!` from `update_state!(::PrescribedAtmosphere)`, which `NestedModel.time_step!` reaches on every nested step. Each probe is a single array read again. Implemented as a `Field` type alias rather than a new struct, so `rnode`, `ColumnView`, `first_above_surface_level`, `Adapt` and `on_architecture` keep working unchanged. `geopotential_data_for_extrema` forwards through `.operand` so `Lz` and the column-mean profile still span every time slice, not just the resident snapshot. Clipping runs before materialization, so every snapshot inherits the clipped columns. Microbenchmark, paired and interleaved in one process, parent 40x40x37 -> child 150x136x50, output bit-identical: interpolate! end-to-end 121.5 ms -> 57.4 ms 2.12x column_fractional_z_index 109.3 ms -> 20.4 ms 5.36x materialize_geopotential! 0.53 ms/tick End-to-end on the ERA5 -> 12 km nested hindcast (`examples/breeze_downscaling_era5.jl`, 150x136x50, Float32), A100-SXM4-40GB, exclusive, 400 iterations, arms alternated A/B/A/B against 2eec75b2b: main 0.0478, 0.0495 s/step this 0.0368, 0.0359 s/step 1.34x Within-arm spread is 0.0017 and 0.0009 s against a between-arm gap of 0.0123 s, and the two arms' ranges do not overlap. Physics envelopes agree: rho in [0.0947, 1.1857], max|u| 98.6-101.9 m/s across all four runs. The accuracy cost is zero rather than small. The parent clock advances only inside `time_step!(atmos, Δt)`, which calls `update_state!` immediately, and the child reads `znode` between parent ticks, so the snapshot always holds Φ at the parent clock time. Over 500 steps of ERA5 data, max |materialized - interpolated| is 0 exactly. One behavior change: reading `rnode` after hand-setting `clock.time` returns the previous snapshot until a refresh runs. The existing clock-following test is updated to assert that. Note also that `FieldTimeSeries` defaults to `Clamp()` time indexing, so past the end of the window the heights freeze on the last snapshot rather than extrapolating; that is now pinned by a test. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YHo11y7uLQQfPZjo3e12tx --- src/Atmospheres/prescribed_atmosphere.jl | 4 + .../ERA5/ERA5_pressure_levels.jl | 5 +- .../pressure_level_vertical_discretization.jl | 46 +++- test/test_pressure_level_grid.jl | 219 +++++++++++++++++- 4 files changed, 264 insertions(+), 10 deletions(-) diff --git a/src/Atmospheres/prescribed_atmosphere.jl b/src/Atmospheres/prescribed_atmosphere.jl index 4ddf250f3..06f83e6c3 100644 --- a/src/Atmospheres/prescribed_atmosphere.jl +++ b/src/Atmospheres/prescribed_atmosphere.jl @@ -158,6 +158,10 @@ end for fts in ftses update_field_time_series!(fts, time) end + + # A no-op unless the grid's vertical is a time-varying geopotential. + NumericalEarth.Grids.materialize_geopotential!(atmos.grid) + return nothing end diff --git a/src/DataWrangling/ERA5/ERA5_pressure_levels.jl b/src/DataWrangling/ERA5/ERA5_pressure_levels.jl index 449ea1912..4359f702b 100644 --- a/src/DataWrangling/ERA5/ERA5_pressure_levels.jl +++ b/src/DataWrangling/ERA5/ERA5_pressure_levels.jl @@ -254,8 +254,9 @@ end Build a [`PressureLevelVerticalDiscretization`](@ref) whose per-column heights follow the ERA5 geopotential in time. The geopotential Φ(λ,φ,p,t) is loaded over the whole date window of `metadata` (all snapshots in memory) and wrapped in a -`TimeSeriesInterpolation` bound to `clock`; `znode` reads Φ interpolated to -`clock.time` and divides by `g`. Sub-surface levels are clipped to the local +`TimeSeriesInterpolation` bound to `clock`, which the discretization materializes +into a `Field` refreshed at each `update_state!(::PrescribedAtmosphere)`; `znode` +reads that snapshot and divides by `g`. Sub-surface levels are clipped to the local surface geopotential so that columns stay monotonic. The single-level surface geopotential (orography × g, static in time) is downloaded from the matching `ERA5*SingleLevel` dataset and used as the clip source. diff --git a/src/Grids/pressure_level_vertical_discretization.jl b/src/Grids/pressure_level_vertical_discretization.jl index b828ee303..0c7d49bfd 100644 --- a/src/Grids/pressure_level_vertical_discretization.jl +++ b/src/Grids/pressure_level_vertical_discretization.jl @@ -34,9 +34,12 @@ Two `znodes` paths: locations, e.g. after `mean(...; dims=(1, 2))`); otherwise it returns a 3-D `Field` of per-cell heights at the field's own location. -`geopotential` (units m²/s²) is a 3-D `Field` or a `TimeSeriesInterpolation` -over a `FieldTimeSeries`. The former gives a static z-coordinate; the latter -gives a time-evolving one driven by an attached `Clock`. +`geopotential` (units m²/s²) is always a 3-D `Field`. Passing a `Field` to the +constructor gives a static z-coordinate. Passing a `TimeSeriesInterpolation` +over a `FieldTimeSeries` gives a time-evolving one driven by an attached +`Clock`: the constructor materializes it into a `Field` that +[`materialize_geopotential!`](@ref) refreshes from +`update_state!(::PrescribedAtmosphere)`. The `LatitudeLongitudeGrid` constructor needs a value for `Lz`; we compute it as `extrema(geopotential) / g` inside `generate_coordinate`. @@ -64,12 +67,18 @@ sub-surface *data* stays on those levels, so `column_fractional_z_index` also clamps interpolation to the first above-ground level (never sampling it). The clip source is retained on the discretization and exposed through [`surface_elevation`](@ref). + +A `TimeSeriesInterpolation` `geopotential` is clipped before it is materialized, +so every snapshot the discretization goes on to compute inherits the clipped +columns. """ function PressureLevelVerticalDiscretization(geopotential; gravitational_acceleration, surface_geopotential = nothing) isnothing(surface_geopotential) || clip_subsurface!(geopotential, surface_geopotential) - return PressureLevelVerticalDiscretization(gravitational_acceleration, geopotential, surface_geopotential) + return PressureLevelVerticalDiscretization(gravitational_acceleration, + geopotential_field(geopotential), + surface_geopotential) end # Skip the generic validator (which would `length`-check the missing 1-D fields). @@ -95,12 +104,26 @@ function generate_coordinate(FT, topo, sz, halo, return Lz, arch_discretization end +""" + MaterializedGeopotential + +A geopotential `Field` computed from a `TimeSeriesInterpolation`: Φ at the time +of the last [`materialize_geopotential!`](@ref), stored as plain data, with the +source series still reachable through the field's `operand`. +""" +const MaterializedGeopotential = Field{<:Any, <:Any, <:Any, <:TimeSeriesInterpolation} + +geopotential_field(Φ::Field) = Φ +geopotential_field(Φ::TimeSeriesInterpolation) = Field(Φ) + geopotential_data_for_extrema(Φ::Field) = interior(Φ) # Use `interior(fts)` — not `parent(fts)` — so halo zeros don't dominate the # extrema / column mean. For a TSI this reads every time slice, so `Lz` spans the # heights reached at any time in the window — the extent a child must be able to # bracket over the whole run. One-time setup, not a hot path. geopotential_data_for_extrema(Φ::TimeSeriesInterpolation) = interior(Φ.time_series) +# A materialized snapshot holds one time; read the source series behind it instead. +geopotential_data_for_extrema(Φ::MaterializedGeopotential) = geopotential_data_for_extrema(Φ.operand) Adapt.adapt_structure(to, z::PressureLevelVerticalDiscretization) = PressureLevelVerticalDiscretization(z.gravitational_acceleration, @@ -126,6 +149,21 @@ Type alias for any underlying grid whose vertical coordinate is a const PressureLevelGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:PressureLevelVerticalDiscretization} +""" + materialize_geopotential!(grid) + +Recompute the geopotential snapshot that a [`PressureLevelGrid`](@ref) with a +time-varying vertical reads, at the current time of the `Clock` its +`TimeSeriesInterpolation` is bound to. `znode` returns Φ at the time of the last +refresh, not at the instant of access; `update_state!(::PrescribedAtmosphere)` +refreshes. + +A no-op for any other grid, and for a static-`Field` geopotential. +""" +materialize_geopotential!(grid) = nothing +materialize_geopotential!(grid::PressureLevelGrid) = materialize_geopotential!(grid.z.geopotential) +materialize_geopotential!(Φ::MaterializedGeopotential) = (compute!(Φ); nothing) + # Override the LLG show, which reads `grid.z.cᵃᵃᶠ` and crashes on PLVD. # We print the horizontal axes and report the z coordinate via the # `PressureLevelVerticalDiscretization` `show` defined above. Targeting the diff --git a/test/test_pressure_level_grid.jl b/test/test_pressure_level_grid.jl index e15da7800..7e29fdd3b 100644 --- a/test/test_pressure_level_grid.jl +++ b/test/test_pressure_level_grid.jl @@ -1,13 +1,16 @@ include("runtests_setup.jl") using Oceananigans +using Oceananigans.BoundaryConditions: fill_halo_regions! using Oceananigans.Fields: instantiated_location using Oceananigans.Grids: Flat, Bounded, topology using Oceananigans.OutputReaders: TimeSeriesInterpolation +using Oceananigans.TimeSteppers: time_step! using Statistics using NumericalEarth.Grids: PressureLevelGrid, PressureLevelVerticalDiscretization, - column_fractional_z_index + column_fractional_z_index, materialize_geopotential!, + surface_elevation # Build a small static-Field-backed `PressureLevelVerticalDiscretization` from # a per-cell geopotential array. Returns the (Φ, Φ_sfc, plvd) triple. @@ -218,7 +221,8 @@ end Φ_fts[2][i, j, k] = 5000.0 * k * g # heights {5, 10, 15, 20} km end - tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock = Clock(time = 0.0)) + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), longitude=(0, 1), latitude=(0, 1), z=plvd) @@ -226,11 +230,17 @@ end @test znodes(grid, Center()) ≈ [3000.0, 6000.0, 9000.0, 12000.0] # Lz = max - min = 20*1000 - 1*1000 = 19000. @test grid.Lz ≈ 19000.0 + + # Regression: reading `interior` of the snapshot instead of forwarding through + # `.operand` would collapse the column mean onto the last materialized time. + clock.time = 1.0 + materialize_geopotential!(grid) + @test znodes(grid, Center()) ≈ [3000.0, 6000.0, 9000.0, 12000.0] end @testset "TimeSeriesInterpolation-backed Φ heights follow the clock" begin - # The whole point of the FTS-backed vertical: `rnode` must return - # different per-cell heights as the shared clock advances. + # `rnode` must return different per-cell heights as the shared clock advances — but + # only at each `materialize_geopotential!`, holding still in between. Nx, Ny, Nz = 2, 2, 3 Φ_grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), longitude=(0, 1), latitude=(0, 1), z=(0, 1)) @@ -251,8 +261,209 @@ end @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 # k=2 at t=0 → 2 km clock.time = 10.0 + @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 # stale: the snapshot has not been refreshed + materialize_geopotential!(grid) @test rnode(1, 1, 2, grid, ℓ...) ≈ 4000.0 # same grid, later snapshot → 4 km clock.time = 5.0 + materialize_geopotential!(grid) @test rnode(1, 1, 2, grid, ℓ...) ≈ 3000.0 # linear-in-time between snapshots end + + @testset "materialized Φ reproduces the interpolated-in-time Φ" begin + # Regression: the materialized snapshot must match, column by column, a static-`Field` + # discretization built from Φ blended to the same time by hand. Terrain cuts a different + # number of levels out of each column, so the clip / first-above-ground path is covered too. + Nx, Ny, Nz = 4, 3, 6 + times = [0.0, 100.0, 200.0] + stretch = [1.0, 1.1, 1.2] # level heights rise with time + raw_height(i, j, k, n) = 500.0 * (i + j) + 900.0 * k * stretch[n] + surface_height(i, j) = 500.0 * (i + j) + 900.0 * (i - 0.5) # clips k < i - 0.5 + + Φ_grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, times) + for n in eachindex(times), i in 1:Nx, j in 1:Ny, k in 1:Nz + Φ_fts[n][i, j, k] = raw_height(i, j, k, n) * g + end + fill_halo_regions!(Φ_fts) + + Φ_sfc_grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, 1), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_sfc = CenterField(Φ_sfc_grid) + interior(Φ_sfc) .= [surface_height(i, j) * g for i in 1:Nx, j in 1:Ny, k in 1:1] + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g, + surface_geopotential = Φ_sfc) + grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), longitude=(0, 1), latitude=(0, 1), + z=plvd, topology=(Bounded, Bounded, Bounded)) + + # Blended from the two bracketing snapshots *after* each is clipped, matching the order + # `clip_subsurface!` and materialization impose. + function reference_grid(t) + n₁ = searchsortedlast(times, t) + n₂ = min(n₁ + 1, length(times)) + w = n₂ == n₁ ? 0.0 : (t - times[n₁]) / (times[n₂] - times[n₁]) + Φ = CenterField(Φ_grid) + interior(Φ) .= [(1 - w) * max(raw_height(i, j, k, n₁), surface_height(i, j)) * g + + w * max(raw_height(i, j, k, n₂), surface_height(i, j)) * g + for i in 1:Nx, j in 1:Ny, k in 1:Nz] + reference = PressureLevelVerticalDiscretization(Φ; gravitational_acceleration = g, + surface_geopotential = Φ_sfc) + return LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), longitude=(0, 1), latitude=(0, 1), + z=reference, topology=(Bounded, Bounded, Bounded)) + end + + rnode = Oceananigans.Grids.rnode + ℓ = (Center(), Center(), Center()) + + for t in (0.0, 37.0, 100.0, 155.5, 200.0) + clock.time = t + materialize_geopotential!(grid) + reference = reference_grid(t) + + for i in 1:Nx, j in 1:Ny + for k in 1:Nz + @test rnode(i, j, k, grid, ℓ...) ≈ rnode(i, j, k, reference, ℓ...) + end + # Targets spanning below the terrain to above the column top. + for z in range(surface_height(i, j) - 2000, raw_height(i, j, Nz, 3) + 2000, length=17) + @test column_fractional_z_index(z, float(i), float(j), grid) ≈ + column_fractional_z_index(z, float(i), float(j), reference) + end + end + end + end + + @testset "materialize_geopotential! is a no-op on a static-Field Φ" begin + grid, Φ, _, plvd = make_plg() + + @test plvd.geopotential === Φ # no copy, no wrapper on the static path + materialize_geopotential!(grid) # warm up + @test (@allocated materialize_geopotential!(grid)) == 0 + end + + @testset "update_state! refreshes the materialized Φ" begin + # Regression: `update_state!` is the only hook that refreshes the snapshot, so dropping + # the call would leave a stepped atmosphere reading Φ at its initial time forever. + Nx, Ny, Nz = 2, 2, 3 + Φ_grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0, 10.0]) + for i in 1:Nx, j in 1:Ny, k in 1:Nz + Φ_fts[1][i, j, k] = 1000.0 * k * g + Φ_fts[2][i, j, k] = 2000.0 * k * g + end + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=plvd) + + atmosphere = PrescribedAtmosphere(grid, [0.0, 10.0]; clock) + rnode = Oceananigans.Grids.rnode + ℓ = (Center(), Center(), Center()) + + time_step!(atmosphere, 5.0) + @test atmosphere.clock.time == 5.0 + @test rnode(1, 1, 2, grid, ℓ...) ≈ 3000.0 + end + + @testset "materialized Φ follows the series' time extrapolation" begin + # Regression: on a node and past either end of the window, the snapshot must reproduce the + # series' own `Clamp()` extrapolation rather than clamp or skip the refresh itself. + Nx, Ny, Nz = 2, 2, 3 + Φ_grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0, 10.0]) + for i in 1:Nx, j in 1:Ny, k in 1:Nz + Φ_fts[1][i, j, k] = 1000.0 * k * g + Φ_fts[2][i, j, k] = 2000.0 * k * g + end + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=plvd) + + rnode = Oceananigans.Grids.rnode + ℓ = (Center(), Center(), Center()) + + for (t, height) in ((10.0, 4000.0), # exactly on the last node + (25.0, 4000.0), # past the end: clamped to the last snapshot + (-5.0, 2000.0), # before the start: clamped to the first + ( 0.0, 2000.0)) # exactly on the first node + clock.time = t + materialize_geopotential!(grid) + @test rnode(1, 1, 2, grid, ℓ...) ≈ height + end + end + + @testset "single-snapshot TimeSeriesInterpolation" begin + # Regression: with one time level there is no interval to blend across, and no + # `surface_geopotential` to clip against — both edges of the construction path. + Nx, Ny, Nz = 2, 2, 3 + Φ_grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0]) + for i in 1:Nx, j in 1:Ny, k in 1:Nz + Φ_fts[1][i, j, k] = 1000.0 * k * g + end + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), longitude=(0, 1), latitude=(0, 1), + z=plvd, topology=(Bounded, Bounded, Bounded)) + + rnode = Oceananigans.Grids.rnode + ℓ = (Center(), Center(), Center()) + + @test znodes(grid, Center()) ≈ [1000.0, 2000.0, 3000.0] + @test grid.Lz ≈ 2000.0 + @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 + + clock.time = 500.0 + materialize_geopotential!(grid) + @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 # one snapshot: nowhere else to go + + @test surface_elevation(grid) === nothing + @test column_fractional_z_index(0.0, 1.0, 1.0, grid) == 1 # no clip ⇒ no plateau to skip + end + + @testset "materialize_geopotential! refills the snapshot's halos" begin + # Regression: a refresh that wrote the interior without refilling halos would strand + # halo Φ on the previous snapshot. + Nx, Ny, Nz = 4, 4, 4 + Φ_grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0, 10.0]) + for i in 1:Nx, j in 1:Ny, k in 1:Nz + Φ_fts[1][i, j, k] = 1000.0 * k * g + Φ_fts[2][i, j, k] = 2000.0 * k * g + end + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(CPU(); size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=plvd) + + clock.time = 2.5 # quarter of the way between snapshots + materialize_geopotential!(grid) + + # Same values on the same grid with the same default boundary conditions, so the halos + # have to match as well as the interior. + Φ_reference = CenterField(Φ_grid) + interior(Φ_reference) .= [1250.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz] + fill_halo_regions!(Φ_reference) + @test parent(grid.z.geopotential) ≈ parent(Φ_reference) + end end