NestedSimulation abstraction + ERA5 regional hindcast example#233
NestedSimulation abstraction + ERA5 regional hindcast example#233ewquon wants to merge 32 commits into
Conversation
9d673f4 to
5463396
Compare
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
This reverts commit 5463396.
|
Okay, the model state should be ready for dynamic initialization. A couple of important notes:
|
|
@ewquon should we add terrain? Terrain following should work for the fully compressible formulation (just not substepping) |
Hallelujah! I'm happy to add that in. Anything that gets us closer to reality sooner rather than later is a win in my book. |
|
@ewquon I added the "build all examples" label so that it will build even if the "build always" tag is false. |
|
@giordano do you know what the GPU error is? |
|
I believe simply happens when the runner was reused from a previous run from somewhere else in the organisation, but the new build uses a new Docker image than the previous run, and then it fails to pull the second one because the combination of the two images fills the partition. In summary, it's an inoffensive, albeit annoying, issue, restarting the job should succeed (unless you rerun in the same situation 😄) |
|
NumericalEarth/Breeze.jl#715 recovers some space on the VMs also in the jobs we run in the Breeze.jl repository (at the moment we do the cleanup only in this repo), hopefully this will make the error less frequent. |
A parent–child wiring for limited-area runs. `PrescribedAtmosphere` now returns 3D `(Center, Center, Center)` field-time-series defaults whenever the grid has `size(grid, 3) > 1`, so it can play the role of a nesting parent without changes to the existing 2D surface-forcing call sites. `src/EarthSystemModels/NestedSimulations/` defines: - a duck-typed parent interface (`parent_clock`, `parent_field`, `parent_time_step!`, `parent_update_state!`, `parent_interpolate`) with methods for Oceananigans `Simulation`, `FieldTimeSeries`, and bare callables; the corresponding `PrescribedAtmosphere` methods live next to the type in `prescribed_atmosphere.jl`, - `NestedSimulation(parent, child::Simulation)` which installs a callback that syncs the parent clock to the child after each iteration, and forwards `time_step!` / `run!` to the child, - `parent_boundary_conditions(parent; variables, sides, grid, schemes)` which builds a NamedTuple of `FieldBoundaryConditions` populated with `OpenBoundaryCondition`s closed over per-side `ParentBoundaryFunction` callables (dispatched via `Val(side)` for the right (y,z,t)/(x,z,t)/ (x,y,t) arity). Test: `test/test_nested_simulation.jl` exercises both the 2D/3D defaults and a translating Lamb-Oseen vortex driving an Oceananigans `NonhydrostaticModel` child. Also bumps `Breeze` compat to "0.4, 0.5" so the example/regional-hindcast work can move to Breeze 0.5 once that lands. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Resolves conflicts in: - src/NumericalEarth.jl — adopt main's relative-`using .EarthSystemModels` style and add `using .EarthSystemModels.NestedSimulations: NestedSimulation, parent_boundary_conditions`. - src/Atmospheres/Atmospheres.jl — adopt main's `using ..EarthSystemModels: EarthSystemModels, AbstractPrescribedComponent` style; drop the previous `import NumericalEarth.EarthSystemModels.NestedSimulations: parent_*` imports along with the rest of the parent-interface scaffolding. - src/Atmospheres/prescribed_atmosphere.jl — keep main's `Oceananigans.`- qualified restore_prognostic_state! signature; drop the parent_clock / parent_field / parent_time_step! / parent_update_state! methods on PrescribedAtmosphere (no longer needed, see below). Simplifies NestedSimulations to lean on existing Oceananigans machinery: - `parent_boundary_conditions(grid; variables, sides, schemes)` now takes a NamedTuple of child_field => FieldTimeSeries directly. Per-side closures capture the boundary-edge coordinate and call `Oceananigans.Fields. interpolate(X, Time(t), fts, …)` from inside an `OpenBoundaryCondition`. Standard `ContinuousBoundaryFunction` regularization tags the side — no bespoke `ParentBoundaryFunction` struct or `Val(side)` dispatch. - `NestedSimulation`'s sync callback uses `parent.clock.time` and `Oceananigans.TimeSteppers.time_step!(parent, Δt)` directly. The whole `parent_interface.jl` (parent_clock, parent_field, parent_time_step!, parent_update_state!, parent_interpolate) is removed. Also renames the regional-hindcast example to "ERA5 downscaling with Breeze and NestedSimulation" (both the literate header and the docs/make.jl entry). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
working syntax (updated for the latest API on the branch — using NumericalEarth, Oceananigans, Breeze
using Breeze.AnelasticEquations: AnelasticDynamics
using Breeze.Thermodynamics: ReferenceState
# Lamb-Oseen vortex translating at U in x
const Γ, a, U, x₀, y₀ = 5e3, 200.0, 10.0, 600.0, 1000.0
const ρ̄ = 1.225
function lamb_oseen_uv(x, y, t)
dx, dy = x - x₀ - U*t, y - y₀
r² = dx^2 + dy^2
uθ_over_r = r² < eps() ? 0.0 : (Γ / (2π * r²)) * (1 - exp(-r² / a^2))
return (U - uθ_over_r * dy, uθ_over_r * dx)
end
# Parent: coarser and strictly larger than the child, so the FTS brackets every
# child sampling node (Interpolated BC + Relaxation-on-FTS both require this).
parent_grid = RectilinearGrid(size = (40, 16, 8),
x = (-500, 4500), y = (-500, 2500), z = (-20, 120),
topology = (Bounded, Periodic, Bounded))
parent = PrescribedAtmosphere(parent_grid, collect(0.0:5.0:305.0))
set!(parent.velocities.u, (x, y, z, t) -> ρ̄ * lamb_oseen_uv(x, y, t)[1]) # ρ̄·u → Breeze momentum BC
set!(parent.velocities.v, (x, y, z, t) -> ρ̄ * lamb_oseen_uv(x, y, t)[2])
# Child: finer LAM grid
child_grid = RectilinearGrid(size = (128, 32, 4),
x = (0, 4000), y = (0, 2000), z = (0, 100),
topology = (Bounded, Periodic, Bounded))
# Sponge mask: ramps from 0 in the interior to 1 in the outer 15% of x.
const Lx, x_sponge = 4000.0, 0.85 * 4000.0
sponge_mask(x, y, z) = x < x_sponge ? 0.0 : ((x - x_sponge) / (Lx - x_sponge))^2
# child_simulation builds the Breeze child model, wires Open BCs from the parent
# on the requested sides, and adds Relaxation-on-FTS interior nudging.
# It returns the constructed AtmosphereModel.
ref = ReferenceState(child_grid; surface_pressure = 101325.0, potential_temperature = 300.0)
child = child_simulation(AtmosphereModel, child_grid, parent;
sides = (:west, :east),
relaxation_rate = 1 / 60, # 60-second sponge timescale
relaxation_mask = sponge_mask,
dynamics = AnelasticDynamics(ref),
formulation = :LiquidIcePotentialTemperature,
advection = WENO(order = 5))
set!(child; θ = 300.0,
ρu = (x, y, z) -> ρ̄ * lamb_oseen_uv(x, y, 0)[1],
ρv = (x, y, z) -> ρ̄ * lamb_oseen_uv(x, y, 0)[2],
ρw = 0.0)
# NestedSimulation(parent, child; sim_kwargs...) is a convenience for
# `Simulation(NestedModel(parent, child); sim_kwargs...)`. The NestedModel's
# `time_step!` steps the child, then ticks the parent's clock to match.
nested = NestedSimulation(parent, child; Δt = 0.2, stop_time = 300.0)
run!(nested)What's in play:
Needs Breeze on |
`child_simulation(modeltype, grid, parent; …)` builds a `NestedSimulation` in a single call: it wires `parent_boundary_conditions` for the open sides, optionally `parent_forcings` for interior relaxation, constructs the child model and `Simulation`, and installs the parent-sync callback. The variable mapping (which parent FieldTimeSeries drives which child field) is chosen by dispatch on `default_parent_variables(modeltype, parent)`: - Oceananigans `NonhydrostaticModel` → `(u, v) ← parent.velocities.(u, v)`. - Breeze `AtmosphereModel` → `(ρu, ρv) ← parent.velocities.(u, v)` (assumes the parent stores ρ̄·u in the velocity slots). Lives in `NumericalEarthBreezeExt` so Breeze stays a weakdep. `parent_forcings(; variables, rate, mask)` returns a NamedTuple of `Oceananigans.Forcings.Relaxation`, one per variable, whose target is a closure reusing the same `Oceananigans.Fields.interpolate(X, Time(t), fts, …)` path as the BC closures. `rate` and `mask` accept scalars, callables, fields, or per-variable NamedTuples; the existing Oceananigans mask types (`GaussianMask`, etc.) work straight through. `child_simulation` exposes this via `relaxation_rate` / `relaxation_mask` kwargs and merges with any user-supplied `forcing` already in `model_kwargs`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Oceananigans main natively supports `Relaxation(rate, mask, target=fts)`: `materialize_forcing` wraps the FTS in a `FieldTimeSeriesTarget` that handles space/time interpolation and GPU adaptation. Dropping the closure wrapper we used before — that path silently extrapolates near boundaries, whereas the FTS-direct path runs `validate_fts_target_extent` to ensure the parent FTS strictly brackets every child sampling position. `parent_boundary_conditions` still uses its per-side closures since the BC machinery doesn't yet have an analogous native FTS path. Compat note: needs the Oceananigans release that bundles FTS-Relaxation support (currently main, post-0.108.0). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Lets `OpenBoundaryCondition(fts)` work directly, instead of requiring user
code to wrap the FTS in per-side closures. The pattern mirrors what
Oceananigans `Relaxation` does for FTS targets:
- `InterpolatedFTSBoundary{Dim, SideType, LX, LY, LZ, FTS}` carries the
boundary-normal axis, side, and the child field's location at the
boundary face as type parameters.
- `Oceananigans.BoundaryConditions.regularize_boundary_condition(::FlavorOfFTS, …)`
builds the side-tagged condition during normal Oceananigans regularization,
and runs the same strict bracketing check Oceananigans uses for
`Relaxation`-on-FTS (FTS node extents must contain every child sampling
position).
- `getbc` methods for dims 1/2/3 compute `node(i, j, k, grid, LX(), LY(), LZ())`
at the boundary face and call `Oceananigans.Fields.interpolate(X, Time(t), fts, …)`.
`parent_boundary_conditions` collapses to one line per side — no closures,
no `Val(side)` dispatch, no edge-coord precomputation in user code. The
implementation lives in NestedSimulations (mild type piracy on
`regularize_boundary_condition(::FieldTimeSeries, …)`) until an analogous
path lands upstream in Oceananigans, after which this file becomes a
deprecation shim.
Test updated to use a parent grid that strictly brackets the child (required
by the same validation used for Relaxation). 13/13 still pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ries
The previous version added a `regularize_boundary_condition(::FieldTimeSeries, …)`
method — type piracy on an Oceananigans-owned function with an
Oceananigans-owned argument type. Replace with our own wrapper type:
OpenBoundaryCondition(Interpolated(fts))
`Interpolated{Dim, Side, LX, LY, LZ, FTS}` carries the side / dimension /
field-location tags as type parameters. The user-facing constructor leaves
them all `Nothing`; Oceananigans' regularization runs through our
`regularize_boundary_condition(::Interpolated{Nothing}, …)` method and
fills in the type parameters. Same bracketing validation, same `getbc`
machinery, no piracy.
`Interpolated` is intentionally not exported — it's an internal wrapper
used by `parent_boundary_conditions`. The public surface stays
`NestedSimulation`, `parent_boundary_conditions`, `parent_forcings`,
`child_simulation`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…s a wrapper Parent-sync moves out of a callback on the child `Simulation` and into the integrator itself. The new layering: - `NestedModel(parent, child::AbstractModel) <: AbstractModel` pairs the two and overrides `time_step!` to step the child, then advance the parent to match the new child clock. Property access (`clock`, `grid`, `velocities`, …) forwards to the child; `fields`, `prognostic_fields`, `architecture`, `iteration`, `update_state!`, and the checkpointing pair forward via explicit methods. The model now satisfies the `AbstractModel` protocol — plain `Simulation(NestedModel(…))` just works. - `NestedSimulation(parent, child_model::AbstractModel; sim_kwargs...)` is a one-line convenience: `Simulation(NestedModel(parent, child_model); …)`. No callback installation, no hidden integrator state. - `child_simulation` → `child_model`: renamed (and the file moved), returns the constructed child model instead of a `Simulation`. The caller follows up with `NestedSimulation(parent, child; Δt, stop_time, …)`. Matching rename on the Breeze ext file (`breeze_child_simulation.jl` → `breeze_child_model.jl`). `NumericalEarth.jl` exports both `NestedModel` and `NestedSimulation`. Test collapses to the new two-call form; 13/13 still pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ation
Keeps the name `child_simulation` (matching `ocean_simulation` /
`atmosphere_simulation` convention) while still returning an `AbstractModel`
— the Breeze `atmosphere_simulation` helper returns a model, not a
`Simulation`, and we follow that convention so the call site composes:
child = child_simulation(modeltype, grid, parent; …)
nested = NestedSimulation(parent, child; Δt, stop_time, …)
The underlying model constructor is now an extension point:
`_build_child_model(modeltype, grid; kwargs...)` defaults to
`modeltype(grid; kwargs...)`; the Breeze ext overrides it for
`AtmosphereModel` to dispatch through `atmosphere_simulation`. That way the
Breeze call picks up Breeze's default advection / microphysics when the user
doesn't override them, rather than reimplementing those defaults here.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…tion alias
Three changes:
1. Rename `default_parent_variables` → `parent_variables`. The function is
now multi-dispatch on the (child_modeltype, parent_type) pair — methods
live wherever both types are visible. A generic fallback throws an
`ArgumentError` that names the missing pair and suggests defining a
method or passing `variables=` explicitly. Cross-realm pairs (e.g. ocean
child × atmosphere parent) are intentionally undefined: the fallback
error names them out so silent unit mismatches can't happen.
The Breeze ext now declares
`parent_variables(::Type{<:AtmosphereModel}, ::PrescribedAtmosphere)`
(was just `default_parent_variables(::Type{<:AtmosphereModel}, parent)`),
documenting the density convention in the method's source comment.
2. Generalize `Interpolated` from `FlavorOfFTS` only to any
`Union{FlavorOfFTS, AbstractField}` source — `_query_source` dispatches
on the source type (FTS → space + time interpolation, Field → space
only at the parent's current state). Lets prognostic-parent setups
(Breeze / SpeedyWeather as the parent) reuse the same wiring path.
3. Add `const NestedSimulation = Simulation{<:NestedModel}` type alias
alongside the existing convenience constructor — enables `::NestedSimulation`
dispatch and still gives users `NestedSimulation(parent, child; …)`
construction.
The `NonhydrostaticModel × PrescribedAtmosphere` default was dropped since
ocean models cannot have atmosphere parents in a unit-consistent way. The
existing test still works because it passes `variables=` explicitly.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
OpenBoundaryCondition on cell-centered fields silently overwrites the first interior cell on W/S walls (asymmetric vs E/N) — the vortex validation campaign hit this on \xcf\x81/\xcf\x81\xce\xb8. The new bc_types kwarg lets the caller pick the BC constructor per child variable, e.g. bc_types = (\xcf\x81 = ValueBoundaryCondition, \xcf\x81\xce\xb8 = ValueBoundaryCondition). Default stays OpenBoundaryCondition for back-compat. Errors when a schemes entry is provided for a non-OpenBC field.
…loop Replace the single-snapshot ERA5 → LAM interpolation with a loop that populates a PrescribedAtmosphere parent + user-owned qᶜ/qⁱ FTSs on the ERA5-native bbox grid, then builds the LAM IC by horizontal regrid of snapshot 1. Same IC behavior as before; this readies the parent for the NestedSimulation BC/Davies wiring that follows. Φ₀ hoisted out of the loop (surface elevation is time-invariant). The small bbox override pins era5_region to the cached 41.2 bbox; the natural era5_bbox() computation gives 41.25, which would force re-downloads. Also drops the dangling `using Breeze.Thermodynamics` import — `vapor_gas_constant` is now re-exported at the top level of Breeze (PR #699).
…pling fluxes Two adjacent framework fixes the era5_breeze example needs to wire OBCs and Davies relaxation on a child driven by a PrescribedAtmosphere parent. 1. `Interpolated` getbc: default clock to nothing → t=0. `fill_halo_regions!` is called without a clock during set!-time IC setup; the kernel launches with an empty `args` tuple, so the existing getbc methods (which required clock as a positional arg) fell through to the generic 'not callable' fallback. Methods now accept clock=nothing and default time to 0 for IC evaluation. 2. `atmosphere_simulation`: per-side merge for user-supplied BCs. The previous shallow `merge(coupling_bcs, user_bcs)` replaced whole `FieldBoundaryConditions` objects per prognostic, clobbering the coupling's bottom-flux BCs (ρτˣ, ρτʸ, Jᵉ, Jᵛ) whenever the user added lateral BCs on the same field. The new `_merge_boundary_conditions` merges field-by-field, then side-by-side within each FieldBoundaryConditions — the user's non-default sides override the coupling's defaults; the coupling's bottom-flux survives where the user leaves the bottom default.
Drive the LAM's lateral boundaries from the parent FTSs:
- ρu, ρv get OpenBC(Interpolated(fts)) at Face stagger
- ρ, ρe (≡ ρθ post-conversion), ρqᵉ get ValueBC at Center stagger,
since OpenBC silently overwrites the first interior cell on the
W/S walls of Center-located fields (validated against the vortex-
transit campaign)
- Davies inner-fringe forcings under specific keys (u, v, θ, qᵉ) so
Breeze's SpecificForcing applies the ρ multiply at kernel time
Cosine ramp over 5 lateral cells; τ_relax = 5·Δx/U_scale at the domain
center latitude. populate_parent_snapshot! now also derives ρ, ρu, ρv,
ρθˡⁱ, ρqᵗ, θˡⁱ, qᵗ on the parent grid and stores them in their FTSs.
Long. labels shifted from the ERA5 [0°, 360°] convention to [-180°, 180°]
to match the LAM grid's coordinate range (data arrays unchanged).
- Drop unused `pl_vars` list (4 lines). - Capture snapshot-1 raw ERA5 arrays from `populate_parent_snapshot!`'s return value instead of re-fetching via `Metadatum` in the plot block (single source of truth, saves ~10 lines). - Rename `Φ₀_arr_snap1` → `Φ₀_arr` now that the plot block doesn't shadow it. - Add a comment tying the existence of `column_interp_z!` / `interp_z_masked` to the still-pending terrain checklist item, with a pointer to PR #241's native `set!(target, metadatum)` path that replaces them once terrain support lands.
Threads an `ARCH=GPU` env-var switch through grid + FieldTimeSeries construction so the example can integrate on CUDA: - `arch = GPU(CUDA.CUDABackend(always_inline=true))` when ARCH=GPU, else `CPU()`. Passed as the first arg to both the LAM and parent `LatitudeLongitudeGrid` constructors. - `column_interp_z!` now does its per-column compute on a host buffer and `copyto!`s into `interior(inter_field)`. The previous loop indexed `interior` directly, which is a no-op on CPU but errors on GPU (scalar device indexing). - Per-snapshot FTS writes use `copyto!(interior(fts, …), host_arr)` in place of the `.=` broadcast — `CuArray .= Array` isn't supported. - `lateral_mask` becomes a `let`-closure that captures the domain extents and fringe width as closure-local bindings. The previous free-function form referenced non-const module globals, which produces dynamic-dispatch IR that CUDA's GPU compiler can't lower. Also fixes `NestedSimulations.Interpolated` to carry the source's grid as a struct field rather than reaching through `source.grid`. `GPUAdaptedFieldTimeSeries` drops the `.grid` field during Adapt, so the old `fts.grid` access in `_query_source` failed GPU compilation with `unsupported call to jl_f_getfield`. Mirrors Oceananigans's own `FieldTimeSeriesTarget` pattern (PR #5575).
… the plot - Replace the placeholder `Δt = 1.0` with an acoustic-CFL-derived value computed from the surface c_sound and `minimum_zspacing(grid)` ($\approx$ 0.044 s for our 50 m bottom cell). - `stop_iteration = 100` instead of `stop_time = parent_times[end]`. A short smoke run \xe2\x80\x94 just enough to verify the BC + Davies machinery doesn't NaN. Substepping (next step) lifts the acoustic-CFL cap. - Add a per-10-iter progress callback logging max|u|, max|v|, max|w|, and \xcf\x81 range. - Move `run!(nested)` from the end of the script to before the plot block and extract post-run LAM state from the model (\xcf\x81, \xcf\x81u, \xcf\x81v, \xcf\x81\xce\xb8, \xcf\x81q\xe1\xb5\x89), with \xce\xb8 = \xcf\x81\xce\xb8/\xcf\x81 and q\xe1\xb5\x97 = \xcf\x81q\xe1\xb5\x89/\xcf\x81. - Overlay a t=t_stop red scatter on each profile axis alongside the t=0 blue scatter; add an axislegend on the rightmost top axis with the actual `model.clock.time` value; drop the redundant 'LAM: blue' from the figure title. - Surface-T and surface-q\xe1\xb5\x9b placeholders (`ValueBoundaryCondition` at the bottom of `:\xcf\x81e` and `:\xcf\x81q\xe1\xb5\x89`) replace `atmosphere_simulation`'s coupling J\xe1\xb5\x89/J\xe1\xb5\x9b bottom-flux BCs. Without this, the bottom :\xcf\x81e flux gets wrapped in an `EnergyFluxBoundaryCondition` whose runtime kernel calls `\xf0\x9d\x92\xac_to_J\xe1\xb6\xbf(... fields.q\xe1\xb5\x9b[i,j,k])` and crashes \xe2\x80\x94 the bulk-flux state isn't populated until SlabLand lands. The placeholders also document the shape of the values land coupling will eventually supply.

Summary
@glwagner @kaiyuan-cheng
Adds
examples/era5_breeze.jl— a building block for a regional modeling example that will eventually couple Breeze (compressible solver, in development) to forthcoming SlabLand and SlabOcean components.We download ERA5 reanalysis over an SGP-centered LAM bounding box (HI-SCALE 2016-09-10 case day) and interpolate onto a
LatitudeLongitudeGridsized for ~3 km horizontal cells at the domain center latitude.Tᵥis computed as a derived field usingBreeze.ThermodynamicConstantsfor theRᵥ/Rd − 1coefficient.Punch list
Notes
Nz=1rather than a 2-D(Bounded, Bounded, Flat)grid, sidesteppingCliMA/Oceananigans.jl#5473(Flat↔non-Flatinterpolate!errors with a crypticBoundsError; #5474 will convert that to a clearArgumentErrorbut does not lift the restriction). Mirrors the pattern inexamples/ERA5_hourly_data.jl.LatitudeLongitudeGrid. A future variant on a projected Cartesian grid would benefit from theset!projection support proposed in Add map projection support toset!forRectilinearGridtargets #232.Test plan
~/.cdsapircconfigured.