Add support for spatially varying and parameterized model initializers

AGENTS.md

@@ -30,8 +30,10 @@ 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()'`.
+- 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.
@@ -259,6 +261,10 @@ Each document should be prefaced by the following template:
 
 Date of initial draft: YYYY-MM-dd
 
+Author: <Name of model and version (if known)>, <harness name and version (if known)>
+
+Reviewer: <Name of human reviewer> (<reviewer email address>)
+
 Base revision: <SHA1 of HEAD when plan was drafted>
 
 ## Originating prompt
@@ -269,15 +275,26 @@ Base revision: <SHA1 of HEAD when plan was drafted>
 
 > 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.
 
 The remainder of the plan document may be adapted on a case-by-case basis but should generally follow this structure:
 

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

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:

docs/src/models/soil_model.md

@@ -75,6 +75,31 @@ 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).
+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) 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
+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 +120,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

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.
+
+```@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.

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"

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,22 +176,20 @@ 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)[land_mask_cpu]
+t2m_mean_on_grid = RingGrids.interpolate(on_architecture(CPU(), grid.rings), t2m_mean_field)
 
-function initial_soil_temperature(x, z)
-    T₀ = t2m_mean_on_grid[round(Int, x)]
-    T = T₀ - NF(0.05) * z
-    return T
-end
+# 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 = (
-    temperature = initial_soil_temperature,
     saturation_water_ice = (x, z) -> min(one(NF), NF(0.6) - NF(0.05) * z),
 )
 integrator = initialize(land; inputs, initializers)

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)