Generalize VarDims to VarLocation and make domains explicit in all variables definitions

docs/src/extending/coupling_processes.md

@@ -22,8 +22,8 @@ A concrete example of indirect coupling in Terrarium is the `ground_temperature`
 
 ```julia
 variables(energy::SoilThermodynamics) = (
-    prognostic(:internal_energy, XYZ(); closure = energy.closure, units = u"J/m^3", desc = "Internal energy of the soil volume, including both latent and sensible components"),
-    auxiliary(:ground_temperature, XY(), ground_temperature, energy, units = u"°C", desc = "Temperature of the uppermost ground or soil grid cell in °C"),
+    prognostic(:internal_energy, Ground(XYZ()); closure = energy.closure, units = u"J/m^3", desc = "Internal energy of the soil volume, including both latent and sensible components"),
+    auxiliary(:ground_temperature, Ground(Top(z = Center())), ground_temperature, energy, units = u"°C", desc = "Temperature of the uppermost ground or soil grid cell in °C"),
 )
 
 function ground_temperature(grid, clock, fields, energy::SoilThermodynamics)
@@ -36,7 +36,7 @@ end
 This `ground_temperature` is then consumed by several other processes such as the [surface energy balance](@ref surface_energy_balance_docs), [ evapotranspiration](@ref "Evapotranspiration"), and the vegetation stress factors computed for [autotrophic respiration](@ref "Autotrophic respiration"). These processes can simply declare `ground_temperature` as an input variable:
 
 ```julia
-input(:ground_temperature, XY(), default = 10.0, units = u"°C")
+input(:ground_temperature, Ground(Top(z = Center())), default = 10.0, units = u"°C")
 ```
 
 When used standalone, processes declaring `ground_temperature` as an input in this manner will allocate it as an independent input `Field` (here with a uniform initial value of 10°C across space). This can be very helpful for isolated testing of such components under a range of different input values for `ground_temperature`.
@@ -124,7 +124,7 @@ As an example, consider the `compute_auxiliary!` method for [`PALADYNCanopyEvapo
 function compute_auxiliary!(state, grid,
         evap::PALADYNCanopyEvapotranspiration,
         canopy_interception::AbstractCanopyInterception,  # 1. sibling (surface hydrology)
-        atmos::AbstractAtmosphere,                        # 2. atmosphere
+        atmos::AbstractAtmosphere,                        # 2. atmospheric inputs
         constants::PhysicalConstants,                     # 3. constants
         soil::Optional{AbstractSoil} = nothing            # 4. foreign (soil)
     )

docs/src/extending/state_variables.md

@@ -6,6 +6,7 @@ CurrentModule = Terrarium
 
 ```@setup variables
 using Terrarium
+using Terrarium: Ground, Snow, Canopy, Surface, Atmosphere, Top, Bottom
 using Oceananigans
 ```
 
