diff --git a/src/Atmospheres/prescribed_atmosphere.jl b/src/Atmospheres/prescribed_atmosphere.jl index fe1fd40dd..fa98bd59e 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 ca3199000..162cc52ba 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 00f1975bb..e078946c8 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. @@ -224,7 +227,8 @@ for arch in test_architectures set!(Φ_fts[1], [1000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) # heights {1, 2, 3, 4} km set!(Φ_fts[2], [5000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) # heights {5, 10, 15, 20} km - 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(arch; size=(Nx, Ny, Nz), longitude=(0, 1), latitude=(0, 1), z=plvd) @@ -232,11 +236,17 @@ for arch in test_architectures @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(arch; size=(Nx, Ny, Nz), longitude=(0, 1), latitude=(0, 1), z=(0, 1)) @@ -253,13 +263,208 @@ for arch in test_architectures rnode = Oceananigans.Grids.rnode ℓ = (Center(), Center(), Center()) - @allowscalar begin - @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, ℓ...) ≈ 4000.0 # same grid, later snapshot → 4 km - clock.time = 5.0 - @test rnode(1, 1, 2, grid, ℓ...) ≈ 3000.0 # linear-in-time between snapshots + # `rnode` indexes the geopotential field directly, so a host-side read on a device + # grid is a scalar read by construction. + @allowscalar @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 # k=2 at t=0 → 2 km + + clock.time = 10.0 + @allowscalar @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 # stale: the snapshot has not been refreshed + materialize_geopotential!(grid) + @allowscalar @test rnode(1, 1, 2, grid, ℓ...) ≈ 4000.0 # same grid, later snapshot → 4 km + + clock.time = 5.0 + materialize_geopotential!(grid) + @allowscalar @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(arch; 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) + set!(Φ_fts[n], [raw_height(i, j, k, n) * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + end + fill_halo_regions!(Φ_fts) + + Φ_sfc_grid = LatitudeLongitudeGrid(arch; size=(Nx, Ny, 1), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_sfc = CenterField(Φ_sfc_grid) + set!(Φ_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(arch; 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) + set!(Φ, [(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(arch; 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) + + @allowscalar 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(arch) + + @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(arch; size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0, 10.0]) + set!(Φ_fts[1], [1000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + set!(Φ_fts[2], [2000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(arch; 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 + @allowscalar @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(arch; size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0, 10.0]) + set!(Φ_fts[1], [1000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + set!(Φ_fts[2], [2000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(arch; 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) + @allowscalar @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(arch; size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0]) + set!(Φ_fts[1], [1000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(arch; 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 + @allowscalar @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 + + clock.time = 500.0 + materialize_geopotential!(grid) + @allowscalar @test rnode(1, 1, 2, grid, ℓ...) ≈ 2000.0 # one snapshot: nowhere else to go + + @test surface_elevation(grid) === nothing + @allowscalar @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(arch; size=(Nx, Ny, Nz), + longitude=(0, 1), latitude=(0, 1), z=(0, 1)) + Φ_fts = FieldTimeSeries{Center, Center, Center}(Φ_grid, [0.0, 10.0]) + set!(Φ_fts[1], [1000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + set!(Φ_fts[2], [2000.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + fill_halo_regions!(Φ_fts) + + clock = Clock(time = 0.0) + tsi = TimeSeriesInterpolation(Φ_fts, Φ_fts.grid; clock) + plvd = PressureLevelVerticalDiscretization(tsi; gravitational_acceleration = g) + grid = LatitudeLongitudeGrid(arch; 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) + set!(Φ_reference, [1250.0 * k * g for i in 1:Nx, j in 1:Ny, k in 1:Nz]) + fill_halo_regions!(Φ_reference) + @test Array(parent(grid.z.geopotential)) ≈ Array(parent(Φ_reference)) + end end end