From 87a3c2c0e2950f0fc719b18886be9ae05dd39100 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 14:29:37 +0200 Subject: [PATCH 01/19] Add new plan formatting rule to AGENTS.md --- AGENTS.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index f35d183c8..78b9b91a9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -277,7 +277,11 @@ N.B: Make sure that each revision is given a number and a date. ``` -The revision log should, to the greatest extent possible, briefly summarize changes to the plan that are made on-the-fly during development. Make sure that each revision is given a number and a date. **A human must approve each revision before implementation**. +The revision log should, to the greatest extent possible, briefly summarize changes to the plan that are made on-the-fly during development. Make sure that each revision is given a number and a date. +**A human must approve each revision before implementation**. + +Write plan prose with a line break at the end of every sentence (one sentence per line) rather than wrapping at an arbitrary column width. +This keeps diffs readable when individual sentences are revised. The remainder of the plan document may be adapted on a case-by-case basis but should generally follow this structure: From 39dce2af5e8775a3295b7b77db06f456ad7bbc00 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Tue, 6 Oct 2026 19:46:38 +0200 Subject: [PATCH 02/19] Update implementation plan template --- AGENTS.md | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 78b9b91a9..cdebb9f57 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -259,6 +259,10 @@ Each document should be prefaced by the following template: Date of initial draft: YYYY-MM-dd +Author: , + +Reviewer: () + Base revision: ## Originating prompt @@ -269,16 +273,23 @@ Base revision: > User prompts here -N.B: Make sure that each revision is given a number and a date. - ## Problem description ## Background ``` -The revision log should, to the greatest extent possible, briefly summarize changes to the plan that are made on-the-fly during development. Make sure that each revision is given a number and a date. -**A human must approve each revision before implementation**. +Notes on filling in the template (these notes are guidance about the template; do not copy them into +the plan document itself): + +- Leave the `Reviewer` field blank until a human has approved the plan, then fill in both name and + email after approval. Where the reviewer is the repository owner, `git config user.name` and + `git config user.email` are an acceptable source for those values. +- The revision log should, to the greatest extent possible, briefly summarize changes to the plan + that are made on-the-fly during development. Give each revision a number and a date, and record + the model name and version (if known) for that revision, since a plan may be revised by a + different model than the one that drafted it. **A human must approve each revision before + implementation**. Write plan prose with a line break at the end of every sentence (one sentence per line) rather than wrapping at an arbitrary column width. This keeps diffs readable when individual sentences are revised. From 2afecef7b33e9d84a6972f81fc10fe9cbaf2a9fe Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 16:30:59 +0200 Subject: [PATCH 03/19] Add initial plan for initializers --- ...LAN_spatially_varying_soil_initializers.md | 221 ++++++++++++++++++ 1 file changed, 221 insertions(+) create mode 100644 docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md diff --git a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md new file mode 100644 index 000000000..563ae86e4 --- /dev/null +++ b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md @@ -0,0 +1,221 @@ +# Spatially varying soil initializers via input variables + +> Status: **planned**. Draft awaiting human review; nothing implemented yet. + +Date of initial draft: 2026-10-08 + +Base revision: `ce9ae606ae39063f52bb0ef23a3328a8b1904055` (tip of `main`, branch `bg/improved-init`) + +## Originating prompt + +> We need to generalize the existing soil initializer to allow for spatially varying inputs. The idea +> would be to declare corresponding input variables for T0 and Qgeo and then pass the struct values as +> default value initializers. These initializers could also be functions or callable structs that contain +> additional parameters, e.g. the periodic initialization used for soil temp and currently duplicated +> across many examples. Draft an implementation plan and summarize your findings. Ask questions if needed. + +## Revision log + +- **Rev 1 (2026-10-08):** Initial draft. +- **Rev 2 (2026-10-08):** Open questions resolved by the author: variable names `initial_surface_temperature` and `geothermal_heat_flux` confirmed; hydrology initializers (change 5) are in scope; `LatitudinalClimatology(T_equator, ΔT)` is the only concrete field initializer and no generic function wrapper is added, since a plain function can be assigned directly to the initializer fields; `AbstractFieldInitializer` lives in core `src/initializers.jl`. + Implementation not yet approved. + +## Problem description + +The soil energy initializers in `src/models/soil/soil_model_init.jl` are column-uniform: + +- `ConstantSoilTemperature{NF}` stores a scalar `T₀`. +- `QuasiThermalSteadyState{NF}` stores scalars `T₀`, `Qgeo`, and `k_eff` and sets the temperature profile with a host-side closure, `set!(state.temperature, (x, z) -> T₀ - Qgeo / k_eff * z)`. + +Neither can express a surface temperature or geothermal heat flux that varies between columns. +As a result, every global example re-implements the same thing by hand: a latitude-dependent climatology `mean_annual_temperature(lat) = 20 - abs(40 * sin(lat))` and a closure `(x, z) -> T₀[round(Int, x)] - 0.05 * z` that indexes a masked latitude (or regridded air temperature) vector by the column index `x`. +This pattern appears, with small variations, in + +- `examples/simulations/soil_heat_global.jl` (which explicitly notes that `QuasiThermalSteadyState` "does not (yet) support spatially variable parameters"), +- `examples/simulations/soil_heat_global_soilgrids.jl`, +- `examples/simulations/speedy_wet_land.jl`, +- `examples/simulations/land_global_era5.jl` (surface temperature from the ERA5 `t2m` mean), +- `examples/autodiff/global_sensitivity.jl` (surface temperature from the first ERA5 `t2m` slice), and +- `examples/autodiff/differentiating_terrarium_reactant.jl` (a plain closure). + +These closures capture host arrays, run through Oceananigans' `set_to_function!` (which detours via the CPU on GPU and Reactant grids), and hardcode the geothermal gradient as a number instead of deriving it from the geothermal heat flux and thermal conductivity. +The geothermal heat flux used for initialization is also disconnected from the `geothermal_heat_flux` variable that the `GeothermalHeatFlux()` boundary condition reads, so the initial profile and the bottom boundary condition can silently disagree. + +## Background + +Three pieces of existing infrastructure shape the design. + +**Input variables with defaults.** `input(name, loc; default, units, ...)` declares an `InputVariable` whose `default` is applied once, at `Field` construction time, in `initialize(var, grid, clock, fields, bcs)` (`src/state_variables.jl`) via `set!(field, var.default)`. +The `Def` type parameter is currently restricted to `Union{Nothing, Number, Function}`. +`reset!(state)` zeroes prognostic, auxiliary, and tendency fields but leaves inputs alone, so defaults survive `initialize!(integrator)`. +Existing processes already use this (e.g. `sand_fraction` in `soil_horizon.jl`, `ground_temperature` in `autotrophic_respiration.jl`). + +**Input sources override declared inputs by name.** `StateVariables(model; input_variables)` builds `Variables(tuplejoin(variables(model), input_variables))`. +For duplicate names, `register!` keeps the first declaration (the model's) provided the metadata (`name`, `dims`, `units`, `eltype`) match and the domains are compatible; otherwise it throws. +At `initialize!(integrator)`, `InputSource`s are applied *before* the user `initializers` NamedTuple and the model initializer, so an `InputSource` named like a declared input simply overwrites its default field. +This is exactly the mechanism needed to let users supply a spatially varying `T₀` or `Qgeo` from data. + +**Initializers are invisible to `variables`.** `variables(::AbstractModel)` collects variables from `processes(model)`, a generated function that only picks fields whose type is `<: AbstractProcess`. +The `initializer` field is an `AbstractInitializer` and is skipped, so an initializer currently has no way to declare state variables. +`variables(::Any) = ()` already provides a safe fallback for `DefaultInitializer`. + +Two further facts matter for the design: + +- Initializer structs are not `@parameterized`, and parameters are opt-in via `@param`, so an initializer may hold non-scalar fields (functions, callable structs, `Field`s) without affecting `parameters(model)` or `ParameterEditing.reconstruct`. +- On a `ColumnRingGrid`, the horizontal coordinate `x` passed to `set!(field, f)` is the (1-based, float-valued) index of the active land column; latitude and longitude must be looked up through `grid.rings` and `grid.mask`. + Any reusable "latitude-dependent" initializer therefore needs access to the grid, which the `initialize(var, grid, ...)` hook already has. + +## Summary of changes + +### 1. Let initializers declare variables + +- In `src/abstract_model.jl`, split the shared default so that models also collect their initializer's variables: + + ```julia + variables(obj::AbstractCoupledProcesses) = tuplejoin(fastmap(variables, processes(obj))...) + variables(model::AbstractModel) = tuplejoin(fastmap(variables, processes(model))..., variables(get_initializer(model))) + ``` + + `variables(::Any) = ()` keeps `DefaultInitializer` and all third-party initializers working unchanged. +- In `src/models/soil/soil_model_init.jl`, add `variables(init::SoilInitializer) = tuplejoin(variables(init.energy), variables(init.hydrology), variables(init.biogeochem))`. +- Update the `AbstractInitializer` docstring in `src/initializers.jl` to document the optional `variables(::AbstractInitializer)` extension point. + +`LandModel` has its own `StateVariables` method but it also goes through `variables(model)`, so it picks the change up automatically. + +### 2. Generalize `InputVariable` defaults + +- Widen the `Def` type parameter of `InputVariable` (`src/abstract_variables.jl`) from `Union{Nothing, Number, Function}` to an unconstrained parameter. + It remains concrete per instance. + Document the accepted forms: `nothing`, a number, any argument accepted by `Oceananigans.set!` (functions of the node coordinates, arrays, `Field`s), or an `AbstractFieldInitializer` (below). +- Introduce a small hook in `initialize(var, grid, clock, fields, bcs)` so that the default is applied through a single overloadable method instead of a bare `set!`: + + ```julia + set_default!(field, grid, default) = set!(field, default) + set_default!(field, grid, ::Nothing) = nothing + ``` + +### 3. `AbstractFieldInitializer`: reusable, grid-aware, parameterized field initializers + +Add to `src/initializers.jl`: + +```julia +""" +Base type for reusable initializers of a single `Field`. Implementations are callable structs that +provide `initialize!(field, grid, init::AbstractFieldInitializer)`; the default evaluates `init` at +the node coordinates via `set!`. They can be passed as `default` of an `input` variable, as an entry of +the `initializers` NamedTuple of `initialize`, or as a parameter of a model initializer. +""" +abstract type AbstractFieldInitializer end + +initialize!(field::AbstractField, grid::AbstractGrid, init::AbstractFieldInitializer) = set!(field, (coords...) -> init(coords...)) +set_default!(field, grid, init::AbstractFieldInitializer) = initialize!(field, grid, init) +``` + +and make the user-facing `initialize!(state, inits::NamedTuple)` path accept them as well (it currently calls `set!` directly, which would not know what to do with a callable struct). +Only one concrete implementation is in scope here, to replace the duplicated example code: + +```julia +""" +Latitude-dependent surface temperature climatology on a `ColumnRingGrid`: +T(φ) = T_equator - |ΔT sin φ|. Evaluated per active land column. +""" +@kwdef struct LatitudinalClimatology{NF} <: AbstractFieldInitializer + T_equator::NF = 20 + ΔT::NF = 40 +end +(init::LatitudinalClimatology)(lat) = init.T_equator - abs(init.ΔT * sin(lat)) +function initialize!(field, grid::ColumnRingGrid, init::LatitudinalClimatology) + # host-side: gather latitudes of active columns, evaluate, then set!(field, values) +end +``` + +No generic wrapper around a user function is provided: a plain function of the node coordinates can be assigned directly to `T₀`/`Qgeo` (or to any `input` default), and the implementation must keep that path working. +The `periodic_bc` *boundary* condition in the same examples (diurnal cycle shifted by longitude) is a boundary-condition concern and is left out of this plan; see "Future work". + +### 4. Generalize the soil energy initializers + +Rewrite `ConstantSoilTemperature` and `QuasiThermalSteadyState` in `src/models/soil/soil_model_init.jl`: + +```julia +struct QuasiThermalSteadyState{NF, T0, QG} <: AbstractInitializer{NF} + "Initial surface temperature (°C): number, function of column coordinate, `Field`, or `AbstractFieldInitializer`" + T₀::T0 + "Geothermal heat flux (W/m²), same accepted forms" + Qgeo::QG + "Bulk thermal conductivity (W/m/K)" + k_eff::NF +end +QuasiThermalSteadyState(::Type{NF}; T₀ = zero(NF), Qgeo = NF(0.02), k_eff = one(NF)) where {NF} + +variables(init::QuasiThermalSteadyState) = ( + input(:initial_surface_temperature, Ground(Top()); default = init.T₀, units = u"°C", desc = "..."), + input(:geothermal_heat_flux, Ground(Bottom()); default = init.Qgeo, units = u"W/m^2", desc = "..."), +) +``` + +`initialize!(state, model, init::QuasiThermalSteadyState)` then reads the two input fields and fills `temperature` with a kernel launched via `launch!`, `T[i, j, k] = T₀[i, j] - Qgeo[i, j] / k_eff * z_k`, instead of a host-side closure. +This is GPU- and Reactant-native (no `set_to_function!` detour) and keeps `k_eff` as a scalar parameter. +`ConstantSoilTemperature` gets the same treatment with only the `initial_surface_temperature` input. + +Reusing the name `geothermal_heat_flux` means that a model with `GeothermalHeatFlux()` as bottom boundary condition and `QuasiThermalSteadyState` as initializer shares one field for both, and that an `InputSource(grid, Qgeo_field; name = :geothermal_heat_flux, domain = Ground(Bottom()), units = u"W/m^2")` drives both consistently. +The units in such an `InputSource` must match the declaration, otherwise `Variables` reports a conflict; this is documented. + +Numeric literals (`0.0`, `0.02`, `1.0`) in the current `@kwdef` structs are replaced by `NF`-typed defaults as required by the project rules. + +### 5. Hydrology initializers (same pattern, smaller scope) + +`SaturationWaterTable` (`water_table_depth`, `vadose_zone_saturation`) and `ConstantSaturation` (`sat`) are generalized the same way (`input(:water_table_depth, Ground(XY()); default = ...)` etc.) so the whole `SoilInitializer` is consistent (confirmed in scope, Rev 2). + +### 6. Examples + +Replace the hand-rolled closures with the new initializers: + +- `soil_heat_global.jl`, `soil_heat_global_soilgrids.jl`, `speedy_wet_land.jl`: `SoilInitializer(NF; energy = QuasiThermalSteadyState(NF; T₀ = LatitudinalClimatology(NF), Qgeo = ..., k_eff = ...))`, removing the duplicated `mean_annual_temperature`/`initial_soil_temperature` definitions and the "does not (yet) support spatially variable parameters" remark. +- `land_global_era5.jl`, `global_sensitivity.jl`: pass the regridded air temperature `Field` as `T₀`, or as an `InputSource` named `initial_surface_temperature`. +- `differentiating_terrarium_reactant.jl`: the plain closure stays valid through the `initializers` NamedTuple; optionally switch to the initializer to exercise it under Reactant. + +### 7. Exports + +Export `AbstractFieldInitializer` and `LatitudinalClimatology` from `src/models/models.jl` or `src/Terrarium.jl` next to the existing initializer exports. + +## Testing and verification + +- `test/soil/soil_energy_tests.jl`: extend the existing `SoilInitializer` tests to + - verify that `QuasiThermalSteadyState` with scalar `T₀`/`Qgeo` reproduces the current profile bit for bit (regression), + - verify a column-varying `T₀` given as a function and as a `Field`, + - verify that an `InputSource` named `geothermal_heat_flux` overrides the default and that the profile uses the overridden values, and + - verify that `initialize!(integrator)` (re-initialization) leaves the result unchanged. +- `test/state_variables.jl`: `variables(model)` includes the initializer's inputs; `DefaultInitializer` contributes none; mismatched units in an `InputSource` raise the existing conflict error. +- `test/grids.jl` or a new `test/initializers.jl`: `LatitudinalClimatology` on a masked `ColumnRingGrid` assigns each active column the value at its latitude. +- `test/differentiability/soil_energy_diff.jl`: run the Enzyme energy test with the new initializer to confirm the kernel-based `initialize!` is differentiable. +- `test/reactant/`: add the kernel-based initializer to one of the existing model configurations in `setup.jl` so it is exercised on the Reactant backend. +- Full suite via `Pkg.test()`, then the draft doc build. + +## Documentation changes + +- `docs/src/models/soil_model.md` ("Initializers" section): document the accepted forms of `T₀`/`Qgeo` and the declared input variables, with an `@example` using `LatitudinalClimatology`. +- `docs/src/running/initialization.md`: add a "Field initializers" subsection for `AbstractFieldInitializer`, and note that model initializers may declare input variables. +- `docs/src/extending/core_interfaces.md`: mention `variables(::AbstractInitializer)`. +- Docstrings for all new types and methods with `jldoctest` blocks where a short verifiable output exists. + +## Known limitations + +- `k_eff` stays a scalar. + A true thermal steady state would need the stratigraphy's conductivity; this is the same caveat the current docstring already states. +- `AbstractFieldInitializer` defaults are applied at `Field` construction (as all input defaults are today), i.e. before any `InputSource` runs, and only once. + That is the desired order, but it means a default that depends on other inputs is not supported. +- The `Variables` conflict rule requires matching units between an `InputSource` and the declared input; users must pass `units = u"W/m^2"` (resp. `u"°C"`) to the source. +- `LandModel` with a `SoilInitializer` is not addressed here; whichever way `initialize!` dispatches for that combination today is unchanged. + +## Future work + +- A reusable periodic *boundary condition* (`T₀(lat) + A sin(2π t / P - lon)`) to remove the matching `get_temperature_bc` duplication in the examples. +- Applying auxiliary-variable initializers in `reset!` (existing TODO in `state_variables.jl`). + +## Resolved questions (Rev 2) + +1. Variable names: `initial_surface_temperature` and `geothermal_heat_flux` (shared with the bottom BC). +2. Hydrology initializers are included. +3. `LatitudinalClimatology(T_equator, ΔT)` with the fixed form `T_equator - |ΔT sin φ|`; no generic function wrapper. + Users assign functions directly to initializer fields. +4. `AbstractFieldInitializer` is defined in core `src/initializers.jl`. From 87d4ca34a7b9be12bceae46f04627c825d8d2699 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:13:59 +0200 Subject: [PATCH 04/19] Add AGENTS.md instruction to not run Enzyme/Reactant tests locally --- AGENTS.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index cdebb9f57..2630039eb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -30,8 +30,9 @@ dynamics are allowed except in very special cases where they must be clearly doc using TestEnv; TestEnv.activate() include("test/soil/soil_energy_tests.jl") ``` -- To run the full suite, use `julia --project=. -e 'using Pkg; Pkg.test()'` (Enzyme/AD tests run - via `Pkg.test(; test_args=["enzyme"])`). +- To run the full suite, use `julia --project=. -e 'using Pkg; Pkg.test()'`. +- **Never run the Enzyme (`Pkg.test(; test_args=["enzyme"])`) or Reactant (`test/reactant/`) test suites + locally unless explicitly requested**. These tests should be left to the CI/CD pipeline. - Julia errors and associated stack traces are often very long due to long type signatures. To mitigate this, always write test output to temporary files and analyze this output using `grep` and similar tools. From cd6dc86206640bf95cd2241eb465144141c664b2 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:28:45 +0200 Subject: [PATCH 05/19] =?UTF-8?q?Add=20=CE=BBnodes=20and=20=CF=86nodes=20i?= =?UTF-8?q?mplementation=20for=20ColumnRingGrid?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/grids/column_ring_grid.jl | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/grids/column_ring_grid.jl b/src/grids/column_ring_grid.jl index 1c3d0fef8..d43bfbc31 100644 --- a/src/grids/column_ring_grid.jl +++ b/src/grids/column_ring_grid.jl @@ -118,6 +118,32 @@ Oceananigans.Grids.rname(grid::ColumnRingGrid) = Oceananigans.Grids.rname(getfie @inline Oceananigans.Grids.ηnode(i, j, k, grid::ColumnRingGrid, ℓx, ℓy, ℓz) = ηnode(i, j, k, getfield(grid, :grid), ℓx, ℓy, ℓz) @inline Oceananigans.Grids.rnode(i, j, k, grid::ColumnRingGrid, ℓx, ℓy, ℓz) = rnode(i, j, k, getfield(grid, :grid), ℓx, ℓy, ℓz) +""" + $SIGNATURES + +Longitudes (degrees) of the active (masked) columns of `grid`, as a host `Vector` ordered like the +horizontal index of the grid. This extends the Oceananigans accessor for curvilinear grids so that +`λnodes`/`φnodes` give the geographic coordinates of a horizontal position on both a +`LatitudeLongitudeGrid` and a `ColumnRingGrid`; on the latter, both are indexed by column. +""" +Oceananigans.Grids.λnodes(grid::ColumnRingGrid, ::CenterOrFace = Center(), args...; kwargs...) = first(active_column_lonlats(grid)) + +""" + $SIGNATURES + +Latitudes (degrees) of the active (masked) columns of `grid`, as a host `Vector` ordered like the +horizontal index of the grid. See [`λnodes`](@ref). +""" +Oceananigans.Grids.φnodes(grid::ColumnRingGrid, ::CenterOrFace = Center(), args...; kwargs...) = last(active_column_lonlats(grid)) + +# Host-side gather of the (longitude, latitude) pairs of the active columns, in degrees. +function active_column_lonlats(grid::ColumnRingGrid) + rings = on_architecture(CPU(), grid.rings) + mask = on_architecture(CPU(), grid.mask) + lond, latd = RingGrids.get_londlatds(rings) + return lond[mask], latd[mask] +end + # Horizontal grid metrics. Oceananigans defines these four spacings per concrete grid type and derives # all areas and volumes from them, e.g. the `Az / volume` scaling of flux boundary conditions in # `compute_z_bcs!`. The vertical spacings need no forwarding since they read `grid.z`, which is From 854d3e32e4138945073e9e5c4192a028b39d8db0 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:29:24 +0200 Subject: [PATCH 06/19] Add set! dispatches for KernelFunction --- src/utils/kernel_utils.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/utils/kernel_utils.jl b/src/utils/kernel_utils.jl index 6c5694cd9..2569efd86 100644 --- a/src/utils/kernel_utils.jl +++ b/src/utils/kernel_utils.jl @@ -86,6 +86,17 @@ function kernel(func, args...; clock = false) return KernelFunction{!clock, typeof(func), typeof(args)}(func, args) end +""" + set!(field::Field, kf::KernelFunction, fields) + set!(field::Field, kf::KernelFunction, clock, fields) + +Set `field` by evaluating the [`kernel`](@ref) function `kf` at every node of `field`, reading from +`fields` (and `clock`, for non-autonomous kernel functions). The operation is constructed on the +grid of `field` and its location. +""" +Oceananigans.Fields.set!(field::Field, kf::KernelFunction{true}, fields) = set!(field, kf(field, field.grid, nothing, fields)) +Oceananigans.Fields.set!(field::Field, kf::KernelFunction{false}, clock, fields) = set!(field, kf(field, field.grid, clock, fields)) + """ findfirst_z(i, j, condition_func, z_nodes, field) From ae351a6ccb7ec52a229981e0e94e7172cb27e591 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:30:05 +0200 Subject: [PATCH 07/19] Allow intializers to declare variables and parameters --- src/Terrarium.jl | 8 +-- src/abstract_model.jl | 19 +++++-- src/abstract_variables.jl | 8 ++- src/initializers.jl | 88 ++++++++++++++++++++++++++++- test/initializers.jl | 116 ++++++++++++++++++++++++++++++++++++++ test/reactant/runtests.jl | 1 + test/reactant/setup.jl | 14 +++++ test/runtests.jl | 4 ++ 8 files changed, 245 insertions(+), 13 deletions(-) create mode 100644 test/initializers.jl diff --git a/src/Terrarium.jl b/src/Terrarium.jl index 71647e629..c8771bb2f 100644 --- a/src/Terrarium.jl +++ b/src/Terrarium.jl @@ -22,8 +22,8 @@ using Oceananigans.Architectures: Architectures, AbstractArchitecture, CPU, GPU, using Oceananigans.Fields: Field, FunctionField, AbstractField, Center, Face, set!, compute!, interior, indices, location using Oceananigans.Forcings: Forcing, ContinuousForcing, DiscreteForcing using Oceananigans.Grids: AbstractGrid, RectilinearGrid, CallableDiscretization, ExponentialDiscretization, - Periodic, Flat, Bounded, halo_size, isrectilinear, nodes, topology, xnodes, ynodes, znodes, znode, zspacings, - ξnode, ηnode, rnode + LatitudeLongitudeGrid, Periodic, Flat, Bounded, halo_size, isrectilinear, nodes, topology, + xnodes, ynodes, znodes, λnodes, φnodes, znode, zspacings, ξnode, ηnode, rnode using Oceananigans.Operators: ∂zᵃᵃᶜ, ∂zᵃᵃᶠ, ℑzᵃᵃᶠ, Δzᵃᵃᶜ using Oceananigans.OutputReaders: FieldTimeSeries using Oceananigans.Simulations: Simulation, run!, timestepper, TimeStepWizard, conjure_time_step_wizard!, Callback, add_callback! @@ -92,7 +92,7 @@ const CenterOrFace = Union{Center, Face} export Simulation, Clock, Field, FieldTimeSeries, KernelFunctionOperation, Center, Face export CPU, GPU, ReactantState, architecture, on_architecture export Value, Flux, Gradient, ValueBoundaryCondition, GradientBoundaryCondition, FluxBoundaryCondition, NoFluxBoundaryCondition -export run!, time_step!, set!, reset!, compute!, interior, znodes, zspacings, location +export run!, time_step!, set!, reset!, compute!, interior, znodes, λnodes, φnodes, zspacings, location export TimeStepWizard, conjure_time_step_wizard!, Callback, add_callback!, IterationInterval, TimeInterval # Re-export selected types from FreezeCurves @@ -151,7 +151,7 @@ export StateVariables, get_fields include("state_variables.jl") # default initializers -export DefaultInitializer +export DefaultInitializer, LatitudinalClimatology include("initializers.jl") # boundary condition helper functions diff --git a/src/abstract_model.jl b/src/abstract_model.jl index 2ad217174..908c46e92 100644 --- a/src/abstract_model.jl +++ b/src/abstract_model.jl @@ -113,13 +113,22 @@ function compute_tendencies! end variables(::Any) = () """ - variables(obj::Union{AbstractCoupledProcesses, AbstractModel}) + variables(obj::AbstractCoupledProcesses) -Default implementation of [`variables`](@ref) for composite [`AbstractModel`](@ref) and -[`AbstractCoupledProcesses`](@ref) types that automatically collects all variables from all processes defined -as properties/fields on the given `obj`. +Default implementation of [`variables`](@ref) for composite [`AbstractCoupledProcesses`](@ref) types +that automatically collects all variables from all processes defined as properties/fields on the given `obj`. """ -variables(obj::Union{AbstractCoupledProcesses, AbstractModel}) = tuplejoin(fastmap(variables, processes(obj))...) +variables(obj::AbstractCoupledProcesses) = tuplejoin(fastmap(variables, processes(obj))...) + +""" + variables(model::AbstractModel) + +Default implementation of [`variables`](@ref) for [`AbstractModel`](@ref) types that collects all +variables from all processes defined as properties/fields on the given `model`, followed by any +variables declared by its initializer (see [`get_initializer`](@ref)). Initializers may declare +`input` variables for their parameters; see [`AbstractInitializer`](@ref). +""" +variables(model::AbstractModel) = tuplejoin(fastmap(variables, processes(model))..., variables(get_initializer(model))) # Fallback dispatches to make implementing compute_boundary_conditions! optional compute_boundary_conditions!(state, ::AbstractModel) = nothing diff --git a/src/abstract_variables.jl b/src/abstract_variables.jl index 04e5ae315..e239dff48 100644 --- a/src/abstract_variables.jl +++ b/src/abstract_variables.jl @@ -339,6 +339,10 @@ end Represents a spatially varying input (e.g. forcing) variable with the given `name` and spatial `dims`. Input variables can also be made to vary in time through the use of [`InputSource`](@ref)s. + +The `default` is applied once, when the input `Field` is constructed, via `set!(field, default)`. +It may be `nothing` (zero), a number, or any other value accepted by `set!`: a function of the node +coordinates, an array, a `Field`, or an [`AbstractFieldInitializer`](@ref). """ struct InputVariable{ name, @@ -346,12 +350,12 @@ struct InputVariable{ UT <: Units, Var <: Variable{name, VL, UT}, BT <: DomainSets.AbstractInterval, - Def <: Union{Nothing, Number, Function}, + Def, } <: AbstractProcessVariable{name, VL, UT} "State variable" var::Var - "Default value or function initializer" + "Default value or initializer; `nothing`, a number, or any argument accepted by `set!`" default::Def "Variable bounds" diff --git a/src/initializers.jl b/src/initializers.jl index 301ec943e..fe92ffcab 100644 --- a/src/initializers.jl +++ b/src/initializers.jl @@ -2,8 +2,9 @@ """ Base type for model initializers. Implementations should provide a dispatch of the `initialize!(state, model::M, init::I)` method where -`M` corresponds to the model type and `I` to the initializer. An implementation of `get_field_initializers` can also be provided which -returns a `NamedTuple` of initializer functions for individual state variable fields. +`M` corresponds to the model type and `I` to the initializer. Initializers may additionally implement +[`variables`](@ref) to declare `input` variables for their (possibly spatially varying) parameters; +these are collected by `variables(model)` alongside the process variables. """ abstract type AbstractInitializer{NF} end @@ -32,3 +33,86 @@ Marker type for a no-op initializer that leaves all `Field`s set to their defaul struct DefaultInitializer{NF} <: AbstractInitializer{NF} end DefaultInitializer(::Type{NF}) where {NF} = DefaultInitializer{NF}() + +# Field initializers + +""" +Base type for reusable initializers of a single `Field`. Implementations should provide a dispatch of +`initialize!(field, grid, init::AbstractFieldInitializer)`. The default evaluates `init` as a callable +at the node coordinates of `field`, i.e. `init(coords...)`, via `set!`. + +A field initializer is accepted everywhere a `set!`-compatible value is: as the `default` of an `input` +variable, as an entry of the `initializers` keyword argument of [`initialize`](@ref), or as a parameter +of a model [`AbstractInitializer`](@ref). In the latter case it may declare `@param` fields, which are +then exposed as model parameters, and the model initializer re-evaluates it at every `initialize!`. +""" +abstract type AbstractFieldInitializer end + +""" + $TYPEDSIGNATURES + +Initialize `field` on `grid` with the given field initializer. The default implementation evaluates +`init(coords...)` at the node coordinates of `field`. +""" +initialize!(field::AbstractField, grid::AbstractGrid, init::AbstractFieldInitializer) = set!(field, (coords...) -> init(coords...)) + +Oceananigans.Fields.set!(field::Field, init::AbstractFieldInitializer) = initialize!(field, field.grid, init) + +""" + $TYPEDSIGNATURES + +Re-evaluate `init` into `field` if it is an [`AbstractFieldInitializer`](@ref); all other initial values +(numbers, functions, arrays, `Field`s) are applied only once as construction-time defaults and are left +untouched here so that they remain overridable by [`InputSource`](@ref)s. +""" +reinitialize!(field::AbstractField, init::AbstractFieldInitializer) = set!(field, init) +reinitialize!(::AbstractField, ::Any) = nothing + +""" + $TYPEDEF + +Latitude-dependent surface temperature climatology (°C): + + T(φ) = T_equator - |ΔT sin φ| + +where `φ` is the latitude in degrees. Supported on a [`ColumnRingGrid`](@ref), where it is evaluated +on the host at the latitude of each active column, and on a `LatitudeLongitudeGrid`, where it is +evaluated at the latitude of each node. + +Properties: +$TYPEDFIELDS +""" +@parameterized @kwdef struct LatitudinalClimatology{NF} <: AbstractFieldInitializer + "Mean annual surface temperature at the equator (°C)" + @param T_equator::NF = 20.0 + "Temperature difference between the equator and the poles (°C)" + @param ΔT::NF = 40.0 +end + +LatitudinalClimatology(::Type{NF}; kwargs...) where {NF} = LatitudinalClimatology{NF}(; kwargs...) + +(init::LatitudinalClimatology)(latitude) = init.T_equator - abs(init.ΔT * sind(latitude)) + +function initialize!(field::AbstractField, grid::ColumnRingGrid, init::LatitudinalClimatology) + column_values = reshape(init.(φnodes(grid)), :, 1, 1) + # broadcast the per-column values over the remaining (vertical) extent of the field, if any + set!(field, repeat(column_values, 1, size(field, 2), size(field, 3))) + return nothing +end + +initialize!(field::AbstractField, ::LatitudeLongitudeGrid, init::LatitudinalClimatology) = set!(field, (λ, φ, args...) -> init(φ)) + +initialize!(field::AbstractField, grid::AbstractLandGrid, init::LatitudinalClimatology) = initialize!(field, ground_domain(grid), init) + +# Parameters of initializers are collected from their fields, like for processes. Numbers become +# parameters, parameterized field initializers contribute theirs, and functions contribute none. +function ParameterEditing.parameters(::Type{PT}, init::AbstractInitializer; kwargs...) where {PT <: ParameterEditing.AbstractParam} + init_params = map(fieldnames(typeof(init))) do name + name => ParameterEditing.parameters(PT, getproperty(init, name); kwargs...) + end + nonempty_params = filter(p -> length(p[2]) > 0, init_params) + return ParameterEditing.ParameterTable((; nonempty_params...)) +end + +# `Field`s used as initial values are state, not parameters. +ParameterEditing.parameters(::Type{PT}, ::AbstractField; kwargs...) where {PT <: ParameterEditing.AbstractParam} = (;) diff --git a/test/initializers.jl b/test/initializers.jl new file mode 100644 index 000000000..0179683e8 --- /dev/null +++ b/test/initializers.jl @@ -0,0 +1,116 @@ +using Terrarium +using Terrarium: RingGrids, parameters, variables, input_variables, varname +using Oceananigans.Grids: LatitudeLongitudeGrid, Bounded +using Test + +@testset "Initializers declare input variables" begin + NF = Float32 + grid = ColumnGrid(CPU(), NF, ExponentialSpacing(N = 10)) + model = SoilModel(grid) + # the default initializer contributes no variables + @test :initial_surface_temperature ∉ map(varname, input_variables(variables(model))) + model = SoilModel(grid; initializer = SoilInitializer(NF)) + names = map(varname, input_variables(variables(model))) + @test :initial_surface_temperature ∈ names + @test :geothermal_heat_flux ∈ names + @test :vadose_zone_saturation ∈ names + @test :water_table_depth ∈ names +end + +@testset "Scalar QuasiThermalSteadyState and SaturationWaterTable" begin + NF = Float32 + grid = ColumnGrid(CPU(), NF, ExponentialSpacing(N = 10)) + T₀ = NF(-1) + Qgeo = NF(0.05) + energy = QuasiThermalSteadyState(NF; T₀, Qgeo) + hydrology = SaturationWaterTable(NF; water_table_depth = NF(1), vadose_zone_saturation = NF(0.5)) + model = SoilModel(grid; initializer = SoilInitializer(NF; energy, hydrology)) + integrator = initialize(model) + state = integrator.state + z = znodes(state.temperature) + @test interior(state.temperature)[1, 1, :] ≈ T₀ .- Qgeo .* z + @test all(interior(state.initial_surface_temperature) .== T₀) + @test all(interior(state.geothermal_heat_flux) .== Qgeo) + sat = interior(state.saturation_water_ice)[1, 1, :] + @test all(sat[z .<= -1] .== 1) + @test all(sat[z .> -1] .== NF(0.5)) + # re-initialization reproduces the same state + T_before = copy(interior(state.temperature)) + Terrarium.initialize!(integrator) + @test interior(state.temperature) == T_before +end + +@testset "Column-varying initial values" begin + NF = Float32 + grid = ColumnGrid(CPU(), NF, ExponentialSpacing(N = 5), 3) + # function of the column coordinate + energy = QuasiThermalSteadyState(NF; T₀ = x -> NF(x), Qgeo = zero(NF)) + integrator = initialize(SoilModel(grid; initializer = SoilInitializer(NF; energy))) + T = interior(integrator.state.temperature) + @test T[:, 1, end] ≈ interior(integrator.state.initial_surface_temperature)[:, 1, 1] + @test T[1, 1, end] < T[2, 1, end] < T[3, 1, end] + # Field with the same values + T₀_field = Field(grid, Terrarium.Ground(XY())) + set!(T₀_field, x -> NF(x)) + energy = QuasiThermalSteadyState(NF; T₀ = T₀_field, Qgeo = zero(NF)) + integrator_field = initialize(SoilModel(grid; initializer = SoilInitializer(NF; energy))) + @test interior(integrator_field.state.temperature) == T + # fields are not parameters + @test :T₀ ∉ keys(vec(parameters(integrator_field.model)).initializer.energy) +end + +@testset "InputSource overrides the geothermal heat flux default" begin + NF = Float32 + rings = RingGrids.FullGaussianGrid(4) + grid = ColumnRingGrid(CPU(), NF, UniformSpacing(Δz = NF(0.5), N = 4), rings) + model = SoilModel(grid; initializer = SoilInitializer(NF; energy = QuasiThermalSteadyState(NF; T₀ = zero(NF), Qgeo = NF(0.02)))) + Qgeo_field = Field(grid, Terrarium.Ground(XY())) + set!(Qgeo_field, NF(0.1)) + source = InputSource(grid, Qgeo_field; name = :geothermal_heat_flux, domain = Terrarium.Ground(), units = u"W/m^2") + integrator = initialize(model; inputs = InputSources(source)) + @test all(interior(integrator.state.geothermal_heat_flux) .== NF(0.1)) + z = znodes(integrator.state.temperature) + @test interior(integrator.state.temperature)[1, 1, :] ≈ -NF(0.1) .* z + # mismatched units are reported as a conflict + bad_source = InputSource(grid, Qgeo_field; name = :geothermal_heat_flux, domain = Terrarium.Ground()) + @test_throws ErrorException initialize(model; inputs = InputSources(bad_source)) +end + +@testset "LatitudinalClimatology" begin + NF = Float32 + climatology = LatitudinalClimatology(NF; T_equator = NF(10), ΔT = NF(40)) + @test climatology(0) == NF(10) + @test climatology(90) ≈ NF(-30) + @test climatology(-90) ≈ NF(-30) + + rings = RingGrids.FullGaussianGrid(4) + grid = ColumnRingGrid(CPU(), NF, UniformSpacing(Δz = NF(0.5), N = 4), rings) + @test length(λnodes(grid)) == length(φnodes(grid)) == size(grid, 1) + model = SoilModel(grid; initializer = SoilInitializer(NF; energy = QuasiThermalSteadyState(NF; T₀ = climatology))) + integrator = initialize(model) + T_surface = interior(integrator.state.initial_surface_temperature)[:, 1, 1] + @test T_surface ≈ climatology.(φnodes(grid)) + + # also as a direct field initializer + integrator_direct = initialize(SoilModel(grid); initializers = (temperature = climatology,)) + @test interior(integrator_direct.state.temperature)[:, 1, 1] ≈ climatology.(φnodes(grid)) + + # parameters are exposed and re-evaluated on reconstruction + ps = vec(parameters(model)) + @test haskey(ps.initializer.energy.T₀, :T_equator) + ps.initializer.energy.T₀.T_equator = NF(30) + integrator = initialize(integrator, ps) + @test interior(integrator.state.initial_surface_temperature)[:, 1, 1] ≈ LatitudinalClimatology(NF; T_equator = NF(30)).(φnodes(grid)) + + # masked grid: only active columns are evaluated + mask = convert.(Bool, ones(rings)) + mask[1:5] .= false + masked_grid = ColumnRingGrid(CPU(), NF, UniformSpacing(Δz = NF(0.5), N = 4), rings, mask) + integrator_masked = initialize(SoilModel(masked_grid; initializer = SoilInitializer(NF; energy = QuasiThermalSteadyState(NF; T₀ = climatology)))) + @test interior(integrator_masked.state.initial_surface_temperature)[:, 1, 1] ≈ climatology.(φnodes(masked_grid)) + + # LatitudeLongitudeGrid + llg = LatitudeLongitudeGrid(CPU(), NF; size = (2, 3, 5), longitude = (0, 10), latitude = (0, 60), z = (-1, 0), topology = (Bounded, Bounded, Bounded)) + integrator_llg = initialize(SoilModel(llg; initializer = SoilInitializer(NF; energy = QuasiThermalSteadyState(NF; T₀ = climatology)))) + @test interior(integrator_llg.state.initial_surface_temperature)[1, :, 1] ≈ climatology.(φnodes(llg, Center())) +end diff --git a/test/reactant/runtests.jl b/test/reactant/runtests.jl index 0a62046e1..b08ab6cf5 100644 --- a/test/reactant/runtests.jl +++ b/test/reactant/runtests.jl @@ -37,6 +37,7 @@ include("setup.jl") test_model(:soil_heat_column) test_model(:soil_heat_column_stretched) # array-valued (ExponentialSpacing) vertical coordinates test_model(:soil_heat_global) + test_model(:soil_heat_global_initializer) # model initializer with input variables and a LatitudinalClimatology test_model(:snow_column) # standalone single-layer snow (closure + tendencies) test_model(:land_soil_snow) test_model(:vegetation_column) # standalone PrescribedVegetation, constant LAI diff --git a/test/reactant/setup.jl b/test/reactant/setup.jl index 3a15ca9d3..39312f771 100644 --- a/test/reactant/setup.jl +++ b/test/reactant/setup.jl @@ -69,6 +69,20 @@ function build_model(::Val{:soil_heat_global}, arch, NF) return (; model, boundary_conditions = bcs, initializers = inits, Δt = NF(600)) end +# --- :soil_heat_global_initializer — SoilModel on a ColumnRingGrid with a model initializer ---- +# Initial temperature from `QuasiThermalSteadyState` with a `LatitudinalClimatology` surface +# temperature, and saturation from `SaturationWaterTable`. Exercises the initializer input +# variables and the `KernelFunction`-based `set!` on the Reactant backend. +function build_model(::Val{:soil_heat_global_initializer}, arch, NF) + rings = RingGrids.FullGaussianGrid(4) + grid = ColumnRingGrid(arch, NF, UniformSpacing(Δz = NF(0.2), N = 20), rings) + energy = QuasiThermalSteadyState(NF; T₀ = LatitudinalClimatology(NF), Qgeo = NF(0.05)) + hydrology = SaturationWaterTable(NF; water_table_depth = NF(2)) + model = SoilModel(grid; initializer = SoilInitializer(NF; energy, hydrology)) + bcs = PrescribedSurfaceTemperature(:T_ub, NF(1)) + return (; model, boundary_conditions = bcs, initializers = (;), Δt = NF(600)) +end + # --- :snow_column — minimal single-column standalone SnowModel -------------------------- # Exercises the snow closure (energy↔temperature) and the mass/energy tendencies under Reactant. # Boundary heat fluxes and SWE/temperature are prescribed as constant input/initial fields. diff --git a/test/runtests.jl b/test/runtests.jl index 432b8ff9c..4ca67f0ca 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -28,6 +28,10 @@ if MAIN_TESTS include("state_variables.jl") end + @testset "Initializers" begin + include("initializers.jl") + end + @testset "Boundary conditions" begin include("boundary_conditions.jl") end From bf767b3844aee7651f14386e006477706147b855 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:30:25 +0200 Subject: [PATCH 08/19] Generalize soi linitializers to accept spatially varying inputs --- src/models/soil/soil_model_init.jl | 183 ++++++++++++++++++++++++----- 1 file changed, 155 insertions(+), 28 deletions(-) diff --git a/src/models/soil/soil_model_init.jl b/src/models/soil/soil_model_init.jl index 880edbb0a..51045c01b 100644 --- a/src/models/soil/soil_model_init.jl +++ b/src/models/soil/soil_model_init.jl @@ -28,6 +28,8 @@ function SoilInitializer( return SoilInitializer(energy, hydrology, biogeochem) end +variables(init::SoilInitializer) = tuplejoin(variables(init.energy), variables(init.hydrology), variables(init.biogeochem)) + function initialize!(state, model::AbstractModel, init::SoilInitializer) initialize!(state, model, init.hydrology) initialize!(state, model, init.biogeochem) @@ -40,46 +42,121 @@ end """ $TYPEDEF -Initializer for soil/ground temperature that sets the temperature profile to a constant value. +Initializer for soil/ground temperature that sets the temperature profile of each column to the +value of the `initial_surface_temperature` input variable, whose default is `T₀`. + +`T₀` may be a number, a function of the node coordinates, an array, a `Field`, or an +[`AbstractFieldInitializer`](@ref). All but the last are applied once as the default of the input +variable and can be overridden by an [`InputSource`](@ref) named `initial_surface_temperature`; +a field initializer is re-evaluated at every `initialize!` so that its parameters take effect. Properties: $TYPEDFIELDS """ -@kwdef struct ConstantSoilTemperature{NF} <: AbstractInitializer{NF} - T₀::NF = 0.0 +struct ConstantSoilTemperature{NF, T0} <: AbstractInitializer{NF} + "Initial surface temperature (°C)" + T₀::T0 end """ -Creates a constant soil temperature initializer. + $TYPEDSIGNATURES + +Creates a constant soil temperature initializer with the given surface temperature `T₀` (°C). +""" +ConstantSoilTemperature(::Type{NF}; T₀ = zero(NF)) where {NF} = ConstantSoilTemperature{NF, typeof(T₀)}(T₀) +ConstantSoilTemperature(T₀::NF) where {NF <: AbstractFloat} = ConstantSoilTemperature{NF, NF}(T₀) + +# Preserve `NF` when reconstructing from parameters (see `ParameterEditing.reconstruct`) +ConstructionBase.constructorof(::Type{<:ConstantSoilTemperature{NF}}) where {NF} = T₀ -> ConstantSoilTemperature{NF, typeof(T₀)}(T₀) + +variables(init::ConstantSoilTemperature) = ( + input(:initial_surface_temperature, Ground(XY()), default = init.T₀, units = u"°C", desc = "Initial surface temperature of the soil column"), +) + +function initialize!(state, model::AbstractModel, init::ConstantSoilTemperature) + reinitialize!(state.initial_surface_temperature, init.T₀) + set!(state.temperature, kernel(compute_constant_temperature, init), state) + return nothing +end + """ -ConstantSoilTemperature(::Type{NF}; kwargs...) where {NF} = ConstantSoilTemperature{NF}(; kwargs...) + $TYPEDSIGNATURES -initialize!(state, ::AbstractModel, init::ConstantSoilTemperature) = set!(state.temperature, init.T₀) +Initial temperature at cell `(i, j, k)` taken from the `initial_surface_temperature` (°C) field. +""" +@propagate_inbounds compute_constant_temperature(i, j, k, grid, fields, ::ConstantSoilTemperature) = fields.initial_surface_temperature[i, j, 1] """ $TYPEDEF Initializer that sets soil/ground temperature to a thermal quasi-steady state based on the given -surface temperature, geothermal heat flux, and bulk (constant) thermal conductivity. Note that this is -not a *true* thermal steady state, which would require iterative calculation of the thermal conductivity -from the soil properties and initial temperature profile. +surface temperature, geothermal heat flux, and bulk (constant) thermal conductivity: + + T(z) = T₀ - Qgeo / k_eff * z + +with depth `z ≤ 0` (m). Note that this is not a *true* thermal steady state, which would require +iterative calculation of the thermal conductivity from the soil properties and initial temperature profile. + +The surface temperature and the geothermal heat flux are declared as the input variables +`initial_surface_temperature` (°C) and `geothermal_heat_flux` (W/m²) with defaults `T₀` and `Qgeo`. +The latter is the same variable read by the [`GeothermalHeatFlux`](@ref) bottom boundary condition, +so that the initial profile and the boundary condition are consistent. Each default may be a number, +a function of the node coordinates, an array, a `Field`, or an [`AbstractFieldInitializer`](@ref). +All but the last are applied once and can be overridden by an [`InputSource`](@ref) of the same name +(with matching units); a field initializer is re-evaluated at every `initialize!` so that its +parameters take effect. Properties: $TYPEDFIELDS """ -@kwdef struct QuasiThermalSteadyState{NF} <: AbstractInitializer{NF} - T₀::NF = 0.0 - Qgeo::NF = 0.02 - k_eff::NF = 1.0 +struct QuasiThermalSteadyState{NF, T0, QG} <: AbstractInitializer{NF} + "Initial surface temperature (°C)" + T₀::T0 + + "Geothermal heat flux (W/m²)" + Qgeo::QG + + "Bulk thermal conductivity (W/m/K)" + k_eff::NF +end + +""" + $TYPEDSIGNATURES + +Creates a quasi-steady-state soil temperature initializer with surface temperature `T₀` (°C), +geothermal heat flux `Qgeo` (W/m²), and bulk thermal conductivity `k_eff` (W/m/K). +""" +function QuasiThermalSteadyState(::Type{NF}; T₀ = zero(NF), Qgeo = convert(NF, 1 // 50), k_eff = one(NF)) where {NF} + return QuasiThermalSteadyState{NF, typeof(T₀), typeof(Qgeo)}(T₀, Qgeo, convert(NF, k_eff)) end -QuasiThermalSteadyState(::Type{NF}; kwargs...) where {NF} = QuasiThermalSteadyState{NF}(; kwargs...) +ConstructionBase.constructorof(::Type{<:QuasiThermalSteadyState{NF}}) where {NF} = (T₀, Qgeo, k_eff) -> QuasiThermalSteadyState{NF, typeof(T₀), typeof(Qgeo)}(T₀, Qgeo, k_eff) + +variables(init::QuasiThermalSteadyState) = ( + input(:initial_surface_temperature, Ground(XY()), default = init.T₀, units = u"°C", desc = "Initial surface temperature of the soil column"), + input(:geothermal_heat_flux, Ground(XY()), default = init.Qgeo, units = u"W/m^2", desc = "Geothermal heat flux at the bottom of the soil column"), +) -function initialize!(state, ::AbstractModel, init::QuasiThermalSteadyState) - set!(state.temperature, (x, z) -> init.T₀ - init.Qgeo / init.k_eff * z) +function initialize!(state, model::AbstractModel, init::QuasiThermalSteadyState) + reinitialize!(state.initial_surface_temperature, init.T₀) + reinitialize!(state.geothermal_heat_flux, init.Qgeo) + set!(state.temperature, kernel(compute_quasi_steady_state_temperature, init), state) return nothing end +""" + $TYPEDSIGNATURES + +Quasi-steady-state temperature `T₀ - Qgeo / k_eff * z` at cell `(i, j, k)` from the +`initial_surface_temperature` (°C) and `geothermal_heat_flux` (W/m²) fields. +""" +@propagate_inbounds function compute_quasi_steady_state_temperature(i, j, k, grid, fields, init::QuasiThermalSteadyState) + z = znode(i, j, k, grid, Center(), Center(), Center()) + T₀ = fields.initial_surface_temperature[i, j, 1] + Qgeo = fields.geothermal_heat_flux[i, j, 1] + return T₀ - Qgeo / init.k_eff * z +end + """ $TYPEDEF @@ -117,36 +194,86 @@ end """ $TYPEDEF -Initializer for soil water/ice sets the saturation profile to a constant value. +Initializer for soil water/ice that sets the saturation profile of each column to the value of the +`initial_saturation` input variable, whose default is `sat`. See [`QuasiThermalSteadyState`](@ref) +for the accepted forms of the default. Properties: $TYPEDFIELDS """ -@kwdef struct ConstantSaturation{NF} <: AbstractInitializer{NF} - sat::NF = 1.0 +struct ConstantSaturation{NF, S} <: AbstractInitializer{NF} + "Initial water/ice saturation (-)" + sat::S +end + +ConstantSaturation(::Type{NF}; sat = one(NF)) where {NF} = ConstantSaturation{NF, typeof(sat)}(sat) + +ConstructionBase.constructorof(::Type{<:ConstantSaturation{NF}}) where {NF} = sat -> ConstantSaturation{NF, typeof(sat)}(sat) + +variables(init::ConstantSaturation) = ( + input(:initial_saturation, Ground(XY()), default = init.sat, bounds = UnitInterval, desc = "Initial water/ice saturation of the soil column"), +) + +function initialize!(state, model::AbstractModel, init::ConstantSaturation) + reinitialize!(state.initial_saturation, init.sat) + set!(state.saturation_water_ice, kernel(compute_constant_saturation, init), state) + return nothing end -ConstantSaturation(::Type{NF}; kwargs...) where {NF} = ConstantSaturation{NF}(; kwargs...) +""" + $TYPEDSIGNATURES -initialize!(state, ::AbstractModel, init::ConstantSaturation) = set!(state.saturation_water_ice, init.sat) +Initial saturation at cell `(i, j, k)` taken from the `initial_saturation` field. +""" +@propagate_inbounds compute_constant_saturation(i, j, k, grid, fields, ::ConstantSaturation) = fields.initial_saturation[i, j, 1] """ $TYPEDEF Simple initialization scheme for soil/ground saturation that sets the initial water table at the -given depth and the saturation level in all layers in the vadose (unsaturated) to a constant value. +given depth and the saturation level in all layers in the vadose (unsaturated) zone to a constant value. +Both are declared as the input variables `water_table_depth` (m) and `vadose_zone_saturation` (-) +with defaults taken from the corresponding properties. See [`QuasiThermalSteadyState`](@ref) for the +accepted forms of the defaults. Properties: $TYPEDFIELDS """ -@kwdef struct SaturationWaterTable{NF} <: AbstractInitializer{NF} - vadose_zone_saturation::NF = 0.75 - water_table_depth::NF = 5.0 +struct SaturationWaterTable{NF, S, D} <: AbstractInitializer{NF} + "Saturation in the vadose zone above the water table (-)" + vadose_zone_saturation::S + + "Depth of the water table below the surface (m)" + water_table_depth::D end -SaturationWaterTable(::Type{NF}; kwargs...) where {NF} = SaturationWaterTable{NF}(; kwargs...) +function SaturationWaterTable(::Type{NF}; vadose_zone_saturation = convert(NF, 3 // 4), water_table_depth = NF(5)) where {NF} + return SaturationWaterTable{NF, typeof(vadose_zone_saturation), typeof(water_table_depth)}(vadose_zone_saturation, water_table_depth) +end + +ConstructionBase.constructorof(::Type{<:SaturationWaterTable{NF}}) where {NF} = (sat, depth) -> SaturationWaterTable{NF, typeof(sat), typeof(depth)}(sat, depth) + +variables(init::SaturationWaterTable) = ( + input(:vadose_zone_saturation, Ground(XY()), default = init.vadose_zone_saturation, bounds = UnitInterval, desc = "Initial saturation in the vadose zone above the water table"), + input(:water_table_depth, Ground(XY()), default = init.water_table_depth, units = u"m", bounds = Nonnegative, desc = "Initial depth of the water table below the surface"), +) -function initialize!(state, ::AbstractModel, init::SaturationWaterTable{NF}) where {NF} - set!(state.saturation_water_ice, (x, z) -> z <= -init.water_table_depth ? one(NF) : init.vadose_zone_saturation) +function initialize!(state, model::AbstractModel, init::SaturationWaterTable) + reinitialize!(state.vadose_zone_saturation, init.vadose_zone_saturation) + reinitialize!(state.water_table_depth, init.water_table_depth) + set!(state.saturation_water_ice, kernel(compute_water_table_saturation, init), state) return nothing end + +""" + $TYPEDSIGNATURES + +Initial saturation at cell `(i, j, k)`: unity at and below the `water_table_depth`, otherwise the +`vadose_zone_saturation`. +""" +@propagate_inbounds function compute_water_table_saturation(i, j, k, grid, fields, ::SaturationWaterTable) + z = znode(i, j, k, grid, Center(), Center(), Center()) + depth = fields.water_table_depth[i, j, 1] + sat = fields.vadose_zone_saturation[i, j, 1] + return ifelse(z <= -depth, one(sat), sat) +end From 0060aefb767b321e23168d6444ffb94d17ce231a Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:31:02 +0200 Subject: [PATCH 09/19] Update example scripts to use new initializers --- examples/simulations/land_global_era5.jl | 15 +++--- examples/simulations/soil_heat_global.jl | 47 +++++++++---------- .../simulations/soil_heat_global_soilgrids.jl | 30 +++++------- examples/simulations/speedy_wet_land.jl | 26 ++++------ 4 files changed, 51 insertions(+), 67 deletions(-) diff --git a/examples/simulations/land_global_era5.jl b/examples/simulations/land_global_era5.jl index 668c15f3c..cee7225cf 100644 --- a/examples/simulations/land_global_era5.jl +++ b/examples/simulations/land_global_era5.jl @@ -178,16 +178,17 @@ inputs = InputSources( t2m_mean = mean(t2m; dims = Ti) # average over time axis t2m_mean_C = Terrarium.kelvin_to_celsius.(constants, t2m_mean) t2m_mean_field = RingGrids.Field(Array(reshape(t2m_mean_C, :)), Terrarium.native_grid(ERA5LandForcings())) -t2m_mean_on_grid = RingGrids.interpolate(on_architecture(CPU(), grid.rings), t2m_mean_field)[land_mask_cpu] +t2m_mean_on_grid = RingGrids.interpolate(on_architecture(CPU(), grid.rings), t2m_mean_field) +T₀_field = Field(on_architecture(arch, t2m_mean_on_grid), grid) # surface field over the land columns -function initial_soil_temperature(x, z) - T₀ = t2m_mean_on_grid[round(Int, x)] - T = T₀ - NF(0.05) * z - return T -end +# The regridded mean air temperature `Field` is passed directly as the surface temperature of +# [`QuasiThermalSteadyState`](@ref); the geothermal gradient follows from `Qgeo / k_eff`. Since the +# initializer is part of the model, we rebuild the `LandModel` with it. +energy_initializer = QuasiThermalSteadyState(NF; T₀ = T₀_field, Qgeo = NF(0.05)) +initializer = SoilInitializer(NF; energy = energy_initializer, hydrology = DefaultInitializer(NF)) +land = LandModel(grid; soil, vegetation, snow, atmosphere, initializer) initializers = ( - temperature = initial_soil_temperature, saturation_water_ice = (x, z) -> min(one(NF), NF(0.6) - NF(0.05) * z), ) integrator = initialize(land; inputs, initializers) diff --git a/examples/simulations/soil_heat_global.jl b/examples/simulations/soil_heat_global.jl index e54864c56..0a9b87d46 100644 --- a/examples/simulations/soil_heat_global.jl +++ b/examples/simulations/soil_heat_global.jl @@ -41,49 +41,47 @@ land_mask = land_sea_frac_N72 .> 0.5 land_mask_cpu = on_architecture(CPU(), land_mask) grid = ColumnRingGrid(arch, NF, ExponentialSpacing(N = 30), land_mask.grid, land_mask) grid_lon, grid_lat = RingGrids.get_lonlats(grid.rings) # in radians, on CPU +grid_latd = rad2deg.(grid_lat) # latitude in degrees, as expected by `LatitudinalClimatology` # Remember from the documentation section on [grids](@ref Grids), that the `x`-axis of the Oceananigans [`RectilinearGrid`](@extref Oceananigans.Grids.RectilinearGrid) # corresponds to a single index following the ring order (for more details, see the [corresponding section in the # RingGrids.jl documentation](https://speedyweather.github.io/SpeedyWeatherDocumentation/stable/ringgrids/#Indexing-Fields)). -# Now we create our [`SoilModel`](@ref), this time without a model initializer: -model = SoilModel(grid) - # To make the simulation a bit more interesting, we will use spatially periodic initial and boundary conditions. # The climatology will be determined by latitude with a maximum of 20 °C at the equator and minimum of -20°C at -# the poles. -mean_annual_temperature(lat) = 20 - abs(40 * sin(lat)) # maximum at equator +# the poles, as provided by [`LatitudinalClimatology`](@ref): +climatology = LatitudinalClimatology(NF) -fig = heatmap(RingGrids.Field(mean_annual_temperature.(grid_lat), grid.rings)) +fig = heatmap(RingGrids.Field(climatology.(grid_latd), grid.rings)) DisplayAs.PNG(fig) #hide -# The initial temperature profiles will be linear temperature profiles similar to what we would get from -# [`QuasiThermalSteadyState`](@ref) but with a hardcoded geothermal gradient of 0.05 K/m. We don't use -# the `QuasiThermalSteadyState` initializer here because it does not (yet) support spatially variable parameters. -function initial_soil_temperature(x, z) - latᵢ = lat_masked[round(Int, x)] - T₀ = mean_annual_temperature(latᵢ) - T = T₀ - 0.05 * z - return T -end +# The initial temperature profiles are linear in depth, as given by [`QuasiThermalSteadyState`](@ref). +# Its surface temperature `T₀` accepts spatially varying values, so we pass the climatology directly; +# with a geothermal heat flux of 0.05 W/m² and a bulk thermal conductivity of 1 W/m/K, this yields a +# geothermal gradient of 0.05 K/m. We leave the (unused) hydrology at its default (dry) state. +# Now we create our [`SoilModel`](@ref) with this initializer: +energy_initializer = QuasiThermalSteadyState(NF; T₀ = climatology, Qgeo = NF(0.05)) +initializer = SoilInitializer(NF; energy = energy_initializer, hydrology = DefaultInitializer(NF)) +model = SoilModel(grid; initializer) # We will impose a periodic temperature boundary condition at the surface to represent the daily cycle. # We can specify it directly as a continuous function thanks to the power of `Oceananigans` `Field`s. # However, we will need to use an enclosing function here to i) copy the vector of latitudes onto the # device specified by `arch`, and ii) ensure that the compiler is able to infer the correct type of the # coordinate values in the boundary condition function `periodic_bc`, which returns a temperature value -# based on the coordinate `x` of the `RectiLinearGrid` and the time `t` (s). -function get_temperature_bc(lon::AbstractVector, lat::AbstractVector, amplitude = 10.0) +# based on the coordinate `x` of the `RectiLinearGrid` and the time `t` (s). The climatology struct +# is captured by the closure as well; it is a plain `isbits` value, so this is safe on the GPU. +function get_temperature_bc(lon::AbstractVector, latd::AbstractVector, amplitude = 10.0) ## make sure coordinate arrays are on the same device lon_device = on_architecture(arch, NF.(lon)) - lat_device = on_architecture(arch, NF.(lat)) + lat_device = on_architecture(arch, NF.(latd)) ## function matching the expected signature for boundary conditions on a column-based grid function periodic_bc(x::NF, t::NF) where {NF} ## x coordinate is just the grid cell index lonₓ = lon_device[round(Int, x)] latₓ = lat_device[round(Int, x)] - ## use climatology at latₓ as the mean of BC - T₀ = mean_annual_temperature(latₓ) + ## use climatology at latₓ (degrees) as the mean of BC + T₀ = climatology(latₓ) seconds_per_day = NF(24 * 3600) ## shift BC by longitude in radians to (roughly) mimic the global daily cycle T = T₀ + NF(amplitude) * sin(2π * t / seconds_per_day - lonₓ) @@ -94,12 +92,11 @@ end lon_masked = grid_lon[land_mask_cpu] -lat_masked = grid_lat[land_mask_cpu] # mask out non-land points -bc = PrescribedSurfaceTemperature(:T_ub, get_temperature_bc(lon_masked, lat_masked)) -inits = (temperature = initial_soil_temperature,) +latd_masked = grid_latd[land_mask_cpu] # mask out non-land points +bc = PrescribedSurfaceTemperature(:T_ub, get_temperature_bc(lon_masked, latd_masked)) -# We are finally ready to initialize our model with the above initial and boundary conditions: -integrator = initialize(model, boundary_conditions = bc, initializers = inits) +# We are finally ready to initialize our model with the above boundary conditions: +integrator = initialize(model, boundary_conditions = bc) # Let's already plot the initial surface temperature state to see what it looks like: T_surface_initial = RingGrids.Field(arch, interior(integrator.state.ground_temperature), grid) diff --git a/examples/simulations/soil_heat_global_soilgrids.jl b/examples/simulations/soil_heat_global_soilgrids.jl index b01eabb62..e3b420b40 100644 --- a/examples/simulations/soil_heat_global_soilgrids.jl +++ b/examples/simulations/soil_heat_global_soilgrids.jl @@ -41,6 +41,7 @@ land_sea_frac_field = RingGrids.FullGaussianGrid(Matrix(land_sea_frac), input_as land_mask = land_sea_frac_field .> 0.5 grid = ColumnRingGrid(arch, NF, ExponentialSpacing(N = 30), land_mask.grid, land_mask) grid_lon, grid_lat = RingGrids.get_lonlats(grid.rings) # in radians +grid_latd = rad2deg.(grid_lat) # latitude in degrees # ## Loading soil texture from SoilGrids 2.0 # SoilGrids 2.0 provides global predictions of soil properties on a ~10 km grid for the six @@ -85,29 +86,25 @@ DisplayAs.PNG(fig) #hide porosity = SoilPorositySURFEX(eltype(grid)) strat = SoilGridsStratigraphy(eltype(grid); porosity) soil = SoilEnergyWaterCarbon(eltype(grid); strat) -model = SoilModel(grid; soil) # We reuse the simple latitude-dependent climatology from the [global example](@ref soil_heat_global) -# for the initial and boundary conditions. -mean_annual_temperature(lat) = 20 - abs(40 * sin(lat)) +# for the initial and boundary conditions. The initial temperature profile is set by +# [`QuasiThermalSteadyState`](@ref) with the climatology as its surface temperature. +climatology = LatitudinalClimatology(NF) +energy_initializer = QuasiThermalSteadyState(NF; T₀ = climatology, Qgeo = NF(0.05)) +initializer = SoilInitializer(NF; energy = energy_initializer, hydrology = DefaultInitializer(NF)) +model = SoilModel(grid; soil, initializer) lon_masked = grid_lon[land_mask] -lat_masked = grid_lat[land_mask] +latd_masked = grid_latd[land_mask] -function initial_soil_temperature(x, z) - latᵢ = lat_masked[round(Int, x)] - T₀ = mean_annual_temperature(latᵢ) - T = T₀ - 0.05 * z - return T -end - -function get_temperature_bc(lon::AbstractVector, lat::AbstractVector, amplitude = 10.0) +function get_temperature_bc(lon::AbstractVector, latd::AbstractVector, amplitude = 10.0) lon_device = on_architecture(arch, NF.(lon)) - lat_device = on_architecture(arch, NF.(lat)) + lat_device = on_architecture(arch, NF.(latd)) function periodic_bc(x::NF, t::NF) where {NF} lonₓ = lon_device[round(Int, x)] latₓ = lat_device[round(Int, x)] - T₀ = mean_annual_temperature(latₓ) + T₀ = climatology(latₓ) seconds_per_day = NF(24 * 3600) T = T₀ + NF(amplitude) * sin(2π * t / seconds_per_day - lonₓ) return T @@ -115,12 +112,11 @@ function get_temperature_bc(lon::AbstractVector, lat::AbstractVector, amplitude return periodic_bc end -bc = PrescribedSurfaceTemperature(:T_ub, get_temperature_bc(lon_masked, lat_masked)) -inits = (temperature = initial_soil_temperature,) +bc = PrescribedSurfaceTemperature(:T_ub, get_temperature_bc(lon_masked, latd_masked)) # Initialize the model; the input sources are matched to the namespaced input variables of # the prescribed soil horizons. -integrator = initialize(model, ForwardEuler(NF); inputs = soilgrids_inputs, boundary_conditions = bc, initializers = inits) +integrator = initialize(model, ForwardEuler(NF); inputs = soilgrids_inputs, boundary_conditions = bc) # We can verify that the SoilGrids texture has been correctly assigned to the first horizon: sand1 = RingGrids.Field(arch, interior(integrator.state.namespaces.horizon1.sand_fraction), grid) diff --git a/examples/simulations/speedy_wet_land.jl b/examples/simulations/speedy_wet_land.jl index 7da920cf9..5e35203db 100644 --- a/examples/simulations/speedy_wet_land.jl +++ b/examples/simulations/speedy_wet_land.jl @@ -28,22 +28,12 @@ plot_land_field(field, z_idx = 1; kwargs...) = heatmap(RingGrids.Field(CPU(), in plot_speedy_field(field; kwargs...) = heatmap(on_architecture(CPU(), field); kwargs...) -# To make the simulation a bit more interesting, we will use spatially periodic initial and boundary conditions. +# To make the simulation a bit more interesting, we will use a spatially varying initial condition. # The climatology will be determined by latitude with a maximum of 20 °C at the equator and minimum of -20°C at -# the poles. -mean_annual_temperature(lat) = 20 - abs(40 * sin(lat)) # maximum at equator - -function initial_soil_temperature(grid) - _, grid_lat = RingGrids.get_lonlats(grid.rings) # in radians - grid_lat_masked = grid_lat[on_architecture(CPU(), grid.mask)] - function init(x, z) - latᵢ = grid_lat_masked[round(Int, x)] - T₀ = mean_annual_temperature(latᵢ) - T = T₀ - 0.05 * z - return T - end - return init -end +# the poles, as provided by [`LatitudinalClimatology`](@ref), and the initial temperature profile is +# linear in depth via [`QuasiThermalSteadyState`](@ref). +climatology = LatitudinalClimatology(Float32) +energy_initializer = QuasiThermalSteadyState(Float32; T₀ = climatology, Qgeo = 0.05f0) function Terrarium.InputSources(dataset::SoilGrids2, grid::ColumnRingGrid, horizons = (Symbol(:horizon, i) for i in 1:6); name = nameof(typeof(dataset)), verbose = true) soilgrids_vars = (:sand_fraction, :silt_fraction, :clay_fraction, :bulk_density) @@ -89,7 +79,8 @@ strat = SoilGridsStratigraphy(eltype(land_grid); porosity) soil = SoilEnergyWaterCarbon(eltype(land_grid); strat) vegetation = PrescribedVegetation(eltype(land_grid)) surface_energy_balance = SurfaceEnergyBalance(eltype(land_grid); albedo = DiagnosticAlbedo(eltype(land_grid))) -terrarium_model = Terrarium.LandModel(land_grid; vegetation, soil, surface_energy_balance) +initializer = SoilInitializer(eltype(land_grid); energy = energy_initializer, hydrology = DefaultInitializer(eltype(land_grid))) +terrarium_model = Terrarium.LandModel(land_grid; vegetation, soil, surface_energy_balance, initializer) # Prepare the input sources: first the SoilGrids data, then the LAI data for prescribed vegetation. initial_date = DateTime(2024) @@ -98,9 +89,8 @@ lai_highveg_fts = FieldTimeSeries(cat(lai_highveg_fields..., dims = 2), land_gri lai_inputs = InputSource(lai_highveg_fts, name = :leaf_area_index, reftime = Speedy.DEFAULT_DATE; domain = Terrarium.Canopy()) inputs = InputSources(lai_inputs, soilgrids_inputs.sources...) # combine input sources -# Here we set our initial conditions for the soil +# Here we set the remaining initial conditions for the soil initializers = ( - temperature = initial_soil_temperature(land_grid), saturation_water_ice = 1.0, # fully saturated soil everywhere ) From af9d2afab2e7a378a7efe9be782210e8dba12d8a Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:31:48 +0200 Subject: [PATCH 10/19] Update docs --- docs/src/extending/core_interfaces.md | 2 +- docs/src/models/soil_model.md | 37 +++++++++++++++++++++++++++ docs/src/running/initialization.md | 17 ++++++++++++ 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/docs/src/extending/core_interfaces.md b/docs/src/extending/core_interfaces.md index f6db332e1..112758f7b 100644 --- a/docs/src/extending/core_interfaces.md +++ b/docs/src/extending/core_interfaces.md @@ -143,7 +143,7 @@ Both models and processes may optionally implement [`compute_boundary_conditions ### The `AbstractInitializer` interface -Standardized model initialization routines can be defined using the [`AbstractInitializer`](@ref) interface (see also [Initialization](@ref)). Each implementation of `AbstractModel` must allow for a user-defined `initializer` (the type can be constrained where appropriate). The simplest initializer is `DefaultInitializer`, which is a no-op that leaves all `Field`s at their default (zero) values. More complex models define their own composite initializers; for example, `SoilInitializer` composes separate initializers for the energy, hydrology, and biogeochemistry state variables. See [Soil models](@ref) for the full list of available initializer types. +Standardized model initialization routines can be defined using the [`AbstractInitializer`](@ref) interface (see also [Initialization](@ref)). Each implementation of `AbstractModel` must allow for a user-defined `initializer` (the type can be constrained where appropriate). The simplest initializer is `DefaultInitializer`, which is a no-op that leaves all `Field`s at their default (zero) values. More complex models define their own composite initializers; for example, `SoilInitializer` composes separate initializers for the energy, hydrology, and biogeochemistry state variables. See [Soil models](@ref) for the full list of available initializer types. Initializers may declare `input` variables by implementing [`variables`](@ref); these are included in the variables of the model. ```@docs; canonical = false AbstractInitializer diff --git a/docs/src/models/soil_model.md b/docs/src/models/soil_model.md index 9eac97ecb..c542cb8f5 100644 --- a/docs/src/models/soil_model.md +++ b/docs/src/models/soil_model.md @@ -75,6 +75,30 @@ initializer = SoilInitializer(Float32; model = SoilModel(grid; initializer) ``` +### Spatially varying initial values + +The parameters of the energy and hydrology initializers are declared as `input` variables of the model, with the values given to the initializer as their defaults: + +| Initializer | Input variables | +|:--|:--| +| [`ConstantSoilTemperature`](@ref) | `initial_surface_temperature` (°C) | +| [`QuasiThermalSteadyState`](@ref) | `initial_surface_temperature` (°C), `geothermal_heat_flux` (W/m²) | +| [`ConstantSaturation`](@ref) | `initial_saturation` | +| [`SaturationWaterTable`](@ref) | `vadose_zone_saturation`, `water_table_depth` (m) | + +Each value may therefore be a number, a function of the horizontal node coordinates, an array, a `Field`, or an [`AbstractFieldInitializer`](@ref). +Plain values are applied once and can be overridden by an [`InputSource`](@ref) with the same name and units. +Field initializers such as [`LatitudinalClimatology`](@ref) are re-evaluated at every initialization, so their parameters are exposed as model parameters. +Note that `geothermal_heat_flux` is the same variable read by the [`GeothermalHeatFlux`](@ref) bottom boundary condition, so the initial profile and the boundary condition stay consistent. + +```@example soilmodel +column_grid = ColumnGrid(arch, Float32, ExponentialSpacing(N = 10), 3) # three columns +energy = QuasiThermalSteadyState(Float32; T₀ = x -> 2.0f0 * x, Qgeo = 0.05f0) +model = SoilModel(column_grid; initializer = SoilInitializer(Float32; energy)) +integrator = initialize(model) +interior(integrator.state.initial_surface_temperature) +``` + ### Energy initializers ```@docs; canonical = false @@ -95,6 +119,19 @@ PiecewiseLinearInitialSoilTemperature SaturationWaterTable ``` +```@docs; canonical = false +ConstantSaturation +``` + +### Kernel functions + +```@docs; canonical = false +compute_quasi_steady_state_temperature +compute_constant_temperature +compute_water_table_saturation +compute_constant_saturation +``` + ### Fallback ```@docs; canonical = false diff --git a/docs/src/running/initialization.md b/docs/src/running/initialization.md index ea545cb77..d86ce9c3a 100644 --- a/docs/src/running/initialization.md +++ b/docs/src/running/initialization.md @@ -33,10 +33,27 @@ As a general rule, these initializers are invoked in the order that they are lis The keyword argument must be a `NamedTuple` where the keys correspond to the name of the state variable and the values are either scalars, arrays matching the size of the model `grid`, or functions of the form `f(coords...)` where `coords` are the non-[`Flat`](@extref Oceananigans.Grids.Flat) dimensions of `grid`. For column-based grids, this is generally `f(x,z)` with `x` corresponding to a column index. +## Field initializers + +Reusable, parameterized initializers of a single `Field` subtype [`AbstractFieldInitializer`](@ref). +They are accepted everywhere a `set!`-compatible value is: in the `initializers` keyword argument, as the default of an `input` variable, or as a parameter of a model initializer. +In the latter case, their `@param` fields become parameters of the model and they are re-evaluated at every initialization. + +```@docs; canonical = false +AbstractFieldInitializer +LatitudinalClimatology +``` + +Field initializers that depend on geographic position can use `λnodes` and `φnodes`, which return longitudes and latitudes (degrees) on both an Oceananigans `LatitudeLongitudeGrid` and a [`ColumnRingGrid`](@ref). +On the latter, they are indexed by active column. + ## Model initializers These `Initializer` types can be supplied to subtypes of [`AbstractModel`](@ref) during construction. Models can/should typically define corresponding [`AbstractInitializer`](@ref) types that represent common initialization strategies appropriate for the processes included in that model; e.g. the [`SoilModel`](@ref) defines [`SoilInitializer`](@ref) with process-specific initialization types like [`QuasiThermalSteadyState`](@ref) and [`SaturationWaterTable`](@ref). +Model initializers may also implement [`variables`](@ref) to declare `input` variables for their parameters. +These are collected by `variables(model)` alongside the process variables, which allows initial values to vary in space and to be supplied by [`InputSource`](@ref)s. + ## Built-in initialization routines Model/process specific dispatches of [`initialize!`](@ref) can be defined for process initialization logic that should be run regardless of the choice of prognostic state variable initialization. From e181f086e3ebc3ff1bd82549bf042b503db207f0 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 17:31:59 +0200 Subject: [PATCH 11/19] Update implementation plan for initializers --- ...LAN_spatially_varying_soil_initializers.md | 68 ++++++++++++------- docs/src/extending/state_variables.md | 5 +- 2 files changed, 47 insertions(+), 26 deletions(-) diff --git a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md index 563ae86e4..bb857a590 100644 --- a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md +++ b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md @@ -1,6 +1,6 @@ # Spatially varying soil initializers via input variables -> Status: **planned**. Draft awaiting human review; nothing implemented yet. +> Status: **completed**. Implemented on `bg/improved-init`; the full test suite passes locally, and the Enzyme and Reactant suites are left to CI. Date of initial draft: 2026-10-08 @@ -8,17 +8,31 @@ Base revision: `ce9ae606ae39063f52bb0ef23a3328a8b1904055` (tip of `main`, branch ## Originating prompt -> We need to generalize the existing soil initializer to allow for spatially varying inputs. The idea -> would be to declare corresponding input variables for T0 and Qgeo and then pass the struct values as -> default value initializers. These initializers could also be functions or callable structs that contain -> additional parameters, e.g. the periodic initialization used for soil temp and currently duplicated -> across many examples. Draft an implementation plan and summarize your findings. Ask questions if needed. +> We need to generalize the existing soil initializer to allow for spatially varying inputs. The idea would be to declare corresponding input variables for T0 and Qgeo and then pass the struct values as default value initializers. These initializers could also be functions or callable structs that contain additional parameters, e.g. the periodic initialization used for soil temp and currently duplicated across many examples. Draft an implementation plan and summarize your findings. Ask questions if needed. ## Revision log - **Rev 1 (2026-10-08):** Initial draft. - **Rev 2 (2026-10-08):** Open questions resolved by the author: variable names `initial_surface_temperature` and `geothermal_heat_flux` confirmed; hydrology initializers (change 5) are in scope; `LatitudinalClimatology(T_equator, ΔT)` is the only concrete field initializer and no generic function wrapper is added, since a plain function can be assigned directly to the initializer fields; `AbstractFieldInitializer` lives in core `src/initializers.jl`. - Implementation not yet approved. +- **Rev 3 (2026-10-08):** Dropped the `set_default!` hook from change 2. + An `AbstractFieldInitializer` default is instead handled by a dispatch of `initialize(var::InputVariable, grid, ...)` on the default's type at `Field` construction. +- **Rev 4 (2026-10-08):** Settled the role of `AbstractFieldInitializer` after questioning whether it duplicates the `initializers` keyword of `initialize`. + It does not: as a field of the model's `initializer` component it can carry `@param` fields that become model parameters, which closures and precomputed `Field`s cannot. + To make those parameters live, the model initializer re-evaluates `AbstractFieldInitializer` values at every `initialize!`, while all other default forms are applied once at `Field` construction. + The construction-time dispatch from Rev 3 is dropped again. + *Approved by the author on 2026-10-08.* +- **Rev 5 (2026-10-08):** Implementation notes, recorded during development on `bg/improved-init`. + Deviations from Rev 4, each requested or agreed by the author during implementation: + 1. **No dedicated kernels.** The initializers fill their prognostic fields with `set!(field, kernel(compute_..., init), state)`, using the existing `kernel`/`KernelFunction` utility in `src/utils/kernel_utils.jl`, which gained a `set!` dispatch for `KernelFunction`. + The kernel functions read their inputs from `fields` (e.g. `fields.initial_surface_temperature[i, j, 1]`) following the usual `compute_something(i, j, k, grid, fields, process)` convention. + 2. **Input variables live on `Ground(XY())`, not `Ground(Top())`/`Ground(Bottom())`.** A windowed declaration of `geothermal_heat_flux` conflicts with the plain surface `Field` that an `InputSource` infers, so all initializer inputs are declared as 2D surface variables. + Functions passed as defaults therefore take the horizontal coordinates only, e.g. `x -> ...` on a `ColumnGrid`. + 3. **`LatitudinalClimatology` takes latitude in degrees and supports `LatitudeLongitudeGrid`.** Geographic coordinates come from a unified interface: Oceananigans already provides `λnodes`/`φnodes` for `LatitudeLongitudeGrid`, and `ColumnRingGrid` now extends them to return the longitudes/latitudes (degrees) of its active columns. + Both are re-exported. + 4. **Parameters of initializers.** `ParameterEditing.parameters` is defined for `AbstractInitializer` (mirroring `AbstractProcess`), `Field`s used as initial values contribute no parameters, and the initializers define `ConstructionBase.constructorof` so that `reconstruct` preserves `NF`. + 5. **Exports.** Abstract types (`AbstractInitializer`, `AbstractFieldInitializer`) are not exported; only `LatitudinalClimatology` is added to the exports. + 6. **Enzyme and Reactant tests are left to CI.** A new Reactant configuration `:soil_heat_global_initializer` is registered in `test/reactant/`, and the existing Enzyme soil energy test already builds its model with `SoilInitializer`; neither suite is run locally. + 7. **Hydrology inputs.** `ConstantSaturation` declares `initial_saturation`, and `SaturationWaterTable` declares `vadose_zone_saturation` and `water_table_depth`. ## Problem description @@ -87,12 +101,8 @@ Two further facts matter for the design: - Widen the `Def` type parameter of `InputVariable` (`src/abstract_variables.jl`) from `Union{Nothing, Number, Function}` to an unconstrained parameter. It remains concrete per instance. Document the accepted forms: `nothing`, a number, any argument accepted by `Oceananigans.set!` (functions of the node coordinates, arrays, `Field`s), or an `AbstractFieldInitializer` (below). -- Introduce a small hook in `initialize(var, grid, clock, fields, bcs)` so that the default is applied through a single overloadable method instead of a bare `set!`: - - ```julia - set_default!(field, grid, default) = set!(field, default) - set_default!(field, grid, ::Nothing) = nothing - ``` +- The generic default path (`set!(field, var.default)`) is unchanged, and no special case is needed at `Field` construction. + An `AbstractFieldInitializer` default is applied at construction through its own `set!` fallback (change 3) and then re-evaluated by the model initializer at every `initialize!` (change 4), so that its parameters take effect after `reconstruct`. ### 3. `AbstractFieldInitializer`: reusable, grid-aware, parameterized field initializers @@ -100,28 +110,30 @@ Add to `src/initializers.jl`: ```julia """ -Base type for reusable initializers of a single `Field`. Implementations are callable structs that -provide `initialize!(field, grid, init::AbstractFieldInitializer)`; the default evaluates `init` at -the node coordinates via `set!`. They can be passed as `default` of an `input` variable, as an entry of -the `initializers` NamedTuple of `initialize`, or as a parameter of a model initializer. +Base type for reusable initializers of a single `Field`. +Implementations provide `initialize!(field, grid, init::AbstractFieldInitializer)`; the default +evaluates `init` at the node coordinates via `set!`. +Field initializers held by a model initializer may declare `@param` fields, which are then exposed +as model parameters and re-evaluated at every `initialize!`. """ abstract type AbstractFieldInitializer end initialize!(field::AbstractField, grid::AbstractGrid, init::AbstractFieldInitializer) = set!(field, (coords...) -> init(coords...)) -set_default!(field, grid, init::AbstractFieldInitializer) = initialize!(field, grid, init) +Oceananigans.set!(field::AbstractField, init::AbstractFieldInitializer) = initialize!(field, field.grid, init) ``` -and make the user-facing `initialize!(state, inits::NamedTuple)` path accept them as well (it currently calls `set!` directly, which would not know what to do with a callable struct). +The `set!` extension makes a field initializer acceptable everywhere a `set!`-compatible value is, including input defaults and the `initializers` keyword, without further dispatch sites. Only one concrete implementation is in scope here, to replace the duplicated example code: ```julia """ Latitude-dependent surface temperature climatology on a `ColumnRingGrid`: -T(φ) = T_equator - |ΔT sin φ|. Evaluated per active land column. +T(φ) = T_equator - |ΔT sin φ|. +Evaluated per active land column. """ -@kwdef struct LatitudinalClimatology{NF} <: AbstractFieldInitializer - T_equator::NF = 20 - ΔT::NF = 40 +@parameterized @kwdef struct LatitudinalClimatology{NF} <: AbstractFieldInitializer + @param T_equator::NF = 20 + @param ΔT::NF = 40 end (init::LatitudinalClimatology)(lat) = init.T_equator - abs(init.ΔT * sin(lat)) function initialize!(field, grid::ColumnRingGrid, init::LatitudinalClimatology) @@ -129,6 +141,9 @@ function initialize!(field, grid::ColumnRingGrid, init::LatitudinalClimatology) end ``` +Why this is not redundant with the `initializers` keyword of `initialize`: that keyword, like input defaults, accepts any `set!`-compatible value, but such values are opaque to `ParameterEditing`. +A field initializer stored in the model's `initializer` component can carry `@param` fields, so `T_equator` and `ΔT` appear in `parameters(model)` and can be changed through `reconstruct` or differentiated. + No generic wrapper around a user function is provided: a plain function of the node coordinates can be assigned directly to `T₀`/`Qgeo` (or to any `input` default), and the implementation must keep that path working. The `periodic_bc` *boundary* condition in the same examples (diurnal cycle shifted by longitude) is a boundary-condition concern and is left out of this plan; see "Future work". @@ -154,6 +169,9 @@ variables(init::QuasiThermalSteadyState) = ( ``` `initialize!(state, model, init::QuasiThermalSteadyState)` then reads the two input fields and fills `temperature` with a kernel launched via `launch!`, `T[i, j, k] = T₀[i, j] - Qgeo[i, j] / k_eff * z_k`, instead of a host-side closure. +Before launching the kernel, `initialize!` re-evaluates any `T₀`/`Qgeo` that is an `AbstractFieldInitializer` into its input field via `initialize!(field, grid, init)`, so that parameter changes made through `reconstruct` take effect. +All other forms (numbers, functions, arrays, `Field`s) are applied only once, as construction-time defaults, and remain overridable by an `InputSource` of the same name. +Supplying both a parameterized field initializer and an `InputSource` for the same variable is therefore not meaningful; the field initializer wins, and this is documented. This is GPU- and Reactant-native (no `set_to_function!` detour) and keeps `k_eff` as a scalar parameter. `ConstantSoilTemperature` gets the same treatment with only the `initial_surface_temperature` input. @@ -184,7 +202,8 @@ Export `AbstractFieldInitializer` and `LatitudinalClimatology` from `src/models/ - verify that `QuasiThermalSteadyState` with scalar `T₀`/`Qgeo` reproduces the current profile bit for bit (regression), - verify a column-varying `T₀` given as a function and as a `Field`, - verify that an `InputSource` named `geothermal_heat_flux` overrides the default and that the profile uses the overridden values, and - - verify that `initialize!(integrator)` (re-initialization) leaves the result unchanged. + - verify that `initialize!(integrator)` (re-initialization) leaves the result unchanged, and + - verify that `initialize(integrator, params)` with a changed `T_equator` of a `LatitudinalClimatology` re-evaluates the initial profile. - `test/state_variables.jl`: `variables(model)` includes the initializer's inputs; `DefaultInitializer` contributes none; mismatched units in an `InputSource` raise the existing conflict error. - `test/grids.jl` or a new `test/initializers.jl`: `LatitudinalClimatology` on a masked `ColumnRingGrid` assigns each active column the value at its latitude. - `test/differentiability/soil_energy_diff.jl`: run the Enzyme energy test with the new initializer to confirm the kernel-based `initialize!` is differentiable. @@ -205,6 +224,7 @@ Export `AbstractFieldInitializer` and `LatitudinalClimatology` from `src/models/ - `AbstractFieldInitializer` defaults are applied at `Field` construction (as all input defaults are today), i.e. before any `InputSource` runs, and only once. That is the desired order, but it means a default that depends on other inputs is not supported. - The `Variables` conflict rule requires matching units between an `InputSource` and the declared input; users must pass `units = u"W/m^2"` (resp. `u"°C"`) to the source. +- An `AbstractFieldInitializer` held by the model initializer overrides an `InputSource` of the same name at every `initialize!` (by design, so its parameters are live). - `LandModel` with a `SoilInitializer` is not addressed here; whichever way `initialize!` dispatches for that combination today is unchanged. ## Future work diff --git a/docs/src/extending/state_variables.md b/docs/src/extending/state_variables.md index dfb217888..b9bc61604 100644 --- a/docs/src/extending/state_variables.md +++ b/docs/src/extending/state_variables.md @@ -23,10 +23,11 @@ Variables are typically constructed using one of the convenience functions [`pro These variable definitions are purely symbolic; they do not hold any data and cannot be used for computation. Constructing [`StateVariables`](@ref) from a model, process, or [`Variables`](@ref) container (see following sections) results in corresponding [Fields](@ref) being allocated for each variable. -A default implementation of `variables` is provided for all [`AbstractModel`](@ref) and `AbstractCoupledProcesses` types that automatically collects variables from all [`AbstractProcess`](@ref) types defined therein: +A default implementation of `variables` is provided for all [`AbstractModel`](@ref) and `AbstractCoupledProcesses` types that automatically collects variables from all [`AbstractProcess`](@ref) types defined therein. For models, any variables declared by the model's initializer are included as well: ```@docs; canonical = false -variables(obj::Union{AbstractCoupledProcesses, AbstractModel}) +variables(obj::AbstractCoupledProcesses) +variables(model::AbstractModel) ``` Most state variables will thus be defined by implementation of `AbstractProcess`. As an example, suppose we are implementing a new process `MyProcess` and we want to define the necessary state variables. We do this by defining a new dispatch of the `variables` method: From 6945f632c8a7eb47bcb7c06d1ad28c0a1182ee47 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 18:49:55 +0200 Subject: [PATCH 12/19] Update Speedy example to load soil temperature climatology --- examples/simulations/speedy_wet_land.jl | 39 ++++++++++++++++++++----- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/examples/simulations/speedy_wet_land.jl b/examples/simulations/speedy_wet_land.jl index 5e35203db..0f0180cb4 100644 --- a/examples/simulations/speedy_wet_land.jl +++ b/examples/simulations/speedy_wet_land.jl @@ -28,13 +28,6 @@ plot_land_field(field, z_idx = 1; kwargs...) = heatmap(RingGrids.Field(CPU(), in plot_speedy_field(field; kwargs...) = heatmap(on_architecture(CPU(), field); kwargs...) -# To make the simulation a bit more interesting, we will use a spatially varying initial condition. -# The climatology will be determined by latitude with a maximum of 20 °C at the equator and minimum of -20°C at -# the poles, as provided by [`LatitudinalClimatology`](@ref), and the initial temperature profile is -# linear in depth via [`QuasiThermalSteadyState`](@ref). -climatology = LatitudinalClimatology(Float32) -energy_initializer = QuasiThermalSteadyState(Float32; T₀ = climatology, Qgeo = 0.05f0) - function Terrarium.InputSources(dataset::SoilGrids2, grid::ColumnRingGrid, horizons = (Symbol(:horizon, i) for i in 1:6); name = nameof(typeof(dataset)), verbose = true) soilgrids_vars = (:sand_fraction, :silt_fraction, :clay_fraction, :bulk_density) metadataset = MetadataSet(soilgrids_vars...; dataset) @@ -79,6 +72,36 @@ strat = SoilGridsStratigraphy(eltype(land_grid); porosity) soil = SoilEnergyWaterCarbon(eltype(land_grid); strat) vegetation = PrescribedVegetation(eltype(land_grid)) surface_energy_balance = SurfaceEnergyBalance(eltype(land_grid); albedo = DiagnosticAlbedo(eltype(land_grid))) + +# The initial soil temperature profiles are linear in depth via [`QuasiThermalSteadyState`](@ref), with the +# surface temperature taken from SpeedyWeather's monthly land surface temperature climatology. We load it +# directly from the SpeedyWeather assets, regrid it onto our ring grid, and average over the twelve months +# to obtain the mean annual surface temperature in °C. +lst_climatology = RingGrids.get_asset( + "data/boundary_conditions/land_surface_temperature.nc"; + from_assets = true, + name = "lst", + ArrayType = RingGrids.FullGaussianField, + FileFormat = NCDataset, + output_grid = on_architecture(CPU(), ring_grid), +) +mean_lst = RingGrids.Field(on_architecture(CPU(), ring_grid)) +mean_lst.data .= vec(mean(lst_climatology.data, dims = 2)) .- 273.15f0 # K to °C + +# The climatology is undefined over the ocean, so some coastal land columns of our grid may not have a value +# after regridding. For these, we fall back to a simple latitude-dependent climatology given by +# [`LatitudinalClimatology`](@ref). +_, ring_lat = RingGrids.get_lonlats(on_architecture(CPU(), ring_grid)) # in radians +fallback = LatitudinalClimatology(Float32) +missing_lst = .!isfinite.(mean_lst.data) +mean_lst.data[missing_lst] .= fallback.(rad2deg.(ring_lat[missing_lst])) +plot_speedy_field(mean_lst) + +# The initializer declares its surface temperature as the input variable `initial_surface_temperature`, +# so we supply the regridded climatology as an [`InputSource`](@ref) with that name below. This keeps the +# data with the inputs rather than in the (grid-agnostic) model. +T₀_source = InputSource(land_grid, mean_lst; name = :initial_surface_temperature, domain = Terrarium.Ground(), units = u"°C") +energy_initializer = QuasiThermalSteadyState(Float32; Qgeo = 0.05f0) initializer = SoilInitializer(eltype(land_grid); energy = energy_initializer, hydrology = DefaultInitializer(eltype(land_grid))) terrarium_model = Terrarium.LandModel(land_grid; vegetation, soil, surface_energy_balance, initializer) @@ -87,7 +110,7 @@ initial_date = DateTime(2024) soilgrids_inputs = InputSources(SoilGrids2(), land_grid) lai_highveg_fts = FieldTimeSeries(cat(lai_highveg_fields..., dims = 2), land_grid, 0.0:1day:365day) lai_inputs = InputSource(lai_highveg_fts, name = :leaf_area_index, reftime = Speedy.DEFAULT_DATE; domain = Terrarium.Canopy()) -inputs = InputSources(lai_inputs, soilgrids_inputs.sources...) # combine input sources +inputs = InputSources(T₀_source, lai_inputs, soilgrids_inputs.sources...) # combine input sources # Here we set the remaining initial conditions for the soil initializers = ( From f803d7e1acbbd6dba3e2879cc03001a01f65619c Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 18:52:06 +0200 Subject: [PATCH 13/19] Update ERA5 example, soil model doc page, and impl plan --- ...LAN_spatially_varying_soil_initializers.md | 4 ++++ docs/src/models/soil_model.md | 1 + examples/simulations/land_global_era5.jl | 23 ++++++++++--------- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md index bb857a590..fc24fc421 100644 --- a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md +++ b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md @@ -33,6 +33,10 @@ Base revision: `ce9ae606ae39063f52bb0ef23a3328a8b1904055` (tip of `main`, branch 5. **Exports.** Abstract types (`AbstractInitializer`, `AbstractFieldInitializer`) are not exported; only `LatitudinalClimatology` is added to the exports. 6. **Enzyme and Reactant tests are left to CI.** A new Reactant configuration `:soil_heat_global_initializer` is registered in `test/reactant/`, and the existing Enzyme soil energy test already builds its model with `SoilInitializer`; neither suite is run locally. 7. **Hydrology inputs.** `ConstantSaturation` declares `initial_saturation`, and `SaturationWaterTable` declares `vadose_zone_saturation` and `water_table_depth`. +- **Rev 6 (2026-10-08):** Spatial data in the examples is supplied through `InputSource`s named `initial_surface_temperature` instead of as a `Field` stored in the initializer. + A `Field` held by the initializer ties the (otherwise grid-agnostic, stateless) model to one grid and architecture. + `Field`s remain accepted as initial values, and the soil model docs recommend input sources for data. + The SpeedyWeather example now initializes from SpeedyWeather's monthly land surface temperature climatology, averaged over the year, with `LatitudinalClimatology` as a fallback for coastal columns without data. ## Problem description diff --git a/docs/src/models/soil_model.md b/docs/src/models/soil_model.md index c542cb8f5..ee59d0d05 100644 --- a/docs/src/models/soil_model.md +++ b/docs/src/models/soil_model.md @@ -88,6 +88,7 @@ The parameters of the energy and hydrology initializers are declared as `input` Each value may therefore be a number, a function of the horizontal node coordinates, an array, a `Field`, or an [`AbstractFieldInitializer`](@ref). Plain values are applied once and can be overridden by an [`InputSource`](@ref) with the same name and units. +Spatial data, such as a regridded climatology, should preferably be supplied this way rather than as a `Field` stored in the initializer: the initializer is part of the model, which should stay independent of any particular grid and free of state. Field initializers such as [`LatitudinalClimatology`](@ref) are re-evaluated at every initialization, so their parameters are exposed as model parameters. Note that `geothermal_heat_flux` is the same variable read by the [`GeothermalHeatFlux`](@ref) bottom boundary condition, so the initial profile and the boundary condition stay consistent. diff --git a/examples/simulations/land_global_era5.jl b/examples/simulations/land_global_era5.jl index cee7225cf..bf2ab6057 100644 --- a/examples/simulations/land_global_era5.jl +++ b/examples/simulations/land_global_era5.jl @@ -136,7 +136,11 @@ soil = SoilEnergyWaterCarbon(NF) vegetation = PrescribedVegetation(NF) snow = SingleLayerSnow(NF) atmosphere = PrescribedAtmosphere(NF; wind = WindVelocity()) -land = LandModel(grid; soil, vegetation, snow, atmosphere) +# The soil temperature is initialized to a quasi-steady state with a geothermal gradient of 0.05 K/m +# (see [`QuasiThermalSteadyState`](@ref)); its surface temperature is supplied as an input below. +energy_initializer = QuasiThermalSteadyState(NF; Qgeo = NF(0.05)) +initializer = SoilInitializer(NF; energy = energy_initializer, hydrology = DefaultInitializer(NF)) +land = LandModel(grid; soil, vegetation, snow, atmosphere, initializer) @show variables(land) # ## Assembling the input sources @@ -172,21 +176,18 @@ inputs = InputSources( # and the surface starts snow-free so the winter forcing builds up the snowpack. The soil temperature # is initialized from the mean annual air temperature computed directly from the ERA5-Land `t2m` # forcing data: we average over all time records, convert to °C, and regridded onto the model's land -# columns. A mild geothermal gradient (0.05 K/m) is added with depth. Computing the mean from the -# actual forcing data (rather than a latitude-based climatology) ensures the initialization is -# consistent with the prescribed atmospheric state. +# columns. Computing the mean from the actual forcing data (rather than a latitude-based climatology) +# ensures the initialization is consistent with the prescribed atmospheric state. t2m_mean = mean(t2m; dims = Ti) # average over time axis t2m_mean_C = Terrarium.kelvin_to_celsius.(constants, t2m_mean) t2m_mean_field = RingGrids.Field(Array(reshape(t2m_mean_C, :)), Terrarium.native_grid(ERA5LandForcings())) t2m_mean_on_grid = RingGrids.interpolate(on_architecture(CPU(), grid.rings), t2m_mean_field) -T₀_field = Field(on_architecture(arch, t2m_mean_on_grid), grid) # surface field over the land columns -# The regridded mean air temperature `Field` is passed directly as the surface temperature of -# [`QuasiThermalSteadyState`](@ref); the geothermal gradient follows from `Qgeo / k_eff`. Since the -# initializer is part of the model, we rebuild the `LandModel` with it. -energy_initializer = QuasiThermalSteadyState(NF; T₀ = T₀_field, Qgeo = NF(0.05)) -initializer = SoilInitializer(NF; energy = energy_initializer, hydrology = DefaultInitializer(NF)) -land = LandModel(grid; soil, vegetation, snow, atmosphere, initializer) +# The initializer declares its surface temperature as the input variable `initial_surface_temperature`, +# so we supply the regridded mean air temperature as an [`InputSource`](@ref) with that name. This keeps +# the data with the inputs rather than in the (grid-agnostic) model. +T₀_source = InputSource(grid, t2m_mean_on_grid; name = :initial_surface_temperature, domain = Terrarium.Ground(), units = u"°C") +inputs = InputSources(T₀_source, inputs.sources...) initializers = ( saturation_water_ice = (x, z) -> min(one(NF), NF(0.6) - NF(0.05) * z), From 0d2b14d73635edc3bd2deb96eba0c218c85734f8 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 18:59:39 +0200 Subject: [PATCH 14/19] Use parameterized macro directly for initializers --- src/abstract_model.jl | 4 ++++ src/initializers.jl | 13 ------------- src/models/soil/soil_model_init.jl | 30 +++++++++++++++--------------- 3 files changed, 19 insertions(+), 28 deletions(-) diff --git a/src/abstract_model.jl b/src/abstract_model.jl index 908c46e92..37fee7afb 100644 --- a/src/abstract_model.jl +++ b/src/abstract_model.jl @@ -263,6 +263,10 @@ function ParameterEditing.parameters(proc::AbstractProcess; kwargs...) return ParameterEditing.ParameterTable((; nonempty_params...)) end +# `Field`s used as initial values (e.g. in a `@param` field of an initializer) are state, not parameters; +# without this, ParameterEditing would treat them as array-valued parameters since `AbstractField <: AbstractArray`. +ParameterEditing.parameters(::Type{PT}, ::AbstractField; kwargs...) where {PT <: ParameterEditing.AbstractParam} = (;) + function Base.show(io::IO, model::AbstractModel{NF}) where {NF} println(io, "$(nameof(typeof(model))){$NF} on $(architecture(get_grid(model)))") for name in propertynames(model) diff --git a/src/initializers.jl b/src/initializers.jl index fe92ffcab..2e07ab575 100644 --- a/src/initializers.jl +++ b/src/initializers.jl @@ -103,16 +103,3 @@ end initialize!(field::AbstractField, ::LatitudeLongitudeGrid, init::LatitudinalClimatology) = set!(field, (λ, φ, args...) -> init(φ)) initialize!(field::AbstractField, grid::AbstractLandGrid, init::LatitudinalClimatology) = initialize!(field, ground_domain(grid), init) - -# Parameters of initializers are collected from their fields, like for processes. Numbers become -# parameters, parameterized field initializers contribute theirs, and functions contribute none. -function ParameterEditing.parameters(::Type{PT}, init::AbstractInitializer; kwargs...) where {PT <: ParameterEditing.AbstractParam} - init_params = map(fieldnames(typeof(init))) do name - name => ParameterEditing.parameters(PT, getproperty(init, name); kwargs...) - end - nonempty_params = filter(p -> length(p[2]) > 0, init_params) - return ParameterEditing.ParameterTable((; nonempty_params...)) -end - -# `Field`s used as initial values are state, not parameters. -ParameterEditing.parameters(::Type{PT}, ::AbstractField; kwargs...) where {PT <: ParameterEditing.AbstractParam} = (;) diff --git a/src/models/soil/soil_model_init.jl b/src/models/soil/soil_model_init.jl index 51045c01b..00f29b599 100644 --- a/src/models/soil/soil_model_init.jl +++ b/src/models/soil/soil_model_init.jl @@ -3,20 +3,20 @@ Initializer for coupled soil energy/hydrology/biogeochemistry models. """ -struct SoilInitializer{ +@parameterized struct SoilInitializer{ NF, EnergyInit <: AbstractInitializer{NF}, HydrologyInit <: AbstractInitializer{NF}, BGCInit <: AbstractInitializer{NF}, } <: AbstractInitializer{NF} "Soil energy/temperature state initializer" - energy::EnergyInit + @component energy::EnergyInit "Soil hydrology state initializer" - hydrology::HydrologyInit + @component hydrology::HydrologyInit "Soil biogeochemistry state initializer" - biogeochem::BGCInit + @component biogeochem::BGCInit end function SoilInitializer( @@ -53,9 +53,9 @@ a field initializer is re-evaluated at every `initialize!` so that its parameter Properties: $TYPEDFIELDS """ -struct ConstantSoilTemperature{NF, T0} <: AbstractInitializer{NF} +@parameterized struct ConstantSoilTemperature{NF, T0} <: AbstractInitializer{NF} "Initial surface temperature (°C)" - T₀::T0 + @param T₀::T0 end """ @@ -109,15 +109,15 @@ parameters take effect. Properties: $TYPEDFIELDS """ -struct QuasiThermalSteadyState{NF, T0, QG} <: AbstractInitializer{NF} +@parameterized struct QuasiThermalSteadyState{NF, T0, QG} <: AbstractInitializer{NF} "Initial surface temperature (°C)" - T₀::T0 + @param T₀::T0 "Geothermal heat flux (W/m²)" - Qgeo::QG + @param Qgeo::QG "Bulk thermal conductivity (W/m/K)" - k_eff::NF + @param k_eff::NF bounds = Positive end """ @@ -201,9 +201,9 @@ for the accepted forms of the default. Properties: $TYPEDFIELDS """ -struct ConstantSaturation{NF, S} <: AbstractInitializer{NF} +@parameterized struct ConstantSaturation{NF, S} <: AbstractInitializer{NF} "Initial water/ice saturation (-)" - sat::S + @param sat::S bounds = UnitInterval end ConstantSaturation(::Type{NF}; sat = one(NF)) where {NF} = ConstantSaturation{NF, typeof(sat)}(sat) @@ -239,12 +239,12 @@ accepted forms of the defaults. Properties: $TYPEDFIELDS """ -struct SaturationWaterTable{NF, S, D} <: AbstractInitializer{NF} +@parameterized struct SaturationWaterTable{NF, S, D} <: AbstractInitializer{NF} "Saturation in the vadose zone above the water table (-)" - vadose_zone_saturation::S + @param vadose_zone_saturation::S bounds = UnitInterval "Depth of the water table below the surface (m)" - water_table_depth::D + @param water_table_depth::D bounds = Nonnegative end function SaturationWaterTable(::Type{NF}; vadose_zone_saturation = convert(NF, 3 // 4), water_table_depth = NF(5)) where {NF} From 3d349bbfef96229ef67c09131a715eb61257e468 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Sun, 27 Sep 2026 13:02:29 +0200 Subject: [PATCH 15/19] Fix bug in ParameterEditing dispatch for process types --- src/abstract_model.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/abstract_model.jl b/src/abstract_model.jl index 37fee7afb..d40c084aa 100644 --- a/src/abstract_model.jl +++ b/src/abstract_model.jl @@ -255,9 +255,9 @@ Convenience constructor for all `AbstractModel` types that accepts `grid` as a p (::Type{Model})(grid::AbstractGrid; kwargs...) where {Model <: AbstractModel} = Model(; grid, kwargs...) # Default parameters collection for processes -function ParameterEditing.parameters(proc::AbstractProcess; kwargs...) +function ParameterEditing.parameters(::Type{PT}, proc::AbstractProcess; kwargs...) where {PT <: ModelParameters.AbstractParam} proc_params = map(fieldnames(typeof(proc))) do name - name => ParameterEditing.parameters(getproperty(proc, name)) + name => ParameterEditing.parameters(PT, getproperty(proc, name)) end nonempty_params = filter(p -> length(p[2]) > 0, proc_params) return ParameterEditing.ParameterTable((; nonempty_params...)) From e477b35d0e510e039508b8345f1ebe53b917228b Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 19:12:06 +0200 Subject: [PATCH 16/19] Include parameter handling tests and update plan --- .../2026-10-08_PLAN_spatially_varying_soil_initializers.md | 4 +++- test/parameter_handling.jl | 7 +++++-- test/runtests.jl | 4 ++++ 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md index fc24fc421..b49fb601b 100644 --- a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md +++ b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md @@ -29,7 +29,9 @@ Base revision: `ce9ae606ae39063f52bb0ef23a3328a8b1904055` (tip of `main`, branch Functions passed as defaults therefore take the horizontal coordinates only, e.g. `x -> ...` on a `ColumnGrid`. 3. **`LatitudinalClimatology` takes latitude in degrees and supports `LatitudeLongitudeGrid`.** Geographic coordinates come from a unified interface: Oceananigans already provides `λnodes`/`φnodes` for `LatitudeLongitudeGrid`, and `ColumnRingGrid` now extends them to return the longitudes/latitudes (degrees) of its active columns. Both are re-exported. - 4. **Parameters of initializers.** `ParameterEditing.parameters` is defined for `AbstractInitializer` (mirroring `AbstractProcess`), `Field`s used as initial values contribute no parameters, and the initializers define `ConstructionBase.constructorof` so that `reconstruct` preserves `NF`. + 4. **Parameters of initializers.** The concrete soil initializers are declared with `@parameterized`: `SoilInitializer` marks its sub-initializers as `@component`, and the energy and hydrology initializers mark their initial values as `@param` with bounds where applicable. + `Field`s used as initial values contribute no parameters, and the initializers define `ConstructionBase.constructorof` so that `reconstruct` preserves `NF`. + The typed-dispatch fix for process parameter collection (`117bf3ff4`, from `bg/differentiability-examples`) is cherry-picked, since without it no process parameters were collected through a model `@component`. 5. **Exports.** Abstract types (`AbstractInitializer`, `AbstractFieldInitializer`) are not exported; only `LatitudinalClimatology` is added to the exports. 6. **Enzyme and Reactant tests are left to CI.** A new Reactant configuration `:soil_heat_global_initializer` is registered in `test/reactant/`, and the existing Enzyme soil energy test already builds its model with `SoilInitializer`; neither suite is run locally. 7. **Hydrology inputs.** `ConstantSaturation` declares `initial_saturation`, and `SaturationWaterTable` declares `vadose_zone_saturation` and `water_table_depth`. diff --git a/test/parameter_handling.jl b/test/parameter_handling.jl index 1837f06c8..ece266577 100644 --- a/test/parameter_handling.jl +++ b/test/parameter_handling.jl @@ -8,7 +8,10 @@ using Test soil = SoilEnergyWaterCarbon(eltype(grid)) soil_params = vec(parameters(soil)) @test length(soil_params) > 0 - @test all(map(Base.Fix1(haskey, soil_params), (:biogeochem, :energy, :hydrology, :strat))) + @test all(map(Base.Fix1(haskey, soil_params), (:biogeochem, :energy, :hydrology))) + # The stratigraphy stores its horizons in a `Tuple`, which parameter collection does not descend into, + # so the parameters of the horizons (e.g. their porosity) are currently not collected. + @test_broken haskey(soil_params, :strat) # check that collected parameters for thermal conductivities match the values in the struct @test all(map(==, soil_params.energy.thermal_properties.conductivities, getproperties(soil.energy.thermal_properties.conductivities))) model = SoilModel(grid; soil) @@ -20,5 +23,5 @@ using Test # test reinitialization of integrator integrator = initialize(model) updated_integrator = initialize(integrator, model_params) - @test all(vec(parameters(integrator.model)) .≈ model_params) + @test all(vec(parameters(updated_integrator.model)) .≈ model_params) end diff --git a/test/runtests.jl b/test/runtests.jl index 4ca67f0ca..f3b77210a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -32,6 +32,10 @@ if MAIN_TESTS include("initializers.jl") end + @testset "Parameter handling" begin + include("parameter_handling.jl") + end + @testset "Boundary conditions" begin include("boundary_conditions.jl") end From 741598182e73dd9cba30680d223589a87048e32b Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 19:18:51 +0200 Subject: [PATCH 17/19] Bump Speedy compat version in examples project --- examples/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/Project.toml b/examples/Project.toml index 5c40ed429..3bb42bd27 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -27,7 +27,7 @@ Terrarium = "80418d68-07fa-499d-ae2b-c12e531f5cd8" Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" [compat] -SpeedyWeather = "0.21, 0.22" +SpeedyWeather = "0.23" [sources.NumericalEarth] rev = "main" From 835bd7c1bde9262e4c7a18cc3580aa6425da8f50 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 8 Oct 2026 19:29:37 +0200 Subject: [PATCH 18/19] Reset inputs in initialize! and remove reinitialize! method --- ...LAN_spatially_varying_soil_initializers.md | 8 ++++++- docs/src/models/soil_model.md | 4 ++-- docs/src/running/initialization.md | 2 +- src/initializers.jl | 12 +--------- src/models/soil/soil_model_init.jl | 20 +++++----------- src/state_variables.jl | 23 +++++++++++++++++++ src/timesteppers/model_integrator.jl | 2 ++ test/initializers.jl | 21 +++++++++++++++++ 8 files changed, 63 insertions(+), 29 deletions(-) diff --git a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md index b49fb601b..a81636d4f 100644 --- a/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md +++ b/docs/dev/2026-10/2026-10-08_PLAN_spatially_varying_soil_initializers.md @@ -39,6 +39,10 @@ Base revision: `ce9ae606ae39063f52bb0ef23a3328a8b1904055` (tip of `main`, branch A `Field` held by the initializer ties the (otherwise grid-agnostic, stateless) model to one grid and architecture. `Field`s remain accepted as initial values, and the soil model docs recommend input sources for data. The SpeedyWeather example now initializes from SpeedyWeather's monthly land surface temperature climatology, averaged over the year, with `LatitudinalClimatology` as a fallback for coastal columns without data. +- **Rev 7 (2026-10-08):** Replaces the re-evaluation rule of Rev 4. + Instead of the model initializer re-evaluating field initializers (`reinitialize!`, now removed), `initialize!(::ModelIntegrator)` first resets every input field to the default declared by the current model's variables (`reset_inputs!`), then applies input sources, user initializers, and the model initializer. + Defaults of every form are therefore re-applied on each initialization, so all initializer parameters, including plain scalars, take effect through `initialize(integrator, params)`, and input sources always take precedence over defaults, including field initializers. + *Approved by the author on 2026-10-08.* ## Problem description @@ -230,11 +234,13 @@ Export `AbstractFieldInitializer` and `LatitudinalClimatology` from `src/models/ - `AbstractFieldInitializer` defaults are applied at `Field` construction (as all input defaults are today), i.e. before any `InputSource` runs, and only once. That is the desired order, but it means a default that depends on other inputs is not supported. - The `Variables` conflict rule requires matching units between an `InputSource` and the declared input; users must pass `units = u"W/m^2"` (resp. `u"°C"`) to the source. -- An `AbstractFieldInitializer` held by the model initializer overrides an `InputSource` of the same name at every `initialize!` (by design, so its parameters are live). +- Defaults are re-applied only by `initialize!(::ModelIntegrator)`; calling `initialize!(state, model)` directly on an existing state does not refresh input fields. - `LandModel` with a `SoilInitializer` is not addressed here; whichever way `initialize!` dispatches for that combination today is unchanged. ## Future work +- Reorder `initialize!(::ModelIntegrator)` to run user `initializers` after the model initializer but before the process `initialize!` dispatches, so that explicit per-field values are not overwritten by model initializers. + This requires splitting the model initializer step out of each model's `initialize!(state, model)`; deferred by the author on 2026-10-08. - A reusable periodic *boundary condition* (`T₀(lat) + A sin(2π t / P - lon)`) to remove the matching `get_temperature_bc` duplication in the examples. - Applying auxiliary-variable initializers in `reset!` (existing TODO in `state_variables.jl`). diff --git a/docs/src/models/soil_model.md b/docs/src/models/soil_model.md index ee59d0d05..2e6986beb 100644 --- a/docs/src/models/soil_model.md +++ b/docs/src/models/soil_model.md @@ -87,9 +87,9 @@ The parameters of the energy and hydrology initializers are declared as `input` | [`SaturationWaterTable`](@ref) | `vadose_zone_saturation`, `water_table_depth` (m) | Each value may therefore be a number, a function of the horizontal node coordinates, an array, a `Field`, or an [`AbstractFieldInitializer`](@ref). -Plain values are applied once and can be overridden by an [`InputSource`](@ref) with the same name and units. +These defaults are re-applied at every initialization, so changes to their parameters take effect, and an [`InputSource`](@ref) with the same name and units always takes precedence over them. Spatial data, such as a regridded climatology, should preferably be supplied this way rather than as a `Field` stored in the initializer: the initializer is part of the model, which should stay independent of any particular grid and free of state. -Field initializers such as [`LatitudinalClimatology`](@ref) are re-evaluated at every initialization, so their parameters are exposed as model parameters. +Field initializers such as [`LatitudinalClimatology`](@ref) expose their own parameters as model parameters. Note that `geothermal_heat_flux` is the same variable read by the [`GeothermalHeatFlux`](@ref) bottom boundary condition, so the initial profile and the boundary condition stay consistent. ```@example soilmodel diff --git a/docs/src/running/initialization.md b/docs/src/running/initialization.md index d86ce9c3a..744887f5c 100644 --- a/docs/src/running/initialization.md +++ b/docs/src/running/initialization.md @@ -37,7 +37,7 @@ The keyword argument must be a `NamedTuple` where the keys correspond to the nam Reusable, parameterized initializers of a single `Field` subtype [`AbstractFieldInitializer`](@ref). They are accepted everywhere a `set!`-compatible value is: in the `initializers` keyword argument, as the default of an `input` variable, or as a parameter of a model initializer. -In the latter case, their `@param` fields become parameters of the model and they are re-evaluated at every initialization. +In the latter case, their `@param` fields become parameters of the model. ```@docs; canonical = false AbstractFieldInitializer diff --git a/src/initializers.jl b/src/initializers.jl index 2e07ab575..2d804f213 100644 --- a/src/initializers.jl +++ b/src/initializers.jl @@ -44,7 +44,7 @@ at the node coordinates of `field`, i.e. `init(coords...)`, via `set!`. A field initializer is accepted everywhere a `set!`-compatible value is: as the `default` of an `input` variable, as an entry of the `initializers` keyword argument of [`initialize`](@ref), or as a parameter of a model [`AbstractInitializer`](@ref). In the latter case it may declare `@param` fields, which are -then exposed as model parameters, and the model initializer re-evaluates it at every `initialize!`. +then exposed as model parameters. """ abstract type AbstractFieldInitializer end @@ -58,16 +58,6 @@ initialize!(field::AbstractField, grid::AbstractGrid, init::AbstractFieldInitial Oceananigans.Fields.set!(field::Field, init::AbstractFieldInitializer) = initialize!(field, field.grid, init) -""" - $TYPEDSIGNATURES - -Re-evaluate `init` into `field` if it is an [`AbstractFieldInitializer`](@ref); all other initial values -(numbers, functions, arrays, `Field`s) are applied only once as construction-time defaults and are left -untouched here so that they remain overridable by [`InputSource`](@ref)s. -""" -reinitialize!(field::AbstractField, init::AbstractFieldInitializer) = set!(field, init) -reinitialize!(::AbstractField, ::Any) = nothing - """ $TYPEDEF diff --git a/src/models/soil/soil_model_init.jl b/src/models/soil/soil_model_init.jl index 00f29b599..fc99fabc3 100644 --- a/src/models/soil/soil_model_init.jl +++ b/src/models/soil/soil_model_init.jl @@ -45,10 +45,9 @@ end Initializer for soil/ground temperature that sets the temperature profile of each column to the value of the `initial_surface_temperature` input variable, whose default is `T₀`. -`T₀` may be a number, a function of the node coordinates, an array, a `Field`, or an -[`AbstractFieldInitializer`](@ref). All but the last are applied once as the default of the input -variable and can be overridden by an [`InputSource`](@ref) named `initial_surface_temperature`; -a field initializer is re-evaluated at every `initialize!` so that its parameters take effect. +`T₀` may be a number, a function of the horizontal node coordinates, an array, a `Field`, or an +[`AbstractFieldInitializer`](@ref). It is re-applied as the default of the input variable at every +initialization and can be overridden by an [`InputSource`](@ref) named `initial_surface_temperature`. Properties: $TYPEDFIELDS @@ -74,7 +73,6 @@ variables(init::ConstantSoilTemperature) = ( ) function initialize!(state, model::AbstractModel, init::ConstantSoilTemperature) - reinitialize!(state.initial_surface_temperature, init.T₀) set!(state.temperature, kernel(compute_constant_temperature, init), state) return nothing end @@ -101,10 +99,9 @@ The surface temperature and the geothermal heat flux are declared as the input v `initial_surface_temperature` (°C) and `geothermal_heat_flux` (W/m²) with defaults `T₀` and `Qgeo`. The latter is the same variable read by the [`GeothermalHeatFlux`](@ref) bottom boundary condition, so that the initial profile and the boundary condition are consistent. Each default may be a number, -a function of the node coordinates, an array, a `Field`, or an [`AbstractFieldInitializer`](@ref). -All but the last are applied once and can be overridden by an [`InputSource`](@ref) of the same name -(with matching units); a field initializer is re-evaluated at every `initialize!` so that its -parameters take effect. +a function of the horizontal node coordinates, an array, a `Field`, or an [`AbstractFieldInitializer`](@ref). +Defaults are re-applied at every initialization, so that changes to parameters take effect, and can be +overridden by an [`InputSource`](@ref) of the same name (with matching units). Properties: $TYPEDFIELDS @@ -138,8 +135,6 @@ variables(init::QuasiThermalSteadyState) = ( ) function initialize!(state, model::AbstractModel, init::QuasiThermalSteadyState) - reinitialize!(state.initial_surface_temperature, init.T₀) - reinitialize!(state.geothermal_heat_flux, init.Qgeo) set!(state.temperature, kernel(compute_quasi_steady_state_temperature, init), state) return nothing end @@ -215,7 +210,6 @@ variables(init::ConstantSaturation) = ( ) function initialize!(state, model::AbstractModel, init::ConstantSaturation) - reinitialize!(state.initial_saturation, init.sat) set!(state.saturation_water_ice, kernel(compute_constant_saturation, init), state) return nothing end @@ -259,8 +253,6 @@ variables(init::SaturationWaterTable) = ( ) function initialize!(state, model::AbstractModel, init::SaturationWaterTable) - reinitialize!(state.vadose_zone_saturation, init.vadose_zone_saturation) - reinitialize!(state.water_table_depth, init.water_table_depth) set!(state.saturation_water_ice, kernel(compute_water_table_saturation, init), state) return nothing end diff --git a/src/state_variables.jl b/src/state_variables.jl index de144bb3f..d91f7f7d2 100644 --- a/src/state_variables.jl +++ b/src/state_variables.jl @@ -116,6 +116,29 @@ function Oceananigans.TimeSteppers.reset!(state::StateVariables) end end +""" + $TYPEDSIGNATURES + +Reset all input `Field`s in `state` to the defaults declared by the corresponding input variables in +`vars` (see [`InputVariable`](@ref)), recursing into namespaces. Inputs without a default are left +unchanged. This is called by `initialize!(::ModelIntegrator)` with the variables of the current model, +before any [`InputSource`](@ref)s are applied, so that defaults always reflect the current model +parameters while input sources take precedence over them. +""" +function reset_inputs!(state::StateVariables, vars::Variables) + for (name, var) in vars.inputs + if !isnothing(var.default) && hasproperty(getfield(state, :inputs), name) + set!(getproperty(getfield(state, :inputs), name), var.default) + end + end + for (name, ns) in vars.namespaces + if hasproperty(getfield(state, :namespaces), name) + reset_inputs!(getproperty(getfield(state, :namespaces), name), Variables(variables(ns))) + end + end + return nothing +end + """ Reset all tendencies in `state` to zero. """ diff --git a/src/timesteppers/model_integrator.jl b/src/timesteppers/model_integrator.jl index 870b6b6c2..c472f45d7 100644 --- a/src/timesteppers/model_integrator.jl +++ b/src/timesteppers/model_integrator.jl @@ -160,6 +160,8 @@ function initialize!(integrator::ModelIntegrator) # reset state variables and clock reset!(integrator.state) reset!(integrator.clock) + # reset inputs to the defaults declared by the (possibly reconstructed) model + reset_inputs!(integrator.state, Variables(variables(integrator.model))) # set inputs based on updated clock/state initialize!(integrator.state, get_grid(integrator.model), integrator.inputs) # evaluate user-specified field initializers diff --git a/test/initializers.jl b/test/initializers.jl index 0179683e8..bb5a21fef 100644 --- a/test/initializers.jl +++ b/test/initializers.jl @@ -40,6 +40,20 @@ end @test interior(state.temperature) == T_before end +@testset "Scalar initial values are parameters" begin + NF = Float32 + grid = ColumnGrid(CPU(), NF, ExponentialSpacing(N = 10)) + energy = QuasiThermalSteadyState(NF; T₀ = NF(1), Qgeo = NF(0.05)) + integrator = initialize(SoilModel(grid; initializer = SoilInitializer(NF; energy))) + ps = vec(parameters(integrator.model)) + ps.initializer.energy.T₀ = NF(4) + # the existing state is reused, so the new default must be re-applied on initialization + integrator = initialize(integrator, ps) + @test all(interior(integrator.state.initial_surface_temperature) .== NF(4)) + z = znodes(integrator.state.temperature) + @test interior(integrator.state.temperature)[1, 1, :] ≈ NF(4) .- NF(0.05) .* z +end + @testset "Column-varying initial values" begin NF = Float32 grid = ColumnGrid(CPU(), NF, ExponentialSpacing(N = 5), 3) @@ -71,6 +85,13 @@ end @test all(interior(integrator.state.geothermal_heat_flux) .== NF(0.1)) z = znodes(integrator.state.temperature) @test interior(integrator.state.temperature)[1, 1, :] ≈ -NF(0.1) .* z + # input sources also take precedence over field initializers + climatology_model = SoilModel(grid; initializer = SoilInitializer(NF; energy = QuasiThermalSteadyState(NF; T₀ = LatitudinalClimatology(NF)))) + T₀_field = Field(grid, Terrarium.Ground(XY())) + set!(T₀_field, NF(3)) + T₀_source = InputSource(grid, T₀_field; name = :initial_surface_temperature, domain = Terrarium.Ground(), units = u"°C") + integrator = initialize(climatology_model; inputs = InputSources(T₀_source)) + @test all(interior(integrator.state.initial_surface_temperature) .== NF(3)) # mismatched units are reported as a conflict bad_source = InputSource(grid, Qgeo_field; name = :geothermal_heat_flux, domain = Terrarium.Ground()) @test_throws ErrorException initialize(model; inputs = InputSources(bad_source)) From 8083e9453ab2b04cf7555b4aff3409136cdeac49 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Fri, 9 Oct 2026 15:45:51 +0200 Subject: [PATCH 19/19] Soften rule on running Enzyme/Reactant tests --- AGENTS.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 2630039eb..daa3f5fb3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -31,8 +31,9 @@ dynamics are allowed except in very special cases where they must be clearly doc include("test/soil/soil_energy_tests.jl") ``` - To run the full suite, use `julia --project=. -e 'using Pkg; Pkg.test()'`. -- **Never run the Enzyme (`Pkg.test(; test_args=["enzyme"])`) or Reactant (`test/reactant/`) test suites - locally unless explicitly requested**. These tests should be left to the CI/CD pipeline. +- Only run the Enzyme (`Pkg.test(; test_args=["enzyme"])`) or Reactant (`test/reactant/`) test suites + locally after directly editing the respective extension modules or unit tests, or when explicitly requested. + These tests should otherwise be left to the CI/CD pipeline. - Julia errors and associated stack traces are often very long due to long type signatures. To mitigate this, always write test output to temporary files and analyze this output using `grep` and similar tools.