@@ -33,13 +34,20 @@ Most state variables will thus be defined by implementation of `AbstractProcess`
 struct MyProcess{NF} <: Terrarium.AbstractProcess{NF} end
 
 Terrarium.variables(::MyProcess) = (
-    Terrarium.prognostic(:progvar, XYZ()),
-    Terrarium.auxiliary(:auxvar, XYZ()),
-    Terrarium.auxiliary(:bc, XY()),
+    Terrarium.prognostic(:progvar, Ground(XYZ())),
+    Terrarium.auxiliary(:auxvar, Ground(XYZ())),
+    Terrarium.auxiliary(:top_flux, Ground(Top())),
+    Terrarium.auxiliary(:bottom_flux, Ground(Bottom())),
     Terrarium.input(:input, XY())
 )
 ```
-This will result in a total of five state variables being allocated upon initialization: one input variable, two auxiliary variables named `auxvar` and `bc` and one prognostic variable named `progvar` along with its corresponding tendency variable which is created automatically. The second argument to the variable metadata constructors `prognostic` and `auxiliary` is a subtype of `VarDims` which specifies on which spatial dimensions the state variable should be defined. [`XYZ()`](@ref) corresponds to a 3D `Field` which varies both laterally and with depth. [`XY()`](@ref) corresponds to a 2D field which is discretized along the lateral X and Y dimensions only.
+This will result in a total of five state variables being allocated upon initialization: one input variable, two auxiliary variables named `auxvar` and `bc` and one prognostic variable named `progvar` along with its corresponding tendency variable which is created automatically. The second argument to the variable metadata constructors `prognostic` and `auxiliary` is a [`VarDims`](@ref) or [`VarLocation`](@ref) which specifies the dimensions, and optionally domain, of the variable in space. Variables with dimensions [`XYZ()`](@ref) are allocated a 3D `Field` varying both laterally in the X and Y dimensions as well as with elevation/depth `Z`, while [`XY()`](@ref) corresponds to a 2D "reduced" `Field` discretized only along the lateral X and Y dimensions.
+
+Both `XY` and `XYZ` are aliases for [`VarDims`](@ref) and accept keyword arguments `x`, `y`, and `z`, each of which can be set to one of [`Center`](@extraref Oceananigans.Fields.Center), [`Face`](@extraref Oceananigans.Fields.Face), or a prespecified [`Coordinate`](@ref). `Center` and `Face` declare the variable's discretized location along Oceananigans' staggered finite volume grids (i.e. cell centers vs. faces) while `Coordinate` restricts the variable to a specific point along the axis. This point may be represented either by a hardcoded integer index (not generally recommended for values $>1$) or a function `f(axis)` that computes the index dynamically from the given `axis` at `Field` construction time. Terrarium provides convenience dispatches covering two common cases: [`Top`](@ref) and [`Bottom`](@ref) which correspond to `XY` fields located respectively at the top or bottom of the vertical domain. This is the suitable for choice for fluxes which are applied as boundary conditions to another `Field`, as implied above by `top_flux` and `bottom_flux`.
+
+The outer [`VarDomain`](@ref), `Ground` in the above example, indicates the spatial domain on which the variables should be discretized. Currently, Terrarium defines five `VarDomain`s: `Ground`, `Snow`, `Canopy`, `Surface`, and `Atmosphere`, with the first three mapping to distinct vertical discretizations in [`LandGrid`](@ref)s. In contrast, the [`Surface`](@ref) domain refers to the interface between the land and atmosphere, while [`Atmosphere`](@ref) is reserved for atmospheric forcing variables. Neither is discretized vertically, so their variables must be declared with dimensions `XY`.
+
+A variable may also be declared with bare dimensions and no domain at all, as show above for the `input` variable. Such a declaration indicates that the code is agnostic to where the variable lives: it is compatible with any domain and will automatically promote the domain to match conflicting definitions of the same variable. [`InputSource`](@ref)s always default to assigning their declared input variables `domain = nothing` unless otherwise specified.
 
 ## Merging and promotion rules
 

docs/src/running/input_sources.md

@@ -27,11 +27,12 @@ update_inputs!(inputs, grid, clock, fields, input::InputSource)
 A [`FieldInputSource`](@ref) holds a single `Field` that is copied into the state once at initialization and is thereafter unchanged. This is the appropriate input source for spatially-varying but time-constant forcings (e.g. maps of soil properties or prescribed climatology).
 
 ```@docs; canonical = false
-InputSource(grid::AbstractLandGrid{NF}, field::FS) where {NF, FS <: AnyField{NF}}
+InputSource(grid::AbstractGrid{NF}, field::FS; name, domain, units) where {NF, FS <: AnyField{NF}}
 ```
 
 ```julia
 using Oceananigans: Field
+using Terrarium: Atmosphere, Ground, Snow, Surface
 
 # Existing Field or array on the model grid
 albedo_field = Field(grid_2d)
@@ -56,13 +57,13 @@ using Oceananigans.Units: hours
 
 # Allocate and populate a FieldTimeSeries
 times = 0.0:3600.0:86400.0 # hourly for one day (seconds)
-fts = FieldTimeSeries(grid, XY(), times)
+fts = FieldTimeSeries(grid, Atmosphere(XY()), times)
 fts.data .= randn(size(fts)) # fill with data
 source = InputSource(fts; name = :air_temperature, units = u"°C")
 ```
 
 ```@docs; canonical = false
-InputSource(grid::AbstractLandGrid{NF}, field::FS) where {NF, FS <: AnyFieldTimeSeries{NF}}
+InputSource(fts::AnyFieldTimeSeries{NF}; name, domain, reftime, units) where {NF}
 ```
 
 The `FieldTimeSeries` can also be loaded from a file using the relevant constructors provided by Oceananigans.
@@ -157,9 +158,9 @@ The minimum set of input fields needed by a process is declared by including `in
 
 ```julia
 Terrarium.variables(snow::DegreeDaySnow{NF}) where {NF} = (
-    input(:air_temperature, XY(), units = u"°C"),
-    input(:snow_fall,       XY(), units = u"m/s"),
-    prognostic(:snow_storage, XY()),
+    input(:air_temperature, Atmosphere(XY()), units = u"°C"),
+    input(:snow_fall,       Snow(XY()), units = u"m/s"),
+    prognostic(:snow_storage, Snow(XY())),
 )
 ```
 
@@ -176,15 +177,15 @@ To add a new input source backend:
 3. Implement `initialize!(fields, source::MySource, clock)` for any one-time setup.
 4. Implement `update_inputs!(fields, source::MySource, clock::Clock)` to update the
    input field at each time step.
-5. Optionally provide a convenience `InputSource(grid, ...; name, units)` constructor
+5. Optionally provide a convenience `InputSource(grid, ...; name, units, domain)` constructor
    dispatch so users do not need to reference the concrete type name.
 
 ```julia
 struct MyInputSource{NF} <: InputSource{NF, :my_var}
     data::Vector{NF}
 end
 
-Terrarium.variables(::MyInputSource{NF}) where {NF} = (input(:my_var, XY()),)
+Terrarium.variables(::MyInputSource{NF}) where {NF} = (input(:my_var, XY())y),)
 
 function Terrarium.update_inputs!(fields, source::MyInputSource, clock::Clock)
     # populate fields.my_var from source.data at clock.time

examples/extending/linear_heat_conduction.jl

@@ -18,6 +18,7 @@
 # this implementation is meant only to serve as an example.
 
 using Terrarium
+using Terrarium: XYZ
 using KernelAbstractions: @kernel, @index
 using Oceananigans.Operators: ∂zᵃᵃᶜ, ∂zᵃᵃᶠ
 using Oceananigans.Utils: launch!
@@ -52,7 +53,7 @@ LinearHeatConduction(::Type{NF}; kwargs...) where {NF} = LinearHeatConduction{NF
 # by this process. Temperature is a 3D column variable ([`XYZ`](@ref)) since it varies with depth.
 
 Terrarium.variables(::LinearHeatConduction) = (
-    Terrarium.prognostic(:temperature, Terrarium.XYZ(); units = u"°C"),
+    Terrarium.prognostic(:temperature, XYZ(); units = u"°C"),
 )
 
 # `prognostic` means the timestepper integrates this variable based on its tendency at each
@@ -133,9 +134,13 @@ end
 # physics and dispatch logic should be defined in the `compute_*` kernel functions.
 #
 # !!! note "Always use 3D indexing in kernel functions"
-#     Even for surface (`XY`) fields, write output as `out.name[i, j, 1]` (with `k = 1`).
-#     2D indexing (`out.name[i, j]`), especially in `setindex!`, will result in errors when
-#     compiling the kernel on GPU.
+#     Write output as `out.name[i, j, k]`. 2D indexing (`out.name[i, j]`), especially in
+#     `setindex!`, will result in errors when compiling the kernel on GPU.
+#
+#     For a field with no vertical extent, write `out.name[i, j, end]` rather than a literal
+#     index. A variable declared at the `Top` or `Bottom` of a domain is stored at that
+#     interface, not at `k = 1`, so `end` is the index which is correct for every such field;
+#     a literal `1` reads or writes outside the field, silently so under `@inbounds`.
 
 @kernel inbounds = true function compute_tendencies_kernel!(
         tendencies, grid, fields, proc::AbstractHeatConduction, args...
@@ -177,7 +182,7 @@ end
 
 @kwdef struct HeatModel{
         NF,
-        Grid <: Terrarium.AbstractLandGrid{NF},
+        Grid <: Terrarium.AbstractGrid{NF},
         Cond <: AbstractHeatConduction{NF},
         Init <: Terrarium.AbstractInitializer{NF},
         TS <: Terrarium.AbstractTimeStepper,

examples/extending/linear_ode_exp_growth.jl

@@ -55,7 +55,7 @@ grid = ColumnGrid(CPU(), Float64, UniformSpacing(Δz = 0.1, N = 1))
     c::NF = 0.1
 end
 #
-@kwdef struct ExpModel{NF, Grid <: Terrarium.AbstractLandGrid{NF}, Dyn, Init, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid}
+@kwdef struct ExpModel{NF, Grid <: Terrarium.AbstractGrid{NF}, Dyn, Init, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid}
     "Spatial grid on which state variables are discretized"
     grid::Grid
     "Linear dynamics process"
@@ -139,7 +139,7 @@ Random.seed!(1234) # set random seed
 t_F = 0:1:300; #seconds
 F = FieldTimeSeries(grid, XY(), t_F);
 F.data .= randn(size(F));
-input = InputSource(grid, F, name = :F)
+input = InputSource(grid, F; name = :F)
 
 # Here we constructed a 2D (`XY()`) time series on our `grid` at times `t_F` with random normal distributed data and defined our `InputSource` for our model based on it.
 

examples/extending/simple_snow_ddm.jl

@@ -56,12 +56,12 @@ import DisplayAs
 end
 #
 Terrarium.variables(model::DegreeDaySnow{NF}) where {NF} = (
-    Terrarium.input(:air_temperature, XY(), default = NF(0), units = u"°C", desc = "Near-surface air temperature in °C"),
-    Terrarium.input(:snow_fall, XY(), default = NF(0), units = u"m/s", desc = "snow fall rate in m/s"),
-    Terrarium.prognostic(:snow_storage, XY(), units = u"m", desc = "Snow water equivalent in m"),
+    Terrarium.input(:air_temperature, Terrarium.Atmosphere(XY()), default = NF(0), units = u"°C", desc = "Near-surface air temperature in °C"),
+    Terrarium.input(:snow_fall, Terrarium.Atmosphere(XY()), default = NF(0), units = u"m/s", desc = "snow fall rate in m/s"),
+    Terrarium.prognostic(:snow_storage, Terrarium.Snow(XY()), units = u"m", desc = "Snow water equivalent in m"),
 )
 
-@kwdef struct SnowModel{NF, Grid <: Terrarium.AbstractLandGrid{NF}, Pro, Init, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid}
+@kwdef struct SnowModel{NF, Grid <: Terrarium.AbstractGrid{NF}, Pro, Init, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid}
     "Spatial grid on which state variables are discretized"
     grid::Grid
     "Snow melting process"
@@ -116,8 +116,8 @@ Terrarium.compute_auxiliary!(state, grid, model::DegreeDaySnow) = nothing
 
 function compute_snow_flux_tendency(i, j, grid, fields, snow_melt)
     ## get the variables we need
-    P = fields.snow_fall[i, j]
-    T = fields.air_temperature[i, j]
+    P = fields.snow_fall[i, j, 1]
+    T = fields.air_temperature[i, j, 1]
     ## get the parameters
     T_melt = snow_melt.T_melt
     k = snow_melt.k

examples/hybrid_models/neural_snow_melt.jl

@@ -96,16 +96,16 @@ Adapt.@adapt_structure NeuralSnowMeltBatched
 # the model.
 const SnowVars{NF} = Union{DegreeDaySnow{NF}, NeuralSnowMelt{NF}, NeuralSnowMeltBatched{NF}}
 Terrarium.variables(::SnowVars{NF}) where {NF} = (
-    Terrarium.input(:air_temperature, XY(), default = NF(0), units = u"°C", desc = "Near-surface air temperature in °C"),
-    Terrarium.input(:snow_fall, XY(), default = NF(0), units = u"m/s", desc = "Snow fall rate in m/s"),
-    Terrarium.prognostic(:snow_storage, XY(), units = u"m", desc = "Snow water equivalent in m"),
+    Terrarium.input(:air_temperature, Terrarium.Atmosphere(XY()), default = NF(0), units = u"°C", desc = "Near-surface air temperature in °C"),
+    Terrarium.input(:snow_fall, Terrarium.Snow(XY()), default = NF(0), units = u"m/s", desc = "Snow fall rate in m/s"),
+    Terrarium.prognostic(:snow_storage, Terrarium.Snow(XY()), units = u"m", desc = "Snow water equivalent in m"),
 )
 
 # ## The model
 #
 # A minimal `AbstractModel` holding one snow-melt process (either kind).
 
-@kwdef struct SnowModel{NF, Grid <: Terrarium.AbstractLandGrid{NF}, Pro, Init, TS} <: Terrarium.AbstractModel{NF, Grid}
+@kwdef struct SnowModel{NF, Grid <: Terrarium.AbstractGrid{NF}, Pro, Init, TS} <: Terrarium.AbstractModel{NF, Grid}
     grid::Grid
     snow_melt::Pro = DegreeDaySnow(eltype(grid))
     initializer::Init = DefaultInitializer(eltype(grid))
@@ -130,11 +130,11 @@ function Terrarium.compute_tendencies!(state, grid, snow_melt::DegreeDaySnow)
 end
 @kernel function ddm_snow_flux_kernel!(tend, grid, fields, snow_melt)
     i, j = @index(Global, NTuple)
-    @inbounds tend.snow_storage[i, j, 1] = ddm_snow_flux(i, j, fields, snow_melt)
+    @inbounds tend.snow_storage[i, j, end] = ddm_snow_flux(i, j, fields, snow_melt)
 end
 @inline function ddm_snow_flux(i, j, fields, snow_melt::DegreeDaySnow)
-    P = @inbounds fields.snow_fall[i, j, 1]
-    T = @inbounds fields.air_temperature[i, j, 1]
+    P = @inbounds fields.snow_fall[i, j, end]
+    T = @inbounds fields.air_temperature[i, j, end]
     return P - melt(snow_melt, T)
 end
 
@@ -147,11 +147,11 @@ function Terrarium.compute_tendencies!(state, grid, snow_melt::NeuralSnowMelt)
 end
 @kernel function nn_snow_flux_kernel!(tend, grid, fields, snow_melt)
     i, j = @index(Global, NTuple)
-    @inbounds tend.snow_storage[i, j, 1] = nn_snow_flux(i, j, fields, snow_melt)
+    @inbounds tend.snow_storage[i, j, end] = nn_snow_flux(i, j, fields, snow_melt)
 end
 @inline function nn_snow_flux(i, j, fields, p::NeuralSnowMelt)
-    P = @inbounds fields.snow_fall[i, j, 1]
-    T = @inbounds fields.air_temperature[i, j, 1]
+    P = @inbounds fields.snow_fall[i, j, end]
+    T = @inbounds fields.air_temperature[i, j, end]
     Tn = (T - p.T_mean) / p.T_std
     ## KernelLux evaluates the MLP inside the kernel; the input is a 1-element static vector.
     y, _ = apply_in_kernel(p.model, SA[Tn], p.ps, p.st)
@@ -264,7 +264,7 @@ snow_fall = fill(NF(1.0e-7), length(_lats))                  # uniform light sno
 ## input sources take a RingGrids.Field over the grid's rings
 inputs = InputSources(
     InputSource(grid, RingGrids.Field(air_temperature, rings), name = :air_temperature, units = u"°C"),
-    InputSource(grid, RingGrids.Field(snow_fall, rings), name = :snow_fall, units = u"m/s"),
+    InputSource(grid, RingGrids.Field(snow_fall, rings), name = :snow_fall, units = u"m/s"; domain = Terrarium.Snow()),
 )
 initializers = (snow_storage = NF(0.5),)                      # start with 0.5 m everywhere
 
@@ -320,7 +320,7 @@ const Nt_ft = 15           # 15-day rollout
 device_grid = ColumnRingGrid(ReactantState(), NF, UniformSpacing(Δz = 0.1, N = 1), rings)
 device_inputs = InputSources(
     InputSource(device_grid, RingGrids.Field(air_temperature, rings), name = :air_temperature, units = u"°C"),
-    InputSource(device_grid, RingGrids.Field(snow_fall, rings), name = :snow_fall, units = u"m/s"),
+    InputSource(device_grid, RingGrids.Field(snow_fall, rings), name = :snow_fall, units = u"m/s"; domain = Terrarium.Snow()),
 )
 build_device_integrator(process) = initialize(
     SnowModel(device_grid; snow_melt = process, timestepper = ForwardEuler(NF));

examples/simulations/land_global_era5.jl

@@ -162,7 +162,8 @@ inputs = InputSources(
         name = :leaf_area_index,
         source_grid = Terrarium.native_grid(lai_asset),
         timedim = :dayofyear,
-        cycle = true
+        cycle = true,
+        domain = Terrarium.Canopy()
     ),
 )
 

examples/simulations/soil_heat_global_soilgrids.jl

@@ -66,7 +66,7 @@ function Terrarium.InputSources(dataset::SoilGrids2, grid::ColumnRingGrid, horiz
             ring_field = RingGrids.on_architecture(arch, RingGrids.FullClenshawField(interior(var_field)[:, (end - 1):-1:2, end - idx + 1], input_as = Matrix))
             target_field = RingGrids.Field(grid.rings)
             RingGrids.interpolate!(target_field, ring_field)
-            layer_inputs[var] = InputSource(grid, Field(target_field, grid); name = horizon => var)
+            layer_inputs[var] = InputSource(grid, Field(target_field, grid); name = horizon => var, domain = Terrarium.Ground())
         end
         # Ensure that mineral texture components with each horizon sum to unity
         Terrarium.normalize_texture!(layer_inputs[:sand_fraction].field, layer_inputs[:silt_fraction].field, layer_inputs[:clay_fraction].field)

examples/simulations/speedy_dry_land.jl

@@ -26,7 +26,7 @@ soil_initializer = SoilInitializer(eltype(grid))
 # Soil model with a prescribed surface-temperature boundary condition driven
 # by SpeedyWeather's near-surface air temperature.
 soil_model = SoilModel(grid; initializer = soil_initializer)
-air_temperature_field = Field(grid, XY())
+air_temperature_field = Field(grid, Terrarium.Surface(XY()))
 Tair_input = InputSource(grid, air_temperature_field; name = :air_temperature)
 bcs = PrescribedSurfaceTemperature(:air_temperature)