diff --git a/docs/src/extending/coupling_processes.md b/docs/src/extending/coupling_processes.md index 83d0ce944..d19518c9e 100644 --- a/docs/src/extending/coupling_processes.md +++ b/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) ) diff --git a/docs/src/extending/state_variables.md b/docs/src/extending/state_variables.md index ee9876982..dfb217888 100644 --- a/docs/src/extending/state_variables.md +++ b/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 diff --git a/docs/src/running/input_sources.md b/docs/src/running/input_sources.md index 992e26fa3..cd16e3410 100644 --- a/docs/src/running/input_sources.md +++ b/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,7 +177,7 @@ 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 @@ -184,7 +185,7 @@ 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 diff --git a/examples/extending/linear_heat_conduction.jl b/examples/extending/linear_heat_conduction.jl index 318d0b736..74942bd2f 100644 --- a/examples/extending/linear_heat_conduction.jl +++ b/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, diff --git a/examples/extending/linear_ode_exp_growth.jl b/examples/extending/linear_ode_exp_growth.jl index b0ed2b024..af6110944 100644 --- a/examples/extending/linear_ode_exp_growth.jl +++ b/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. diff --git a/examples/extending/simple_snow_ddm.jl b/examples/extending/simple_snow_ddm.jl index 0fac68172..569b7daa4 100644 --- a/examples/extending/simple_snow_ddm.jl +++ b/examples/extending/simple_snow_ddm.jl @@ -61,7 +61,7 @@ Terrarium.variables(model::DegreeDaySnow{NF}) where {NF} = ( Terrarium.prognostic(:snow_storage, 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 diff --git a/examples/hybrid_models/neural_snow_melt.jl b/examples/hybrid_models/neural_snow_melt.jl index 51207f9ff..25b3da2e5 100644 --- a/examples/hybrid_models/neural_snow_melt.jl +++ b/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)); diff --git a/examples/simulations/land_global_era5.jl b/examples/simulations/land_global_era5.jl index 38d8dc14b..668c15f3c 100644 --- a/examples/simulations/land_global_era5.jl +++ b/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() ), ) diff --git a/examples/simulations/soil_heat_global_soilgrids.jl b/examples/simulations/soil_heat_global_soilgrids.jl index 9ebb62ec3..b01eabb62 100644 --- a/examples/simulations/soil_heat_global_soilgrids.jl +++ b/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) diff --git a/examples/simulations/speedy_dry_land.jl b/examples/simulations/speedy_dry_land.jl index 3dff8da2d..d91dac142 100644 --- a/examples/simulations/speedy_dry_land.jl +++ b/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) diff --git a/examples/simulations/speedy_wet_land.jl b/examples/simulations/speedy_wet_land.jl index 334d48a42..7da920cf9 100644 --- a/examples/simulations/speedy_wet_land.jl +++ b/examples/simulations/speedy_wet_land.jl @@ -58,7 +58,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) @@ -95,7 +95,7 @@ terrarium_model = Terrarium.LandModel(land_grid; vegetation, soil, surface_energ 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) +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 diff --git a/examples/simulations/vegetation_global.jl b/examples/simulations/vegetation_global.jl index ef30d05ea..af3478d05 100644 --- a/examples/simulations/vegetation_global.jl +++ b/examples/simulations/vegetation_global.jl @@ -41,7 +41,7 @@ lai_asset_path = Terrarium.get_asset(lai_asset) # cycle = true` makes the LAI climatology repeat every year over the whole simulation. lai_raster = RasterStack(lai_asset_path, lazy = true) lai_highveg = convert.(NF, replace_missing(lai_raster[:lai_hv], zero(NF))) -lai_input = InputSource(grid, lai_highveg; source_grid = Terrarium.native_grid(lai_asset), name = :leaf_area_index, cycle = true) +lai_input = InputSource(grid, lai_highveg; source_grid = Terrarium.native_grid(lai_asset), name = :leaf_area_index, cycle = true, domain = Terrarium.Canopy()) # ## Vegetation model # Here we set up a [`PrescribedVegetation`](@ref) model which uses [`PrescribedPhenology`](@ref) to take LAI as a prescribed input. diff --git a/ext/TerrariumRastersExt/TerrariumRastersExt.jl b/ext/TerrariumRastersExt/TerrariumRastersExt.jl index a27a57f50..d36f7bb8d 100644 --- a/ext/TerrariumRastersExt/TerrariumRastersExt.jl +++ b/ext/TerrariumRastersExt/TerrariumRastersExt.jl @@ -31,9 +31,9 @@ When `cycle = true`, the time axis is treated as periodic (e.g. an annual climat whole simulation. This type should generally not be constructed directly but rather via the `InputSource` constructor interface; static (time-invariant) rasters instead yield a plain `FieldInputSource`. """ -struct RasterInputSource{NF, name, VD, TT, IM <: AbstractVector{Int}, RS <: AbstractRaster{NF}, SG, UT, EX <: Extrapolation} <: InputSource{NF, name} - "Variable dimensions" - dims::VD +struct RasterInputSource{NF, name, VL, TT, IM <: AbstractVector{Int}, RS <: AbstractRaster{NF}, SG, UT, EX <: Extrapolation} <: InputSource{NF, name} + "Variable location: spatial dimensions and the model domain the variable lives on" + loc::VL "Physical units" units::UT @@ -67,6 +67,7 @@ function Terrarium.InputSource( raster::AbstractRaster{NF}; source_grid = grid.rings, name = raster.name, + domain::Terrarium.Optional{Terrarium.VarDomain} = nothing, units = NoUnits, timedim = Ti, reftime = nothing, @@ -78,7 +79,7 @@ function Terrarium.InputSource( td = dims(raster, timedim) raster = set(raster, td => Ti(td.val)) end - return RasterInputSource(grid, source_grid, raster, dims(raster, Ti), name, units, reftime, extrapolation) + return RasterInputSource(grid, source_grid, raster, dims(raster, Ti), name, domain, units, reftime, extrapolation) end # Land grids delegate to the discretization of their ground domain, which carries the ring grid. @@ -90,7 +91,7 @@ function RasterInputSource( source_grid::RingGrids.AbstractGrid, raster::AbstractRaster{NF}, ::Nothing, - name, units, args... + name, domain, units, args... ) where {NF} idxmap = findall(Array(grid.mask)) field = similar(grid.mask, NF) @@ -102,7 +103,7 @@ function RasterInputSource( RingGrids.interpolate!(field, source_field) end # Construct and return a FieldInputSource from the RingGrids `field` - return Terrarium.InputSource(grid, field; name, units) + return Terrarium.InputSource(grid, field; name, domain, units) end # Time-varying rasters retain the custom source with lazy time interpolation (and optional cycling). @@ -111,24 +112,24 @@ function RasterInputSource( source_grid::RingGrids.AbstractGrid, raster::AbstractRaster{NF}, ::TimeDim, - name, units, reftime, extrapolation + name, domain, units, reftime, extrapolation ) where {NF} # get indices from grid mask idxmap = on_architecture(architecture(grid), findall(Array(grid.mask))) # infer the VarDims and subsequently the Field location from the data dimensions - rdims = Terrarium.vardims(raster) + rloc = Terrarium.VarLocation(Terrarium.vardims(raster), domain) # infer reference time reftime = default_reftime(raster, reftime) raster = Rasters.setdims(raster, convert_time_axis(dims(raster, Ti))) path = Terrarium.varpath(name) - return RasterInputSource{NF, path, typeof(rdims), typeof(reftime), typeof(idxmap), typeof(raster), typeof(source_grid), typeof(units), typeof(extrapolation)}( - rdims, units, idxmap, reftime, raster, source_grid, extrapolation + return RasterInputSource{NF, path, typeof(rloc), typeof(reftime), typeof(idxmap), typeof(raster), typeof(source_grid), typeof(units), typeof(extrapolation)}( + rloc, units, idxmap, reftime, raster, source_grid, extrapolation ) end Terrarium.variables(source::RasterInputSource) = Terrarium.with_scope( Base.front(Terrarium.varpath(source)), - Terrarium.input(Terrarium.varname(source), source.dims; units = source.units) + Terrarium.input(Terrarium.varname(source), source.loc; units = source.units) ) |> tuple # Infer VarDims based on the axes defined in the Raster diff --git a/ext/TerrariumReactantExt/TerrariumReactantExt.jl b/ext/TerrariumReactantExt/TerrariumReactantExt.jl index f0c205a58..72aac60fc 100644 --- a/ext/TerrariumReactantExt/TerrariumReactantExt.jl +++ b/ext/TerrariumReactantExt/TerrariumReactantExt.jl @@ -15,16 +15,16 @@ using Oceananigans using Oceananigans.Architectures: ReactantState, CPU, architecture, on_architecture -using Terrarium: Terrarium, AbstractLandGrid, ColumnRingGrid, AbstractModel, +using Terrarium: Terrarium, AbstractGrid, ColumnRingGrid, AbstractModel, ModelIntegrator, ground_domain, get_grid, get_timestepper const RARCH = ReactantState @inline Terrarium.uses_reactant(::Terrarium.ReactantMarker) = true -# Land grids that live on the device. -const ReactantLandGrid{NF, TX, TY, TZ} = AbstractLandGrid{NF, TX, TY, TZ, <:RARCH} -const ReactantModel{NF} = AbstractModel{NF, <:ReactantLandGrid{NF}} +# Grids and models that live on the device +const ReactantGrid{NF, TX, TY, TZ, ST} = AbstractGrid{NF, TX, TY, TZ, <:RARCH, ST} +const ReactantModel{NF} = AbstractModel{NF, <:ReactantGrid} # Inside the compiled stepping loop `clock.time` is a `TracedRNumber`, and it reaches host-level input # code through `timestamp`/`convert_dt` (e.g. `FieldTimeSeriesInputSource.update_inputs!`). The generic diff --git a/src/Terrarium.jl b/src/Terrarium.jl index 7516b093e..71647e629 100644 --- a/src/Terrarium.jl +++ b/src/Terrarium.jl @@ -19,7 +19,7 @@ using KernelAbstractions: @kernel, @index # Oceananigans numerics using Oceananigans.AbstractOperations: Average, Integral, ConditionalOperation, KernelFunctionOperation using Oceananigans.Architectures: Architectures, AbstractArchitecture, CPU, GPU, ReactantState, architecture, on_architecture, array_type -using Oceananigans.Fields: Field, FunctionField, AbstractField, Center, Face, set!, compute!, interior, location +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, @@ -83,6 +83,11 @@ Alias for Oceananigans `AbstractBoundaryConditionClassification` """ const BCType = AbstractBoundaryConditionClassification +""" +Alias for Oceananigans location types, i.e. `Center` or `Face`. +""" +const CenterOrFace = Union{Center, Face} + # Re-export selected types and methods from Oceananigans export Simulation, Clock, Field, FieldTimeSeries, KernelFunctionOperation, Center, Face export CPU, GPU, ReactantState, architecture, on_architecture @@ -116,6 +121,8 @@ end export @assert_kernel include("utils/utils.jl") +include("domains.jl") + export XY, XYZ include("abstract_variables.jl") diff --git a/src/abstract_model.jl b/src/abstract_model.jl index 4a978d2ce..2ad217174 100644 --- a/src/abstract_model.jl +++ b/src/abstract_model.jl @@ -38,7 +38,7 @@ Implementations of `AbstractModel` are required to implement, at minimum, three Note that a default implementation of `variables` is provided which automatically collects all variables declared by `AbstractProcess`es defined as fields (properties) of `struct`s that subtype `AbstractModel`. """ -abstract type AbstractModel{NF, Grid <: AbstractLandGrid{NF}} end +abstract type AbstractModel{NF, Grid <: AbstractGrid} end # Method interface for AbstractModel and AbstractProcess @@ -165,7 +165,7 @@ Note that this is a type-stable, `@generated` function that is compiled for each end """ - get_grid(model::AbstractModel)::AbstractLandGrid + get_grid(model::AbstractModel)::AbstractGrid Return the spatial grid associated with the given `model`. """ @@ -238,22 +238,12 @@ defined by the coupling interface for the process type. """ invclosure!(state, grid, closure::AbstractClosureRelation, ::AbstractProcess, args...) = nothing -""" - (::Type{Model})(grid::AbstractLandGrid; kwargs...) where {Model <: AbstractModel} - -Convenience constructor for all `AbstractModel` types that accepts `grid` as a positional argument. -""" -(::Type{Model})(grid::AbstractLandGrid; kwargs...) where {Model <: AbstractModel} = Model(; grid, kwargs...) - """ (::Type{Model})(grid::AbstractGrid; kwargs...) where {Model <: AbstractModel} -Default constructor for all `AbstractModel` types that accepts an ordinary spatial discretization -`grid` and builds the model's [`LandGrid`](@ref) from it via [`create_land_grid`](@ref). Models which -resolve additional vertical domains should define their own constructor which passes the relevant -process components to `create_land_grid`. +Convenience constructor for all `AbstractModel` types that accepts `grid` as a positional argument. """ -(::Type{Model})(grid::AbstractGrid; kwargs...) where {Model <: AbstractModel} = Model(create_land_grid(grid); kwargs...) +(::Type{Model})(grid::AbstractGrid; kwargs...) where {Model <: AbstractModel} = Model(; grid, kwargs...) # Default parameters collection for processes function ParameterEditing.parameters(proc::AbstractProcess; kwargs...) diff --git a/src/abstract_variables.jl b/src/abstract_variables.jl index abe305feb..51a3bd296 100644 --- a/src/abstract_variables.jl +++ b/src/abstract_variables.jl @@ -1,47 +1,197 @@ -# Abstract variable types for declaring fields. +""" + $TYPEDEF + +Helper type used for specifying a single point or cross-section in space along the X, Y, or Z axis. +If `val = nothing`, the coordinate is undefined or integrated over the extent of the axis. +""" +@kwdef struct Coordinate{V, L} + val::V = nothing + loc::L = Face() +end -abstract type VarDims end +# Dispatch for Oceananigans `location` method +Oceananigans.location(dims::Coordinate) = dims.loc + +# A `Coordinate` resolves to a single index along its axis. An integer selects that index directly; +# a function (`firstindex`/`lastindex`) is applied to the range of valid indices along the axis, so +# that the same `Coordinate` means the top face (`Nz + 1`) or the top cell center (`Nz`) according to +# the location it carries. +Oceananigans.Fields.indices(axis, dims::Coordinate{<:Integer}) = dims.val +Oceananigans.Fields.indices(axis, dims::Coordinate{<:Function}) = dims.val(axis) + +# Abstract state variable types """ - XYZ <: VarDims + $TYPEDEF -Indicator type for variables that should be assigned a 3D field on their associated grid. +Indicator type describing the location on which a variable should be instantiated on an Oceananigans grid. +The fields `x`, `y`, and `z` should be one of: `Center()` or `Face()` for variables that should be discretized +on grid cell centers or faces along each axis, `Coordinate` for variables that are defined at a single point along an axis, +or `nothing` for variables that represented quantities integrated over a domain. """ -@kwdef struct XYZ{LX, LY, LZ} <: VarDims +@kwdef struct VarDims{LX, LY, LZ} x::LX = Center() y::LY = Center() z::LZ = Center() end -# Dispatch for Oceananigans `location` method -Oceananigans.location(dims::XYZ) = (dims.x, dims.y, dims.z) +# Resolve one axis of a `VarDims` to its Oceananigans location. `nothing` means the variable has no +# extent along that axis; a `Coordinate` contributes the location it is defined at. +@inline axis_location(::Nothing) = nothing +@inline axis_location(loc::CenterOrFace) = loc +@inline axis_location(coord::Coordinate) = location(coord) + +# The range of valid indices of `grid` along dimension `dim` at location `loc`. This is what the +# `firstindex`/`lastindex` of a `Coordinate` are resolved against, and it is why a `Coordinate` must +# carry its location: along a bounded axis there are `N` cell centers but `N + 1` faces. +@inline axis_range(grid::AbstractGrid, dim::Int, loc) = + Base.OneTo(Oceananigans.Grids.total_length(loc, topology(grid, dim)(), size(grid, dim))) + +# Resolve one axis of a `VarDims` to the corresponding entry of the `indices` argument of the `Field` +# constructor. Anything which is not a `Coordinate` spans the whole axis. +@inline axis_indices(_, ::Union{Nothing, CenterOrFace}) = Colon() +@inline axis_indices(axis, coord::Coordinate) = indices(axis, coord) + +Oceananigans.location(dims::VarDims) = (axis_location(dims.x), axis_location(dims.y), axis_location(dims.z)) + +Oceananigans.Fields.indices(grid::AbstractGrid, dims::VarDims) = ( + axis_indices(axis_range(grid, 1, axis_location(dims.x)), dims.x), + axis_indices(axis_range(grid, 2, axis_location(dims.y)), dims.y), + axis_indices(axis_range(grid, 3, axis_location(dims.z)), dims.z), +) + +# VarDims aliases """ - XY <: VarDims + XY(x = Center(), y = Center(), z = nothing) -Indicator type for variables that should be assigned a 2D (lateral only) field on their associated grid. +Dimensions for a variable with no vertical extent, i.e. one assigned a 2D (lateral only) field on +its associated grid. This covers both genuinely two-dimensional quantities and those integrated or +averaged over a domain's vertical extent. """ -@kwdef struct XY{LX, LY} <: VarDims - x::LX = Center() - y::LY = Center() -end +const XY = VarDims{LX, LY, LZ} where {LX <: CenterOrFace, LY <: CenterOrFace, LZ <: Union{Nothing, Coordinate}} + +XY(x::CenterOrFace, y::CenterOrFace = Center(), z::Union{Nothing, Coordinate} = nothing) = VarDims(x, y, z) +XY(; x::CenterOrFace = Center(), y::CenterOrFace = Center(), z::Union{Nothing, Coordinate} = nothing) = VarDims(x, y, z) -Oceananigans.location(dims::XY) = (dims.x, dims.y, nothing) +""" + Top(x = Center(), y = Center(), z = Face()) + +Dimensions for a variable defined at a single point at the *top* of a domain's vertical axis, such +as a surface flux. The vertical index is resolved with `lastindex`, so a field declared this way is +restricted to one `k`: `Nz + 1` at `Face` (the default, appropriate for fluxes across the interface) +and `Nz` at `Center` (the uppermost cell). + +See also [`Bottom`](@ref). +""" +const Top{TZ} = VarDims{LX, LY, LZ} where {LX <: CenterOrFace, LY <: CenterOrFace, TZ <: CenterOrFace, LZ <: Coordinate{typeof(lastindex), TZ}} + +Top(x::CenterOrFace, y::CenterOrFace = Center(), z::CenterOrFace = Face()) = VarDims(x, y, Coordinate(lastindex, z)) +Top(; x::CenterOrFace = Center(), y::CenterOrFace = Center(), z::CenterOrFace = Face()) = VarDims(x, y, Coordinate(lastindex, z)) + +""" + Bottom(x = Center(), y = Center(), z = Face()) + +Dimensions for a variable defined at a single point at the *bottom* of a domain's vertical axis, +such as a basal flux. The vertical index is resolved with `firstindex`, so a field declared this way +is restricted to `k = 1`. + +See also [`Top`](@ref). +""" +const Bottom{TZ} = VarDims{LX, LY, LZ} where {LX <: CenterOrFace, LY <: CenterOrFace, TZ <: CenterOrFace, LZ <: Coordinate{typeof(firstindex), TZ}} -# TODO: do we need to support state variables not defined on a grid? +Bottom(x::CenterOrFace, y::CenterOrFace = Center(), z::CenterOrFace = Face()) = VarDims(x, y, Coordinate(firstindex, z)) +Bottom(; x::CenterOrFace = Center(), y::CenterOrFace = Center(), z::CenterOrFace = Face()) = VarDims(x, y, Coordinate(firstindex, z)) + +""" + XYZ(x = Center(), y = Center(), z = Center()) + +Dimensions for a variable which is resolved over a domain's full vertical extent, i.e. one assigned +a 3D field on its associated grid. +""" +const XYZ = VarDims{LX, LY, LZ} where {LX <: CenterOrFace, LY <: CenterOrFace, LZ <: CenterOrFace} + +XYZ(x::CenterOrFace, y::CenterOrFace = Center(), z::CenterOrFace = Center()) = VarDims(x, y, z) +XYZ(; x::CenterOrFace = Center(), y::CenterOrFace = Center(), z::CenterOrFace = Center()) = VarDims(x, y, z) """ $SIGNATURES Infer the appropriate `VarDims` from the given `Field`. + +This infers the dimensions of an *externally supplied* field, e.g. one wrapped by an +[`InputSource`](@ref), and returns only `XY` or `XYZ`. It is deliberately not the inverse of the +`Field` constructor: a field created from a [`Top`](@ref) or [`Bottom`](@ref) variable reports `XY`, +because recovering the `Coordinate` would mean inferring intent from the field's indices. Where the +domain or the coordinate matters, read the variable's declared [`VarLocation`](@ref) instead. """ vardims(::AbstractField{LX, LY, Nothing}) where {LX, LY} = XY(LX(), LY()) vardims(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = XYZ(LX(), LY(), LZ()) +""" + $TYPEDEF + +Represents the "location" of an abstract variable, i.e. both the spatial domain and +its dimensionality. + +The `domain` may be `nothing`, which declares a variable as domain-agnostic: it makes no claim about +where the variable lives, so it is compatible with any domain and adopts whichever one another +declaration of the same variable states. This is the default for an [`InputSource`](@ref), which +generally cannot know the domain of the variable it feeds, and it is also the natural choice for a +model discretized on a plain `AbstractGrid`, which has only one discretization. + +On an [`AbstractLandGrid`](@ref) a variable which is still domainless once all declarations are +merged is allocated on the ground domain. That is unambiguous for a 2D variable, since every domain +shares the same horizontal discretization, but a variable with a vertical extent or position also +warns; see `domain_matters`. +""" +struct VarLocation{Dims, Domain} + dims::Dims + domain::Domain +end + +VarLocation(dims::VarDims) = VarLocation(dims, nothing) + +vardims(var::VarLocation) = var.dims +vardomain(var::VarLocation) = var.domain + +Base.summary(::Ground) = "ground" +Base.summary(::Snow) = "snow" +Base.summary(::Canopy) = "canopy" +Base.summary(::Surface) = "surface" +Base.summary(::Atmosphere) = "atmosphere" + +# Variables may be declared without a domain, so the places which name a variable's domain need a +# word for its absence. Defining `Base.summary(::Nothing)` would be type piracy. +@inline domain_summary(domain::VarDomain) = summary(domain) +@inline domain_summary(::Nothing) = "unspecified" + +""" + $SIGNATURES + +Whether two declarations of the same variable agree on its domain. A declaration without a domain +makes no claim and is compatible with any: this is what lets an [`InputSource`](@ref), which +generally cannot know which domain the variable it feeds belongs to, stay domain-agnostic and adopt +whatever the declaring process states. Two *different* stated domains remain a conflict. +""" +@inline domains_compatible(d1::VarDomain, d2::VarDomain) = d1 == d2 +@inline domains_compatible(::Nothing, ::VarDomain) = true +@inline domains_compatible(::VarDomain, ::Nothing) = true +@inline domains_compatible(::Nothing, ::Nothing) = true + +# Aliased constructors for VarLocation on the three domains +Ground(dims::VarDims) = VarLocation(dims, Ground()) +Snow(dims::VarDims) = VarLocation(dims, Snow()) +Canopy(dims::VarDims) = VarLocation(dims, Canopy()) +Surface(dims::XY) = VarLocation(dims, Surface()) +Surface(::XYZ) = error("surface variables must be 2D (XY)") +Atmosphere(dims::XY) = VarLocation(dims, Atmosphere()) +Atmosphere(::XYZ) = error("atmospheric forcing variables must be 2D (XY)") + """ Base type for state variable placeholder types. """ -abstract type AbstractVariable{name, VD, UT} end +abstract type AbstractVariable{name, VL, UT} end """ $SIGNATURES @@ -53,13 +203,30 @@ should return the name of the variable returned by the closure relation. @inline varname(::Type{<:AbstractVariable{name}}) where {name} = name @inline varname(namespace::Pair{Symbol}) = first(namespace) +""" + $SIGNATURES + +Retrieve the [`VarLocation`](@ref) of this variable, i.e. both its grid dimensions and the model +domain it is defined on. +""" +@inline varloc(var::AbstractVariable) = var.loc +@inline varloc(::Type{<:AbstractVariable{name, VL}}) where {name, VL} = VL + """ $SIGNATURES Retrieve the grid dimensions on which this variable is defined. """ -@inline vardims(var::AbstractVariable) = var.dims -@inline vardims(::Type{<:AbstractVariable{name, VD}}) where {name, VD} = VD +@inline vardims(var::AbstractVariable) = vardims(varloc(var)) +@inline vardims(::Type{<:AbstractVariable{name, <:VarLocation{Dims}}}) where {name, Dims} = Dims + +""" + $SIGNATURES + +Retrieve the grid domain on which this variable is defined. +""" +@inline vardomain(var::AbstractVariable) = vardomain(varloc(var)) +@inline vardomain(::Type{<:AbstractVariable{name, <:VarLocation{Dims, Domain}}}) where {name, Dims, Domain} = Domain """ $SIGNATURES @@ -67,33 +234,34 @@ Retrieve the grid dimensions on which this variable is defined. Retrieve the physical units for the given variable. """ @inline varunits(var::AbstractVariable) = var.units -@inline varunits(::Type{<:AbstractVariable{name, VD, UT}}) where {name, VD, UT} = UT +@inline varunits(::Type{<:AbstractVariable{name, VL, UT}}) where {name, VL, UT} = UT -# Test equality between variables by their names, dimensions, and physical units +# Test equality between variables by their names, dimensions, domains, and physical units Base.:(==)(var1::AbstractVariable, var2::AbstractVariable) = varname(var1) == varname(var2) && vardims(var1) == vardims(var2) && + vardomain(var1) == vardomain(var2) && varunits(var1) == varunits(var2) function Base.summary(var::AbstractVariable) unitstr = varunits(var) == NoUnits ? "-" : varunits(var) - text = "$(string(varname(var))) [$(unitstr)] on $(typeof(vardims(var)))" + text = "$(string(varname(var))) [$(unitstr)] on $(typeof(vardims(var))) of the $(domain_summary(vardomain(var))) domain" return text end """ $TYPEDEF -Represents metadata for a generic state variable with the given `name` and spatial `dims`. +Represents metadata for a generic state variable with the given `name` and spatial `loc`. """ -struct Variable{name, VD, UT} <: AbstractVariable{name, VD, UT} - "Variable dimensions" - dims::VD +struct Variable{name, VL, UT} <: AbstractVariable{name, VL, UT} + "Variable location" + loc::VL "Physical units" units::UT - Variable(name::Symbol, dims::VarDims, units::Units = NoUnits) = new{name, typeof(dims), typeof(units)}(dims, units) + Variable(name::Symbol, loc::VarLocation, units::Units = NoUnits) = new{name, typeof(loc), typeof(units)}(loc, units) end """ @@ -113,17 +281,20 @@ abstract type AbstractClosureRelation end """ Baste type for process state variables with specific intents, e.g. `prognostic`, `auxiliary`, or `input`. """ -abstract type AbstractProcessVariable{name, VD, UT} <: AbstractVariable{name, VD, UT} end +abstract type AbstractProcessVariable{name, VL, UT} <: AbstractVariable{name, VL, UT} end +@inline varloc(pv::AbstractProcessVariable) = varloc(pv.var) @inline vardims(pv::AbstractProcessVariable) = vardims(pv.var) +@inline vardomain(pv::AbstractProcessVariable) = vardomain(pv.var) @inline varunits(pv::AbstractProcessVariable) = varunits(pv.var) function Base.show(io::IO, ::MIME"text/plain", var::AbstractVariable) units = varunits(var) + domain = domain_summary(vardomain(var)) return if units != NoUnits - println(io, "$(nameof(typeof(var))) $(varname(var)) with dimensions $(typeof(vardims(var))) and units $(string(varunits(var)))") + println(io, "$(nameof(typeof(var))) $(varname(var)) on the $domain domain with dimensions $(typeof(vardims(var))) and units $(string(varunits(var)))") else - println(io, "$(nameof(typeof(var))) $(varname(var)) with dimensions $(typeof(vardims(var)))") + println(io, "$(nameof(typeof(var))) $(varname(var)) on the $domain domain with dimensions $(typeof(vardims(var)))") end end @@ -136,12 +307,12 @@ indirectly from the values of one or more prognostic variables. """ struct AuxiliaryVariable{ name, - VD <: VarDims, + VL <: VarLocation, UT <: Units, - Var <: Variable{name, VD, UT}, + Var <: Variable{name, VL, UT}, BT <: DomainSets.AbstractInterval, FC, - } <: AbstractProcessVariable{name, VD, UT} + } <: AbstractProcessVariable{name, VL, UT} "State variable" var::Var @@ -163,12 +334,12 @@ Input variables can also be made to vary in time through the use of [`InputSourc """ struct InputVariable{ name, - VD <: VarDims, + VL <: VarLocation, UT <: Units, - Var <: Variable{name, VD, UT}, + Var <: Variable{name, VL, UT}, BT <: DomainSets.AbstractInterval, Def <: Union{Nothing, Number, Function}, - } <: AbstractProcessVariable{name, VD, UT} + } <: AbstractProcessVariable{name, VL, UT} "State variable" var::Var @@ -196,13 +367,13 @@ variable which is used to hold the value of their instantaneous time derivative """ struct PrognosticVariable{ name, - VD <: VarDims, + VL <: VarLocation, UT <: Units, - Var <: Variable{name, VD, UT}, + Var <: Variable{name, VL, UT}, CL <: Union{Nothing, AbstractClosureRelation}, TV <: Union{Nothing, AuxiliaryVariable}, BT <: DomainSets.AbstractInterval, - } <: AbstractProcessVariable{name, VD, UT} + } <: AbstractProcessVariable{name, VL, UT} "State variable" var::Var @@ -293,21 +464,49 @@ with_scope(path::VarPath, var::AbstractVariable) = Variables(obj) = Variables(variables(obj)) Variables(vars::Variables) = vars Variables(vars::Union{AbstractProcessVariable, Namespace}...) = Variables(vars) +""" + $SIGNATURES + +Describe how two declarations of the same variable disagree. Used to report incompatible duplicates +in terms of the attribute which differs rather than by printing both variables and leaving the +reader to spot it. +""" +function describe_conflict(var1::AbstractVariable, var2::AbstractVariable) + differences = String[] + vardims(var1) != vardims(var2) && push!(differences, "dimensions $(typeof(vardims(var1))) vs $(typeof(vardims(var2)))") + !domains_compatible(vardomain(var1), vardomain(var2)) && push!(differences, "domain $(domain_summary(vardomain(var1))) vs $(domain_summary(vardomain(var2)))") + varunits(var1) != varunits(var2) && push!(differences, "units $(varunits(var1)) vs $(varunits(var2))") + return isempty(differences) ? "differing declarations" : join(differences, ", ") +end + function Variables(vars::Tuple{Vararg{Union{AbstractProcessVariable, Namespace}}}) # partition variables into prognostic, auxiliary, input, and namespace groups; # duplicates within each group are automatically merged varmeta(var::AbstractVariable) = (varname(var), vardims(var), varunits(var)) varmeta(ns::Namespace) = varname(ns) - function register!(vardict::AbstractDict, var) - if haskey(vardict, varname(var)) && varmeta(vardict[varname(var)]) != varmeta(var) - error("Found incompatible duplicates of variable $(varname(var)): $(var) $(vars[varname(var)])") - elseif !haskey(vardict, varname(var)) - vardict[varname(var)] = var + # The domain is compared separately from the rest of the metadata because a domainless + # declaration is compatible with any domain rather than equal to it. + compatible(v1, v2) = varmeta(v1) == varmeta(v2) && domains_compatible(vardomain(v1), vardomain(v2)) + function register!(vardict::AbstractDict, var::AbstractVariable) + name = varname(var) + if !haskey(vardict, name) + vardict[name] = var + elseif !compatible(vardict[name], var) + error("Found incompatible duplicates of variable $name: $(describe_conflict(var, vardict[name]))") + elseif isnothing(vardomain(vardict[name])) && !isnothing(vardomain(var)) + # the stated domain wins over the agnostic one + vardict[name] = var else - vardict[varname(var)] = first(merge(vardict[varname(var)], var)) + vardict[name] = first(merge(vardict[name], var)) end return nothing end + # Namespaces carry no domain, so they are merged on their name alone. + function register!(nsdict::AbstractDict, ns::Namespace) + name = varname(ns) + nsdict[name] = haskey(nsdict, name) ? first(merge(nsdict[name], ns)) : ns + return nothing + end # create OrderedDicts for each variable type prognostic_vars = OrderedDict{Symbol, AbstractVariable}() tendency_vars = OrderedDict{Symbol, AuxiliaryVariable}() @@ -488,16 +687,26 @@ end """ $SIGNATURES -Convenience constructor for `Variable`. +Convenience constructor for `Variable`. A variable normally states the domain it lives on by +wrapping its [`VarDims`](@ref) in a [`VarDomain`](@ref), e.g. `var(:temperature, Ground(XYZ()))` or +`var(:snow_temperature, Snow(XY()))`. + +Bare dimensions, e.g. `var(:u, XY())`, declare the variable without a domain. This is for models +discretized on a plain `AbstractGrid`, where there is only one discretization and naming a domain +would say nothing. Prefer stating the domain in any model which runs on an +[`AbstractLandGrid`](@ref): there the domain is a real choice, and a reader of the declaration +should not have to infer it. """ -@inline var(name::Symbol, dims::VarDims, units::Units = NoUnits) = Variable(name, dims, units) +@inline var(name::Symbol, loc::VarLocation, units::Units = NoUnits) = Variable(name, loc, units) + +@inline var(name::Symbol, dims::VarDims, units::Units = NoUnits) = Variable(name, VarLocation(dims), units) """ $SIGNATURES Convenience constructors for `PrognosticVariable`. """ -@inline prognostic(name::Symbol, dims::VarDims; units = NoUnits, closure = nothing, bounds = Unbounded, desc = "") = prognostic(var(name, dims, units); closure, bounds, desc) +@inline prognostic(name::Symbol, loc::Union{VarDims, VarLocation}; units = NoUnits, closure = nothing, bounds = Unbounded, desc = "") = prognostic(var(name, loc, units); closure, bounds, desc) @inline prognostic(var::Variable; closure = nothing, bounds = Unbounded, desc = "") = PrognosticVariable(var, closure, tendency(var), bounds, desc) """ @@ -505,7 +714,7 @@ Convenience constructors for `PrognosticVariable`. Convenience constructor method for `AuxiliaryVariable`. """ -@inline auxiliary(name::Symbol, dims::VarDims, ctor = nothing, params = nothing; units = NoUnits, bounds = Unbounded, desc = "") = auxiliary(var(name, dims, units), ctor, params; bounds, desc) +@inline auxiliary(name::Symbol, loc::Union{VarDims, VarLocation}, ctor = nothing, params = nothing; units = NoUnits, bounds = Unbounded, desc = "") = auxiliary(var(name, loc, units), ctor, params; bounds, desc) @inline auxiliary(var::Variable, ::Nothing, ::Nothing; bounds = Unbounded, desc = "") = AuxiliaryVariable(var, nothing, bounds, desc) @inline auxiliary(var::Variable, ctor::Function, params; bounds = Unbounded, desc = "") = AuxiliaryVariable(var, (_, grid, clock, fields) -> ctor(grid, clock, fields, params), bounds, desc) # `KernelFunction` constructors (from `kernel`) are callable structs, not `Function`s; they define @@ -517,7 +726,7 @@ Convenience constructor method for `AuxiliaryVariable`. Convenience constructor method for `InputVariable`. """ -@inline input(name::Symbol, dims::VarDims; default = nothing, units = NoUnits, bounds = Unbounded, desc = "") = input(var(name, dims, units); default, bounds, desc) +@inline input(name::Symbol, loc::Union{VarDims, VarLocation}; default = nothing, units = NoUnits, bounds = Unbounded, desc = "") = input(var(name, loc, units); default, bounds, desc) @inline input(var::Variable; default = nothing, bounds = Unbounded, desc = "") = InputVariable(var, default, bounds, desc) """ @@ -526,7 +735,7 @@ Convenience constructor method for `InputVariable`. Creates an `AuxiliaryVariable` for the tendency of a prognostic variable with the given name, dimensions, and physical units. This constructor is primarily used internally by other constructors and does not usually need to be called by implementations of `variables`. """ -@inline tendency(var::Variable) = auxiliary(varname(var), vardims(var), units = upreferred(varunits(var)) / u"s") +@inline tendency(var::Variable) = auxiliary(varname(var), varloc(var), units = upreferred(varunits(var)) / u"s") """ $SIGNATURES diff --git a/src/diagnostics/surface_fluxes.jl b/src/diagnostics/surface_fluxes.jl index 8780d2106..fc233d4e8 100644 --- a/src/diagnostics/surface_fluxes.jl +++ b/src/diagnostics/surface_fluxes.jl @@ -12,7 +12,7 @@ function diagnose_skin_temperature_residual( ) return kernel_operation2D(state, model.grid, seb.skin_temperature, model.constants, snow) do i, j, grid, fields, skinT::ImplicitSkinTemperature, args... Ts_implicit = compute_skin_temperature(i, j, grid, fields, skinT, args...) - Ts_prev = state.skin_temperature[i, j, 1] + Ts_prev = state.skin_temperature[i, j, end] return Ts_prev - Ts_implicit end end diff --git a/src/domains.jl b/src/domains.jl new file mode 100644 index 000000000..ce71db79e --- /dev/null +++ b/src/domains.jl @@ -0,0 +1,64 @@ +""" + $TYPEDEF + +Marker type for state variable spatial *domains*; currently, five domains are considered: +`Ground`, `Snow`, `Canopy`, `Surface`, and `Atmosphere`. + +The first three are *vertical* domains, each of which a land grid may discretize in its own right. +`Surface` instead represents the land-atmosphere interface, whose physical position depends on the +surface tile (the top of the soil column over bare ground, the top of the snowpack under snow, the +canopy where there is vegetation). A variable declared on `Surface` therefore never carries a +vertical dimension and must always be 2D (`XY`). + +`Atmosphere` is likewise not a vertical domain; it is the domain of the atmospheric forcings a land +model reads rather than computes. + +A variable need not state a domain at all; see [`VarLocation`](@ref). +""" +abstract type VarDomain end + +""" + $TYPEDEF + +The ground (soil) domain: the vertically resolved subsurface column. This is the domain on which a +land grid's shared horizontal discretization is defined. +""" +struct Ground <: VarDomain end + +""" + $TYPEDEF + +The snow domain, i.e. the snowpack above the ground surface. A grid which does not discretize the +snowpack vertically (a single-layer scheme) still carries its 2D variables. +""" +struct Snow <: VarDomain end + +""" + $TYPEDEF + +The canopy domain, i.e. the vegetation layer. A grid which does not discretize the canopy vertically +(a big-leaf scheme) still carries its 2D variables. +""" +struct Canopy <: VarDomain end + +""" + $TYPEDEF + +The land-atmosphere interface. Unlike [`Ground`](@ref), [`Snow`](@ref), and [`Canopy`](@ref), the +surface is not a vertical domain: its physical position depends on the surface tile, being the top +of the soil column over bare ground, the top of the snowpack under snow, and the canopy where there +is vegetation. `Surface` variables therefore never carry a vertical dimension and must be declared +as `Surface(XY())`; `Surface(XYZ())` raises an error. +""" +struct Surface <: VarDomain end + +""" + $TYPEDEF + +The near-surface atmosphere: the domain of *atmospheric forcing* variables, i.e. quantities a land +model reads rather than computes, such as air temperature, wind, humidity, precipitation, and +downwelling radiation. `Atmosphere` is intended for forcing variables only. Like [`Surface`](@ref), +`Atmosphere` carries no vertical discretization, so its variables must be declared as `Atmosphere(XY())`; +`Atmosphere(XYZ())` raises an error. +""" +struct Atmosphere <: VarDomain end diff --git a/src/grids/column_ring_grid.jl b/src/grids/column_ring_grid.jl index 7f747cf24..e3150f53f 100644 --- a/src/grids/column_ring_grid.jl +++ b/src/grids/column_ring_grid.jl @@ -152,7 +152,7 @@ Converts a `RingGrids.Field` to an Oceananigans `Field` using the given `ColumnRingGrid`. Only masked grid points are copied to the Oceananigans field. For 2D RingGrids fields, returns a 2D Oceananigans field. For 3D fields, returns a 3D field. """ -function Oceananigans.Field(ring_field::RingGrids.AbstractField, grid::ColumnRingGrid; default_value = zero(eltype(ring_field))) +function Oceananigans.Field(ring_field::RingGrids.AbstractField, grid::ColumnRingGrid; domain = nothing) if ndims(ring_field) == 1 # 2D field (horizontal only): treat the data as a single-column matrix so one masked gather # (`data[mask, :]`) serves both the 1D and 2D cases. There's a related Reactant bug that makes this necessary: https://github.com/EnzymeAD/Reactant.jl/issues/3087 @@ -170,7 +170,8 @@ function Oceananigans.Field(ring_field::RingGrids.AbstractField, grid::ColumnRin mask = grid.mask.data # host boolean mask (see note above) gathered = data[mask, :] values = reshape(gathered, size(gathered, 1), 1, size(gathered, 2)) - oceananigans_field = Field(grid, dims) + loc = VarLocation(dims, domain) + oceananigans_field = Field(grid, loc) set!(oceananigans_field, values) return oceananigans_field end @@ -183,7 +184,7 @@ Converts a `RingGrids.Field` whose last dimension is time to an Oceananigans `Fi (horizontal × time) yields a 2D series, a 3D field (horizontal × vertical × time) a 3D series. Additional keyword arguments (e.g. `time_indexing = Cyclical()`) are forwarded to the `FieldTimeSeries` constructor. """ -function Oceananigans.FieldTimeSeries(ring_field::RingGrids.AbstractField, grid::ColumnRingGrid, times::AbstractVector; kwargs...) +function Oceananigans.FieldTimeSeries(ring_field::RingGrids.AbstractField, grid::ColumnRingGrid, times::AbstractVector; domain = nothing, kwargs...) @assert last(size(ring_field)) == length(times) "Last dimension of RingGrids Field must match the length of `times`" arch = architecture(grid) @@ -205,7 +206,8 @@ function Oceananigans.FieldTimeSeries(ring_field::RingGrids.AbstractField, grid: mask = grid.mask.data # host boolean mask (see note above) gathered = data[mask, :, :] values = reshape(gathered, size(gathered, 1), 1, size(gathered)[2:end]...) - oceananigans_fts = FieldTimeSeries(grid, dims, times; kwargs...) + loc = VarLocation(dims, domain) + oceananigans_fts = FieldTimeSeries(grid, loc, times; kwargs...) copyto!(interior(oceananigans_fts), values) return oceananigans_fts end diff --git a/src/grids/grid_utils.jl b/src/grids/grid_utils.jl index edf94ba7a..a35c00788 100644 --- a/src/grids/grid_utils.jl +++ b/src/grids/grid_utils.jl @@ -43,39 +43,72 @@ RingGrids.Architectures.architecture(::CPU) = RingGrids.Architectures.CPU() """ Field( grid::AbstractGrid, - dims::VarDims, + loc::VarLocation, boundary_conditions = nothing, args...; kwargs... ) -Auxiliary constructor for an Oceananigans `Field` on `grid` with the given Terrarium variable `dims` and boundary conditions. -Additional arguments are passed direclty to the `Field` constructor. The location of the `Field` -is determined by `VarDims` defined on `var`. +Auxiliary constructor for an Oceananigans `Field` on `grid` at the given Terrarium variable location +and boundary conditions. Additional arguments are passed directly to the `Field` constructor. -Note that the `Field` is allocated on the discretization of the ground domain rather than on a land -grid itself, so that `field.grid` names the domain the field lives on and Oceananigans' own grid -methods — many of which dispatch on concrete grid types — apply to it unchanged. +[`VarLocation`](@ref) determines both *where* the field is allocated and *what shape* it has: its +[`VarDomain`](@ref) selects the domain discretization via [`variable_grid`](@ref), and its +[`VarDims`](@ref) give the Oceananigans location and, for variables declared at a single point such +as [`Top`](@ref) or [`Bottom`](@ref), the `indices` which restrict the field to that point. + +Note that the `Field` is allocated on the domain discretization rather than on the land grid itself, so +that `field.grid` names the domain the field lives on and Oceananigans' own grid methods, many of +which dispatch on concrete grid types, apply to it unchanged. """ +function Oceananigans.Field( + grid::AbstractLandGrid, + loc::VarLocation, + boundary_conditions = nothing, + args...; + kwargs... + ) + return create_field(variable_grid(grid, loc), loc, boundary_conditions, args...; kwargs...) +end + +# A plain grid has a single discretization, so there is no domain to resolve and `loc`'s domain (if +# any) is ignored; only its dimensions matter. function Oceananigans.Field( grid::AbstractGrid, - dims::VarDims, + loc::VarLocation, + boundary_conditions = nothing, + args...; + kwargs... + ) + return create_field(grid, loc, boundary_conditions, args...; kwargs...) +end + +# Bare dimensions declare a domainless variable; see `var`. +Oceananigans.Field(grid::AbstractGrid, dims::VarDims, args...; kwargs...) = + Field(grid, VarLocation(dims), args...; kwargs...) + +# Allocate the field on an already-resolved `domain` discretization. +function create_field( + domain::AbstractGrid, + loc::VarLocation, boundary_conditions = nothing, args...; kwargs... ) - # infer the location of the Field on the Oceananigans grid from `dims` - loc = location(dims) - FT = Field{map(typeof, loc)...} + dims = vardims(loc) + # infer the location and index restriction of the Field on the Oceananigans grid from `dims` + field_loc = location(dims) + FT = Field{map(typeof, field_loc)...} + field_indices = indices(domain, dims) # Specify BCs if defined field = if isa(boundary_conditions, FieldBoundaryConditions) - FT(ground_domain(grid), args...; boundary_conditions, kwargs...) + FT(domain, args...; indices = field_indices, boundary_conditions, kwargs...) elseif isa(boundary_conditions, NamedTuple) # assume that named tuple corresponds to FieldBoundaryConditions positions - field_bcs = FieldBoundaryConditions(ground_domain(grid), (Center(), Center(), nothing); boundary_conditions...) - FT(ground_domain(grid), args...; boundary_conditions = field_bcs, kwargs...) + field_bcs = FieldBoundaryConditions(domain, (Center(), Center(), nothing); boundary_conditions...) + FT(domain, args...; indices = field_indices, boundary_conditions = field_bcs, kwargs...) else - FT(ground_domain(grid), args...; kwargs...) + FT(domain, args...; indices = field_indices, kwargs...) end return field end @@ -83,22 +116,38 @@ end """ FieldTimeSeries( grid::AbstractGrid, - dims::VarDims, + loc::VarLocation, times=eltype(grid)[]; kwargs... ) -Construct a `FieldTimeSeries` on the given land `grid` with the given `dims` and `times`. Additional -keyword arguments (e.g. `time_indexing = Cyclical()` for a periodically repeating climatology) are -forwarded to the Oceananigans `FieldTimeSeries` constructor. +Construct a `FieldTimeSeries` on the domain discretization of `grid` selected by `loc`, with the +given `times`. Additional keyword arguments (e.g. `time_indexing = Cyclical()` for a periodically +repeating climatology) are forwarded to the Oceananigans `FieldTimeSeries` constructor. """ +function Oceananigans.FieldTimeSeries( + grid::AbstractLandGrid, + loc::VarLocation, + times = eltype(grid)[]; + kwargs... + ) + return create_field_time_series(variable_grid(grid, loc), loc, architecture(grid), times; kwargs...) +end + +# As for `Field`: a plain grid resolves no domain. function Oceananigans.FieldTimeSeries( grid::AbstractGrid, - dims::VarDims, + loc::VarLocation, times = eltype(grid)[]; kwargs... ) - loc = location(dims) - arch = architecture(grid) - return FieldTimeSeries(loc, ground_domain(grid), on_architecture(arch, times); kwargs...) + return create_field_time_series(grid, loc, architecture(grid), times; kwargs...) +end + +Oceananigans.FieldTimeSeries(grid::AbstractGrid, dims::VarDims, times = eltype(grid)[]; kwargs...) = + FieldTimeSeries(grid, VarLocation(dims), times; kwargs...) + +function create_field_time_series(domain::AbstractGrid, loc::VarLocation, arch, times; kwargs...) + dims = vardims(loc) + return FieldTimeSeries(location(dims), domain, on_architecture(arch, times); indices = indices(domain, dims), kwargs...) end diff --git a/src/grids/land_grid.jl b/src/grids/land_grid.jl index 42c17c7b0..d150748bc 100644 --- a/src/grids/land_grid.jl +++ b/src/grids/land_grid.jl @@ -96,18 +96,76 @@ grids are their own ground discretization. $SIGNATURES Return the spatial discretization of the snow domain of `grid`, or `nothing` if the snowpack is not -vertically resolved. +vertically resolved. Grids which are not land grids never resolve a snow domain. """ +@inline snow_domain(::AbstractGrid) = nothing @inline snow_domain(grid::LandGrid) = getfield(grid, :snow) """ $SIGNATURES Return the spatial discretization of the canopy domain of `grid`, or `nothing` if the canopy is not -vertically resolved. +vertically resolved. Grids which are not land grids never resolve a canopy domain. """ +@inline canopy_domain(::AbstractGrid) = nothing @inline canopy_domain(grid::LandGrid) = getfield(grid, :canopy) +""" + $SIGNATURES + +Return the spatial discretization of `grid` for the model domain selected by `domain`, or `nothing` +if `grid` does not resolve that domain vertically. +""" +@inline get_domain(grid::AbstractGrid, ::Ground) = ground_domain(grid) +@inline get_domain(grid::AbstractGrid, ::Snow) = snow_domain(grid) +@inline get_domain(grid::AbstractGrid, ::Canopy) = canopy_domain(grid) +# The surface is an interface, not a vertical domain: it is never discretized in its own right, so +# its (necessarily two-dimensional) variables fall back to the shared horizontal discretization. +@inline get_domain(::AbstractGrid, ::Surface) = nothing +@inline get_domain(::AbstractGrid, ::Atmosphere) = nothing + +""" + $SIGNATURES + +Return the spatial discretization on which a variable at `loc` is allocated. + +A domain which the model does not resolve vertically still has variables: a single-layer snowpack has +a snow water equivalent, a big-leaf canopy has a temperature. Those variables carry no vertical +dimension, so they are allocated on the shared horizontal discretization, which is the ground +domain's. A variable which *is* vertically resolved ([`XYZ`](@ref)) cannot be placed on a domain with +no vertical discretization, and asking for one is a configuration error. + +A grid which is not an [`AbstractLandGrid`](@ref) has a single discretization and no notion of +domains, so the variable's domain is ignored and the grid is returned unchanged. +""" +@inline function variable_grid(grid::AbstractLandGrid, loc::VarLocation) + domain_grid = get_domain(grid, vardomain(loc)) + return isnothing(domain_grid) ? default_domain_grid(grid, vardims(loc), vardomain(loc)) : domain_grid +end + +@inline variable_grid(grid::AbstractGrid, ::VarLocation) = grid + +# A variable declared without a domain falls back to the ground domain. That is unambiguous for a 2D +# variable, but for one with a vertical extent or position it is a real choice being made silently, +# which is how variables end up on the wrong discretization. +function variable_grid(grid::AbstractLandGrid, loc::VarLocation{<:VarDims, Nothing}) + dims = vardims(loc) + !isnothing(dims.z) && @warn "a variable declared without a domain, but with a vertical " * + "position or extent ($(typeof(dims))), is being allocated on the ground domain of this " * + "$(nameof(typeof(grid))). State the domain explicitly, e.g. `Ground(...)`." maxlog = 1 + return ground_domain(grid) +end + +@inline default_domain_grid(grid::AbstractGrid, ::VarDims, ::VarDomain) = ground_domain(grid) + +default_domain_grid(grid::AbstractGrid, dims::XYZ, domain::VarDomain) = throw( + ArgumentError( + "cannot allocate a vertically resolved ($(typeof(dims))) variable on the $(summary(domain)) " * + "domain: this $(nameof(typeof(grid))) does not discretize it vertically. Either declare the " * + "variable on the ground domain, or construct the grid with a $(summary(domain)) discretization." + ) +) + # Properties which are not the land grid's own resolve against the ground domain grid, so that a # ground discretization which is itself a wrapper (e.g. `ColumnRingGrid`) can expose its own # properties (`rings`, `mask`) through the land grid. diff --git a/src/input_output/assets.jl b/src/input_output/assets.jl index 3b8356cea..265cfa7c5 100644 --- a/src/input_output/assets.jl +++ b/src/input_output/assets.jl @@ -41,8 +41,8 @@ struct ERA5LandInvariants <: AbstractLandAsset end artifact_name(::ERA5LandInvariants) = "era5-land-invariants" varnames(::ERA5LandInvariants) = ["cvh", "lsm", "tvl", "cvl", "z", "slt", "dl", "cl", "si10", "tvh"] format(::ERA5LandInvariants) = NetCDF() -indices(::ERA5LandInvariants) = (:, :, 5) native_grid(::ERA5LandInvariants) = RingGrids.FullClenshawGrid(900) +data_indices(::ERA5LandInvariants) = (:, :, 5) """ $TYPEDEF @@ -66,7 +66,7 @@ native_grid(::ERA5LandLeafAreaIndex{N72}) = RingGrids.FullGaussianGrid(72) varnames(::ERA5LandLeafAreaIndex) = ["lai_lv", "lai_hv"] format(::ERA5LandLeafAreaIndex) = NetCDF() -indices(::ERA5LandLeafAreaIndex) = (:, :, :) +data_indices(::ERA5LandLeafAreaIndex) = (:, :, :) """ $TYPEDEF @@ -90,8 +90,8 @@ struct ERA5LandForcings <: AbstractLandAsset end artifact_name(::ERA5LandForcings) = "era5-land-forcings-N72" varnames(::ERA5LandForcings) = ["t2m", "d2m", "tp", "sf", "sp", "ssrd", "strd", "u10", "v10"] format(::ERA5LandForcings) = NetCDF() -indices(::ERA5LandForcings) = (:, :, :) native_grid(::ERA5LandForcings) = RingGrids.FullGaussianGrid(72) +data_indices(::ERA5LandForcings) = (:, :, :) # get_asset @@ -125,7 +125,7 @@ function load_asset(asset::AbstractLandAsset, name::String; NF = Float32, fill_v path = get_asset(asset) fmt = format(asset) grid = native_grid(asset) - return load_asset(path, name, grid, fmt, NF; indices = indices(asset), fill_value) + return load_asset(path, name, grid, fmt, NF; indices = data_indices(asset), fill_value) end """ diff --git a/src/input_output/input_sources.jl b/src/input_output/input_sources.jl index 3a2263588..c2ca3b397 100644 --- a/src/input_output/input_sources.jl +++ b/src/input_output/input_sources.jl @@ -16,6 +16,17 @@ abstract type InputSource{NF, name} end # Default kwarg constructor for convenience InputSource(; kwargs...) = InputSource(kwargs...) +""" + $SIGNATURES + +Retrieve the [`VarLocation`](@ref) of the variable this input source provides, and its constituent +[`VarDims`](@ref) and [`VarDomain`](@ref). The dimensions are inferred from the source data; the +domain is given when the source is constructed and must match the variable it feeds. +""" +@inline varloc(source::InputSource) = source.loc +@inline vardims(source::InputSource) = vardims(varloc(source)) +@inline vardomain(source::InputSource) = vardomain(varloc(source)) + """ $TYPEDSIGNATURES @@ -107,9 +118,9 @@ end Input source that defines `input` state variables with the given names which can then be directly modified by the user. """ -struct FieldInputSource{NF, name, VD <: VarDims, FS <: AnyField{NF}, UT} <: InputSource{NF, name} - "Variable dimensions" - dims::VD +struct FieldInputSource{NF, name, VL <: VarLocation, FS <: AnyField{NF}, UT} <: InputSource{NF, name} + "Variable location: spatial dimensions and the model domain the variable lives on" + loc::VL "Physical units" units::UT @@ -132,8 +143,13 @@ end Create a `FieldInputSource` with the given grid and input variable `fields`. Use it for static input fields. The `name` can either be a plain `Symbol` or a namespaced path; see [`varpath`](@ref). + +The spatial dimensions are inferred from `field`. The `domain` defaults to [`Surface`](@ref), which +is where most forcing data lives; give it explicitly for a source feeding a variable on another +domain. It must match the [`VarDomain`](@ref) that variable is declared with, since a variable +declared on two different domains is a conflict rather than a merge, and is reported as one. """ -function InputSource(grid::AbstractGrid{NF}, field::FS; name, units = NoUnits) where {NF, FS <: AnyField{NF}} +function InputSource(grid::AbstractGrid{NF}, field::FS; name, domain::Optional{VarDomain} = nothing, units = NoUnits) where {NF, FS <: AnyField{NF}} # ensure fields are on the same architecture as the grid field = on_architecture(architecture(grid), field) @@ -141,10 +157,10 @@ function InputSource(grid::AbstractGrid{NF}, field::FS; name, units = NoUnits) w @assert field.grid == ground_domain(grid) "Field must have the same grid as the input grid" # infer the VarDims and subsequently the Field location from the data dimensions - dims = Terrarium.vardims(field) + loc = VarLocation(Terrarium.vardims(field), domain) path = varpath(name) - return FieldInputSource{NF, path, typeof(dims), typeof(field), typeof(units)}(dims, units, field) + return FieldInputSource{NF, path, typeof(loc), typeof(field), typeof(units)}(loc, units, field) end """ @@ -153,17 +169,17 @@ end Convenience function to create a `FieldInputSource` from a `RingGrids.Field`. Converts the RingGrids field to an Oceananigans field and then creates the input source. """ -function InputSource(grid::ColumnRingGrid{NF}, ring_field::RingGrids.AbstractField; name, units = NoUnits) where {NF} +function InputSource(grid::ColumnRingGrid{NF}, ring_field::RingGrids.AbstractField; name, domain::Optional{VarDomain} = nothing, units = NoUnits) where {NF} oceananigans_field = Field(ring_field, grid) - dims = Terrarium.vardims(oceananigans_field) + loc = VarLocation(Terrarium.vardims(oceananigans_field), domain) path = varpath(name) - return FieldInputSource{NF, path, typeof(dims), typeof(oceananigans_field), typeof(units)}(dims, units, oceananigans_field) + return FieldInputSource{NF, path, typeof(loc), typeof(oceananigans_field), typeof(units)}(loc, units, oceananigans_field) end # Land grids delegate to the discretization of the ground domain, which carries the ring grid. InputSource(grid::AbstractLandGrid, ring_field::RingGrids.AbstractField; kwargs...) = InputSource(ground_domain(grid), ring_field; kwargs...) -variables(source::FieldInputSource) = tuple(with_scope(Base.front(varpath(source)), input(varname(source), source.dims; units = source.units))) +variables(source::FieldInputSource) = tuple(with_scope(Base.front(varpath(source)), input(varname(source), source.loc; units = source.units))) """ Type alias for a `FieldTimeSeries` with any X, Y, Z location or grid. @@ -175,9 +191,9 @@ const AnyFieldTimeSeries{NF} = FieldTimeSeries{LX, LY, LZ, TI, K, I, D, G, NF} w Input source that reads input fields from pre-specified Oceananigans `FieldTimeSeries`. """ -struct FieldTimeSeriesInputSource{NF, name, VD <: VarDims, FTS <: AnyFieldTimeSeries{NF}, TT, UT} <: InputSource{NF, name} - "Variable dimensions" - dims::VD +struct FieldTimeSeriesInputSource{NF, name, VL <: VarLocation, FTS <: AnyFieldTimeSeries{NF}, TT, UT} <: InputSource{NF, name} + "Variable location: spatial dimensions and the model domain the variable lives on" + loc::VL "Physical units" units::UT @@ -189,16 +205,16 @@ struct FieldTimeSeriesInputSource{NF, name, VD <: VarDims, FTS <: AnyFieldTimeSe fts::FTS end -function InputSource(fts::AnyFieldTimeSeries{NF}; name, reftime = first(fts.times), units = NoUnits) where {NF} - dims = vardims(fts) +function InputSource(fts::AnyFieldTimeSeries{NF}; name, domain::Optional{VarDomain} = nothing, reftime = first(fts.times), units = NoUnits) where {NF} + loc = VarLocation(vardims(fts), domain) path = varpath(name) - return FieldTimeSeriesInputSource{NF, path, typeof(dims), typeof(fts), typeof(reftime), typeof(units)}(dims, units, reftime, fts) + return FieldTimeSeriesInputSource{NF, path, typeof(loc), typeof(fts), typeof(reftime), typeof(units)}(loc, units, reftime, fts) end # Forward to `InputSource(fts)` to avoid hitting static InputSource(grid, ::AbstractField) constructor InputSource(::AbstractGrid{NF}, fts::AnyFieldTimeSeries{NF}; kwargs...) where {NF} = InputSource(fts; kwargs...) -variables(source::FieldTimeSeriesInputSource) = tuple(with_scope(Base.front(varpath(source)), input(varname(source), source.dims; units = source.units))) +variables(source::FieldTimeSeriesInputSource) = tuple(with_scope(Base.front(varpath(source)), input(varname(source), source.loc; units = source.units))) # to initialize just update the state once at the start time function initialize!(inputs, grid, clock, fields, source::FieldTimeSeriesInputSource) diff --git a/src/models/coupled/land_model.jl b/src/models/coupled/land_model.jl index c71d22d5e..707ad12d4 100644 --- a/src/models/coupled/land_model.jl +++ b/src/models/coupled/land_model.jl @@ -50,6 +50,15 @@ $(TYPEDFIELDS) @component timestepper::Timestepper = default_timestepper(eltype(grid)) end +""" + $SIGNATURES + +Construct a [`LandModel`](@ref) from the given `grid` with the model's multi-domain +[`LandGrid`](@ref) constructed via [`create_land_grid`](@ref). +""" +LandModel(grid::AbstractGrid; kwargs...) = LandModel(create_land_grid(grid); kwargs...) +LandModel(grid::AbstractLandGrid; kwargs...) = LandModel(; grid, kwargs...) + function StateVariables( model::LandModel{NF}; clock = Clock(time = zero(NF)), @@ -97,8 +106,8 @@ function StateVariables( end interface_variables(::LandModel) = ( - auxiliary(:soil_heat_flux, XY(); units = u"W/m^2", desc = "Blended heat flux into the soil top (snow base + bare ground)"), - auxiliary(:snow_surface_heat_flux, XY(); units = u"W/m^2", desc = "Conductive heat flux from the skin into the top of the snowpack (positive upward); drives the snowpack's own energy tendency"), + auxiliary(:soil_heat_flux, Ground(Top()); units = u"W/m^2", desc = "Blended heat flux into the soil top (snow base + bare ground)"), + auxiliary(:snow_surface_heat_flux, Snow(Top()); units = u"W/m^2", desc = "Conductive heat flux from the skin into the top of the snowpack (positive upward); drives the snowpack's own energy tendency"), ) function initialize!(state, model::LandModel) diff --git a/src/models/snow/snow_model.jl b/src/models/snow/snow_model.jl index 49a47182d..2c6c05cdd 100644 --- a/src/models/snow/snow_model.jl +++ b/src/models/snow/snow_model.jl @@ -12,11 +12,11 @@ $(TYPEDFIELDS) """ @parameterized @kwdef struct SnowModel{ NF, - GridType <: AbstractLandGrid{NF}, - Snow <: AbstractSnow{NF}, + GridType <: AbstractGrid, + Snow <: AbstractSnow, Atmosphere <: AbstractAtmosphere, Initializer <: AbstractInitializer, - Timestepper <: AbstractTimeStepper{NF}, + Timestepper <: AbstractTimeStepper, } <: AbstractSnowModel{NF, GridType} "Spatial grid type" grid::GridType @@ -89,6 +89,6 @@ end @kernel inbounds = true function enforce_snow_constraints!(out, grid, snow::AbstractSnow{NF}) where {NF} i, j = @index(Global, NTuple) - out.snow_water_equivalent[i, j, 1] = max(out.snow_water_equivalent[i, j], zero(NF)) - out.snow_energy[i, j, 1] = out.snow_energy[i, j] * (out.snow_water_equivalent[i, j] > zero(NF)) + out.snow_water_equivalent[i, j, end] = max(out.snow_water_equivalent[i, j], zero(NF)) + out.snow_energy[i, j, end] = out.snow_energy[i, j] * (out.snow_water_equivalent[i, j] > zero(NF)) end diff --git a/src/models/soil/soil_model.jl b/src/models/soil/soil_model.jl index 0000e1471..2c7b3ba36 100644 --- a/src/models/soil/soil_model.jl +++ b/src/models/soil/soil_model.jl @@ -8,10 +8,10 @@ $(TYPEDFIELDS) """ @parameterized @kwdef struct SoilModel{ NF, - GridType <: AbstractLandGrid{NF}, - Soil <: AbstractSoil{NF}, + GridType <: AbstractGrid, + Soil <: AbstractSoil, Initializer <: AbstractInitializer, - Timestepper <: AbstractTimeStepper{NF}, + Timestepper <: AbstractTimeStepper, } <: AbstractSoilModel{NF, GridType} "Spatial grid type" grid::GridType diff --git a/src/models/soil/soil_model_bcs.jl b/src/models/soil/soil_model_bcs.jl index 180362e04..2fb35010c 100644 --- a/src/models/soil/soil_model_bcs.jl +++ b/src/models/soil/soil_model_bcs.jl @@ -3,30 +3,30 @@ """ Alias for `FluxBoundaryCondition` on `internal_energy` with name `soil_heat_flux` representing the net heat flux into the top of the soil column. Without snow this equals the surface energy balance `ground_heat_flux`; with snow it is the (blended) conductive flux across the snow base. """ -SoilHeatFlux(value = var(:soil_heat_flux, XY(), u"W/m^2"); kwargs...) = (internal_energy = (top = FluxBoundaryCondition(value; kwargs...),),) +SoilHeatFlux(value = var(:soil_heat_flux, Ground(Top()), u"W/m^2"); kwargs...) = (internal_energy = (top = FluxBoundaryCondition(value; kwargs...),),) """ Alias for `FluxBoundaryCondition` on `internal_energy` with name `geothermal_heat_flux` representing the geothermal heat flux at the bottom boundary of the soil column. """ -GeothermalHeatFlux(value = var(:geothermal_heat_flux, XY(), u"W/m^2"); kwargs...) = (internal_energy = (bottom = FluxBoundaryCondition(value; kwargs...),),) +GeothermalHeatFlux(value = var(:geothermal_heat_flux, Ground(Bottom()), u"W/m^2"); kwargs...) = (internal_energy = (bottom = FluxBoundaryCondition(value; kwargs...),),) """ Alias for `ValueBoundaryCondition` on top `temperature` (in °C) with the given variable name. """ -PrescribedSurfaceTemperature(name::Symbol, value = var(name, XY(), u"°C"); kwargs...) = (temperature = (top = ValueBoundaryCondition(value; kwargs...),),) +PrescribedSurfaceTemperature(name::Symbol, value = var(name, Ground(Top()), u"°C"); kwargs...) = (temperature = (top = ValueBoundaryCondition(value; kwargs...),),) """ Alias for `ValueBoundaryCondition` on top `temperature` (in °C) with the given variable name. """ -PrescribedBottomTemperature(name::Symbol, value = var(name, XY(), u"°C"); kwargs...) = (temperature = (bottom = ValueBoundaryCondition(value; kwargs...),),) +PrescribedBottomTemperature(name::Symbol, value = var(name, Ground(Bottom()), u"°C"); kwargs...) = (temperature = (bottom = ValueBoundaryCondition(value; kwargs...),),) # Hydrology BCs """ Alias for `PrescribedFlux` with name `infiltration` representing liquid water infiltration at the soil surface. """ -InfiltrationFlux(value = var(:infiltration, XY(), u"m/s"); kwargs...) = (saturation_water_ice = (top = FluxBoundaryCondition(value; kwargs...),),) +InfiltrationFlux(value = var(:infiltration, Ground(Top()), u"m/s"); kwargs...) = (saturation_water_ice = (top = FluxBoundaryCondition(value; kwargs...),),) """ Alias for a discrete-form `FluxBoundaryCondition` on `saturation_water_ice` that normalizes the diff --git a/src/models/surface/surface_energy_model.jl b/src/models/surface/surface_energy_model.jl index 54740d6f3..c7a516028 100644 --- a/src/models/surface/surface_energy_model.jl +++ b/src/models/surface/surface_energy_model.jl @@ -9,7 +9,7 @@ conditions. """ @parameterized @kwdef struct SurfaceEnergyModel{ NF, - GridType <: AbstractLandGrid{NF}, + GridType <: AbstractGrid{NF}, SEB <: AbstractSurfaceEnergyBalance, Atmosphere <: AbstractAtmosphere, Initializer <: AbstractInitializer, diff --git a/src/models/surface/surface_hydrology_model.jl b/src/models/surface/surface_hydrology_model.jl index ff17373ea..cdb2b6fcd 100644 --- a/src/models/surface/surface_hydrology_model.jl +++ b/src/models/surface/surface_hydrology_model.jl @@ -8,7 +8,7 @@ $TYPEDFIELDS """ @parameterized @kwdef struct SurfaceHydrologyModel{ NF, - GridType <: AbstractLandGrid{NF}, + GridType <: AbstractGrid{NF}, Atmosphere <: AbstractAtmosphere, CanopyHydrology <: AbstractCanopyInterception, CanopyET <: AbstractEvapotranspiration, diff --git a/src/models/vegetation/vegetation_model.jl b/src/models/vegetation/vegetation_model.jl index 9d650d845..62d3751c8 100644 --- a/src/models/vegetation/vegetation_model.jl +++ b/src/models/vegetation/vegetation_model.jl @@ -12,9 +12,9 @@ $TYPEDFIELDS NF, Vegetation <: AbstractVegetation{NF}, Atmosphere <: AbstractAtmosphere{NF}, - GridType <: AbstractLandGrid{NF}, + GridType <: AbstractGrid, Initializer <: AbstractInitializer, - Timestepper <: AbstractTimeStepper{NF}, + Timestepper <: AbstractTimeStepper, } <: AbstractVegetationModel{NF, GridType} "Spatial grid type" grid::GridType diff --git a/src/processes/atmosphere/prescribed_atmosphere.jl b/src/processes/atmosphere/prescribed_atmosphere.jl index 0e9960e32..c4219ef21 100644 --- a/src/processes/atmosphere/prescribed_atmosphere.jl +++ b/src/processes/atmosphere/prescribed_atmosphere.jl @@ -11,7 +11,7 @@ end Base.nameof(::TracerGas{NF, name}) where {NF, name} = name variables(gas::TracerGas{NF, name}) where {NF, name} = ( - input(name, XY(), default = gas.concentration, units = u"ppm", desc = "Ambient atmospheric $(name) concentration in ppm"), + input(name, Atmosphere(XY()), default = gas.concentration, units = u"ppm", desc = "Ambient atmospheric $(name) concentration in ppm"), ) """ @@ -97,8 +97,8 @@ ParameterEditing.parameters(::PrescribedAtmosphere) = (;) minimum_windspeed(atmos::PrescribedAtmosphere) = atmos.min_windspeed variables(atmos::PrescribedAtmosphere{NF}) where {NF} = ( - input(:air_temperature, XY(), default = NF(10), units = u"°C", desc = "Near-surface air temperature in °C"), - input(:air_pressure, XY(), default = NF(101_325), units = u"Pa", desc = "Atmospheric pressure at the surface in Pa"), + input(:air_temperature, Atmosphere(XY()), default = NF(10), units = u"°C", desc = "Near-surface air temperature in °C"), + input(:air_pressure, Atmosphere(XY()), default = NF(101_325), units = u"Pa", desc = "Atmospheric pressure at the surface in Pa"), variables(atmos.wind)..., variables(atmos.humidity)..., variables(atmos.precip)..., @@ -144,7 +144,7 @@ Retrieve or compute the air pressure at the current time step. Return the current prescribed ambient CO2 concentration level. """ -@propagate_inbounds ambient_co2(i, j, grid, fields, ::PrescribedAtmosphere) = fields.CO2[i, j, 1] +@propagate_inbounds ambient_co2(i, j, grid, fields, ::PrescribedAtmosphere) = fields.CO2[i, j, end] """ air_density(i, j, grid, fields, atmos::AbstractAtmosphere, constants::PhysicalConstants) @@ -167,7 +167,7 @@ Represents a windspeed as direct input/forcing variable. struct Windspeed <: AbstractWind end variables(::Windspeed) = ( - input(:windspeed, XY(), default = 0.1, units = u"m/s", desc = "Wind speed in m/s"), + input(:windspeed, Atmosphere(XY()), default = 0.1, units = u"m/s", desc = "Wind speed in m/s"), ) """ @@ -185,8 +185,8 @@ Represents a windspeed given as `u` (east-west) and `v` (south-north) velocity c struct WindVelocity <: AbstractWind end variables(::WindVelocity) = ( - input(:wind_u, XY(), default = 0.1, units = u"m/s", desc = "Wind velocity u-component in m/s"), - input(:wind_v, XY(), default = 0.1, units = u"m/s", desc = "Wind velocity v-component in m/s"), + input(:wind_u, Atmosphere(XY()), default = 0.1, units = u"m/s", desc = "Wind velocity u-component in m/s"), + input(:wind_v, Atmosphere(XY()), default = 0.1, units = u"m/s", desc = "Wind velocity v-component in m/s"), ) @propagate_inbounds windspeed(i, j, grid, fields, atmos::AbstractAtmosphere{NF, PR, IR, HD, WindVelocity}) where {NF, PR, IR, HD} = max(sqrt(fields.wind_u[i, j]^2 + fields.wind_v[i, j]^2), minimum_windspeed(atmos)) @@ -201,7 +201,7 @@ provided directly as an input field. struct SpecificHumidity <: AbstractHumidity end variables(::SpecificHumidity) = ( - input(:specific_humidity, XY(), default = 1.0e-3, units = u"kg/kg", desc = "Near-surface specific humidity in kg/kg"), + input(:specific_humidity, Atmosphere(XY()), default = 1.0e-3, units = u"kg/kg", desc = "Near-surface specific humidity in kg/kg"), ) """ @@ -233,8 +233,8 @@ are provided as separate input fields. struct RainSnow <: AbstractPrecipitation end variables(::RainSnow) = ( - input(:rainfall, XY(), units = u"m/s", desc = "Liquid precipitation (rainfall) rate"), - input(:snowfall, XY(), units = u"m/s", desc = "Frozen precipitation (snowfall) rate"), + input(:rainfall, Atmosphere(XY()), units = u"m/s", desc = "Liquid precipitation (rainfall) rate"), + input(:snowfall, Atmosphere(XY()), units = u"m/s", desc = "Frozen precipitation (snowfall) rate"), ) """ @@ -261,9 +261,9 @@ length [hr]. struct LongShortWaveRadiation <: AbstractIncomingRadiation end variables(::LongShortWaveRadiation) = ( - input(:surface_shortwave_down, XY(), default = 341, units = u"W/m^2", desc = "Incoming (downwelling) shortwave solar radiation"), - input(:surface_longwave_down, XY(), default = 333, units = u"W/m^2", desc = "Incoming (downwelling) longwave thermal radiation"), - input(:daytime_length, XY(), default = 12, units = u"hr", desc = "Number of daytime hours varying with the season and orbital parameters"), + input(:surface_shortwave_down, Atmosphere(XY()), default = 341, units = u"W/m^2", desc = "Incoming (downwelling) shortwave solar radiation"), + input(:surface_longwave_down, Atmosphere(XY()), default = 333, units = u"W/m^2", desc = "Incoming (downwelling) longwave thermal radiation"), + input(:daytime_length, Atmosphere(XY()), default = 12, units = u"hr", desc = "Number of daytime hours varying with the season and orbital parameters"), ) """ diff --git a/src/processes/snow/energy/snow_energy_closures.jl b/src/processes/snow/energy/snow_energy_closures.jl index 501ff450f..469bf2c74 100644 --- a/src/processes/snow/energy/snow_energy_closures.jl +++ b/src/processes/snow/energy/snow_energy_closures.jl @@ -20,8 +20,8 @@ struct SnowEnergyTemperatureClosure{NF} <: AbstractEnergyClosure end SnowEnergyTemperatureClosure(::Type{NF}) where {NF} = SnowEnergyTemperatureClosure{NF}() variables(::SnowEnergyTemperatureClosure) = ( - auxiliary(:snow_temperature, XY(), units = u"°C", desc = "Depth-averaged snow temperature in °C (≤ 0)"), - auxiliary(:snow_liquid_fraction, XY(), bounds = UnitInterval, desc = "Liquid (unfrozen) fraction of the snow water substance"), + auxiliary(:snow_temperature, Snow(XY()), units = u"°C", desc = "Depth-averaged snow temperature in °C (≤ 0)"), + auxiliary(:snow_liquid_fraction, Snow(XY()), bounds = UnitInterval, desc = "Liquid (unfrozen) fraction of the snow water substance"), ) # Process-level closure entry points (dispatch on the snow process, mirroring the soil interface). @@ -109,14 +109,14 @@ Recover the snow temperature and liquid water fraction from the depth-integrated # Compute liquid water fraction from depth-integrated energy and latent heat of fusion liq = liquid_water_fraction(FreeWater(), Ū_snow, d_snow * ρLθ) # liq will be 1 for W_snow = 0, so gate by W_snow > 0 - out.snow_liquid_fraction[i, j, 1] = liq * (W_snow > 0) + out.snow_liquid_fraction[i, j, end] = liq * (W_snow > 0) # Snow temperature cannot exceed 0°C, so clip the free-water temperature at zero. The energy above # the fully-melted (0°C, all-liquid) reference, i.e. the positive part `U_snow > 0`, is not stored; it # is derived on demand where needed to determine snow melt. C_snow = compute_snow_volumetric_heat_capacity(snow, constants, ρ_snow, liq) U_snow = compute_snow_volumetric_energy(Ū_snow, d_snow, min_snow_conduction_thickness(i, j, grid, fields, snow)) T = energy_to_temperature(FreeWater(), U_snow, ρLθ, C_snow) - out.snow_temperature[i, j, 1] = min(T, zero(T)) + out.snow_temperature[i, j, end] = min(T, zero(T)) return nothing end @@ -139,10 +139,10 @@ Compute the depth-integrated snow energy from a prescribed temperature at grid c d_snow = compute_snow_depth(snow, W_snow, ρ_snow, ρ_w) ρLθ = ρ_snow * L_sl # ρ_snow L_sl = ρ_w θ L_sl by definition liq = zero(eltype(grid)) # always initialize snow fully frozen - out.snow_liquid_fraction[i, j, 1] = liq + out.snow_liquid_fraction[i, j, end] = liq C_snow = compute_snow_volumetric_heat_capacity(snow, constants, ρ_snow, liq) U_snow = T * C_snow - ρLθ * (one(liq) - liq) - out.snow_energy[i, j, 1] = U_snow * d_snow # integrate by d_snow + out.snow_energy[i, j, end] = U_snow * d_snow # integrate by d_snow return nothing end diff --git a/src/processes/snow/snow_interfaces.jl b/src/processes/snow/snow_interfaces.jl index 640211e66..bde5cd217 100644 --- a/src/processes/snow/snow_interfaces.jl +++ b/src/processes/snow/snow_interfaces.jl @@ -128,9 +128,9 @@ for the latent-flux partition). ) f = snow_cover_fraction(i, j, grid, fields, snow) E_subl = compute_snow_sublimation_flux(i, j, grid, fields, snow, atmos, constants, seb.skin_temperature) - out.sublimation[i, j, 1] = f * E_subl - out.basal_heat_flux[i, j, 1] = compute_snow_soil_heat_flux(i, j, grid, fields, snow, constants, soil) - out.surface_heat_flux[i, j, 1] = compute_snow_surface_heat_flux(i, j, grid, fields, snow, constants) + out.sublimation[i, j, end] = f * E_subl + out.basal_heat_flux[i, j, end] = compute_snow_soil_heat_flux(i, j, grid, fields, snow, constants, soil) + out.surface_heat_flux[i, j, end] = compute_snow_surface_heat_flux(i, j, grid, fields, snow, constants) return nothing end diff --git a/src/processes/snow/snow_single_layer.jl b/src/processes/snow/snow_single_layer.jl index 6f3fc293c..b42131eff 100644 --- a/src/processes/snow/snow_single_layer.jl +++ b/src/processes/snow/snow_single_layer.jl @@ -95,13 +95,13 @@ with the bulk density `ρ_snow`. # Process methods variables(snow::SingleLayerSnow) = ( - prognostic(:snow_energy, XY(); closure = get_closure(snow), units = u"J/m^2", desc = "Depth-integrated (column) internal energy of the snowpack relative to water at 0°C"), - prognostic(:snow_water_equivalent, XY(); units = u"m", desc = "Snow water equivalent (ice + retained liquid)"), - auxiliary(:snow_depth, XY(); units = u"m", desc = "Snow layer depth"), - auxiliary(:snow_cover_fraction, XY(); bounds = UnitInterval, desc = "Sub-grid snow-covered area fraction"), - input(:surface_heat_flux, XY(); units = u"W/m^2", desc = "Net heat flux at the snow surface (positive upward)"), - input(:basal_heat_flux, XY(); units = u"W/m^2", desc = "Conductive heat flux at the snow base (positive upward, soil → snow)"), - input(:sublimation, XY(); units = u"m/s", desc = "Sublimation/evaporation rate from the snow surface (SWE)"), + prognostic(:snow_energy, Snow(XY()); closure = get_closure(snow), units = u"J/m^2", desc = "Depth-integrated (column) internal energy of the snowpack relative to water at 0°C"), + prognostic(:snow_water_equivalent, Snow(XY()); units = u"m", desc = "Snow water equivalent (ice + retained liquid)"), + auxiliary(:snow_depth, Snow(XY()); units = u"m", desc = "Snow layer depth"), + auxiliary(:snow_cover_fraction, Snow(XY()); bounds = UnitInterval, desc = "Sub-grid snow-covered area fraction"), + input(:surface_heat_flux, Snow(Top()); units = u"W/m^2", desc = "Net heat flux at the snow surface (positive upward)"), + input(:basal_heat_flux, Snow(Bottom()); units = u"W/m^2", desc = "Conductive heat flux at the snow base (positive upward, soil → snow)"), + input(:sublimation, Snow(Top()); units = u"m/s", desc = "Sublimation/evaporation rate from the snow surface (SWE)"), ) """ @@ -158,8 +158,8 @@ mass and energy balances (see [`compute_snow_water_tendency`](@ref) and [`comput atmos::AbstractAtmosphere, constants::PhysicalConstants ) - tendencies.snow_water_equivalent[i, j, 1] += compute_snow_water_tendency(i, j, grid, fields, snow, atmos) - tendencies.snow_energy[i, j, 1] += compute_snow_energy_tendency(i, j, grid, fields, snow, atmos, constants) + tendencies.snow_water_equivalent[i, j, end] += compute_snow_water_tendency(i, j, grid, fields, snow, atmos) + tendencies.snow_energy[i, j, end] += compute_snow_energy_tendency(i, j, grid, fields, snow, atmos, constants) return nothing end """ @@ -175,8 +175,8 @@ Compute the snow depth, cover fraction, and thermal conductivity at grid cell `i W_snow = fields.snow_water_equivalent[i, j] ρ_w = constants.material.density_water ρ_snow = compute_snow_density(i, j, grid, fields, snow.density) - out.snow_depth[i, j, 1] = compute_snow_depth(snow, W_snow, ρ_snow, ρ_w) - out.snow_cover_fraction[i, j, 1] = compute_snow_cover_fraction(snow.cover, W_snow) + out.snow_depth[i, j, end] = compute_snow_depth(snow, W_snow, ρ_snow, ρ_w) + out.snow_cover_fraction[i, j, end] = compute_snow_cover_fraction(snow.cover, W_snow) return nothing end diff --git a/src/processes/soil/energy/soil_energy.jl b/src/processes/soil/energy/soil_energy.jl index 4458be6bf..538e50071 100644 --- a/src/processes/soil/energy/soil_energy.jl +++ b/src/processes/soil/energy/soil_energy.jl @@ -35,8 +35,8 @@ SoilThermodynamics( Adapt.@adapt_structure SoilThermodynamics 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"), ) # Field constructor for ground_temperature that returns a view of the uppermost soil layer diff --git a/src/processes/soil/energy/soil_energy_closures.jl b/src/processes/soil/energy/soil_energy_closures.jl index 4e6ec33f2..cb0db474b 100644 --- a/src/processes/soil/energy/soil_energy_closures.jl +++ b/src/processes/soil/energy/soil_energy_closures.jl @@ -20,8 +20,8 @@ struct SoilEnergyTemperatureClosure <: AbstractEnergyClosure end Defines `temperature` as the closure variable for `SoilEnergyTemperatureClosure`. """ variables(::SoilEnergyTemperatureClosure) = ( - auxiliary(:temperature, XYZ(), units = u"°C", desc = "Temperature of the soil volume in °C"), - auxiliary(:liquid_water_fraction, XYZ(), bounds = UnitInterval, desc = "Fraction of unfrozen water in the pore space"), + auxiliary(:temperature, Ground(XYZ()), units = u"°C", desc = "Temperature of the soil volume in °C"), + auxiliary(:liquid_water_fraction, Ground(XYZ()), bounds = UnitInterval, desc = "Fraction of unfrozen water in the pore space"), ) function closure!( diff --git a/src/processes/soil/hydrology/soil_hydraulic_closures.jl b/src/processes/soil/hydrology/soil_hydraulic_closures.jl index 971785ef4..b5e93dfdd 100644 --- a/src/processes/soil/hydrology/soil_hydraulic_closures.jl +++ b/src/processes/soil/hydrology/soil_hydraulic_closures.jl @@ -12,7 +12,7 @@ which is here defined in implementations of [`AbstractSoilHydraulics`](@ref). @kwdef struct SoilSaturationPressureClosure <: AbstractSoilWaterClosure end variables(::SoilSaturationPressureClosure) = ( - auxiliary(:pressure_head, XYZ(), units = u"m", desc = "Total hydraulic pressure head in m water displaced at standard pressure"), + auxiliary(:pressure_head, Ground(XYZ()), units = u"m", desc = "Total hydraulic pressure head in m water displaced at standard pressure"), ) """ @@ -95,7 +95,7 @@ end ψz = z - z_ref # compute hydrostatic pressure head assuming impermeable lower boundary # TODO: relax this assumption in the future? - z₀ = fields.water_table[i, j, 1] + z₀ = fields.water_table[i, j, end] ψh = max(0, z₀ - z) # remove hydrostatic and elevation components ψm = ψ - ψh - ψz @@ -129,7 +129,7 @@ end ψz = z - z_ref # compute hydrostatic pressure head assuming impermeable lower boundary # TODO: can we generalize this for arbitrary lower boundaries? - z₀ = fields.water_table[i, j, 1] + z₀ = fields.water_table[i, j, end] ψh = max(0, z₀ - z) # compute total pressure head as sum of ψh + ψm + ψz # note that ψh and ψz will cancel out in the saturated zone diff --git a/src/processes/soil/hydrology/soil_hydrology.jl b/src/processes/soil/hydrology/soil_hydrology.jl index 12e568a43..74e6410b3 100644 --- a/src/processes/soil/hydrology/soil_hydrology.jl +++ b/src/processes/soil/hydrology/soil_hydrology.jl @@ -76,10 +76,10 @@ if not defined for the given configuration. @inline get_closure(hydrology::SoilHydrology) = hydrology.closure variables(hydrology::SoilHydrology{NF}) where {NF} = ( - auxiliary(:saturation_water_ice, XYZ(), bounds = UnitInterval, desc = "Saturation level of water and ice in the pore space"), - auxiliary(:water_table, XY(), units = u"m", desc = "Elevation of the water table in meters"), - auxiliary(:hydraulic_conductivity, XYZ(z = Face()), units = u"m/s", desc = "Hydraulic conductivity of soil volumes in m/s"), - input(:liquid_water_fraction, XYZ(), default = 1, bounds = UnitInterval, desc = "Fraction of unfrozen water in the pore space"), + auxiliary(:saturation_water_ice, Ground(XYZ()), bounds = UnitInterval, desc = "Saturation level of water and ice in the pore space"), + auxiliary(:water_table, Ground(XY()), units = u"m", desc = "Elevation of the water table in meters"), + auxiliary(:hydraulic_conductivity, Ground(XYZ(z = Face())), units = u"m/s", desc = "Hydraulic conductivity of soil volumes in m/s"), + input(:liquid_water_fraction, Ground(XYZ()), default = 1, bounds = UnitInterval, desc = "Fraction of unfrozen water in the pore space"), ) function compute_water_table!(state, grid, hydrology::SoilHydrology) @@ -173,7 +173,7 @@ Kernel function that diagnoses the water table at grid cell `i, j` given the cur @propagate_inbounds function compute_water_table!(water_table, i, j, grid, sat, ::SoilHydrology{NF}) where {NF} zs = znodes(ground_domain(grid), Center(), Center(), Face()) # scan z axis starting from the bottom (index 1) to find first non-saturated grid cell - water_table[i, j, 1] = findfirst_z(i, j, <(one(NF)), zs, sat) + water_table[i, j, end] = findfirst_z(i, j, <(one(NF)), zs, sat) return nothing end @@ -246,7 +246,7 @@ surface into the `surface_excess_water` pool owned by the given surface `runoff` runoff::AbstractSurfaceRunoff ) where {NF} surface_excess = redistribute_saturation_profile!(out.saturation_water_ice, i, j, grid, hydrology) - out.surface_excess_water[i, j, 1] += surface_excess + out.surface_excess_water[i, j, end] += surface_excess return nothing end diff --git a/src/processes/soil/hydrology/soil_hydrology_rre.jl b/src/processes/soil/hydrology/soil_hydrology_rre.jl index c8366dab6..485fe8c4a 100644 --- a/src/processes/soil/hydrology/soil_hydrology_rre.jl +++ b/src/processes/soil/hydrology/soil_hydrology_rre.jl @@ -21,10 +21,10 @@ closure relating saturation and pressure head. @kwdef struct RichardsEq <: AbstractVerticalFlow end variables(hydrology::SoilHydrology{NF, RichardsEq}) where {NF} = ( - prognostic(:saturation_water_ice, XYZ(); closure = get_closure(hydrology), bounds = UnitInterval, desc = "Saturation level of water and ice in the pore space"), - auxiliary(:hydraulic_conductivity, XYZ(z = Face()), units = u"m/s", desc = "Hydraulic conductivity of soil volumes in m/s"), - auxiliary(:water_table, XY(), units = u"m", desc = "Elevation of the water table in meters"), - input(:liquid_water_fraction, XYZ(), default = NF(1), bounds = UnitInterval, desc = "Fraction of unfrozen water in the pore space"), + prognostic(:saturation_water_ice, Ground(XYZ()); closure = get_closure(hydrology), bounds = UnitInterval, desc = "Saturation level of water and ice in the pore space"), + auxiliary(:hydraulic_conductivity, Ground(XYZ(z = Face())), units = u"m/s", desc = "Hydraulic conductivity of soil volumes in m/s"), + auxiliary(:water_table, Ground(XY()), units = u"m", desc = "Elevation of the water table in meters"), + input(:liquid_water_fraction, Ground(XYZ()), default = NF(1), bounds = UnitInterval, desc = "Fraction of unfrozen water in the pore space"), ) # Top-level interface methods diff --git a/src/processes/soil/soil_diffusion_timescales.jl b/src/processes/soil/soil_diffusion_timescales.jl index e4fa615a0..f5dc7b7ef 100644 --- a/src/processes/soil/soil_diffusion_timescales.jl +++ b/src/processes/soil/soil_diffusion_timescales.jl @@ -126,7 +126,7 @@ conductivity (fully dry or frozen) impose no restriction and return `Inf`. z = znode(i, j, k, grid, Center(), Center(), Center()) z_ref = znode(i, j, grid.Nz + 1, grid, Center(), Center(), Face()) ψz = z - z_ref - z₀ = fields.water_table[i, j, 1] + z₀ = fields.water_table[i, j, end] ψh = max(zero(z), z₀ - z) ψm = ψ - ψh - ψz # specific moisture capacity ∂θ/∂ψ (m⁻¹) from the analytic SWRC derivative diff --git a/src/processes/soil/stratigraphy/soil_horizon.jl b/src/processes/soil/stratigraphy/soil_horizon.jl index 67d487603..96eebdf4a 100644 --- a/src/processes/soil/stratigraphy/soil_horizon.jl +++ b/src/processes/soil/stratigraphy/soil_horizon.jl @@ -55,10 +55,10 @@ function PrescribedSoilHorizon(::Type{NF}, name::Symbol; porosity = ConstantSoil end variables(horizon::PrescribedSoilHorizon{NF}) where {NF} = ( - input(:sand_fraction, XY(), default = one(NF), bounds = UnitInterval, desc = "Mass fraction of sand in soil matrix"), - input(:silt_fraction, XY(), bounds = UnitInterval, desc = "Mass fraction of silt in soil matrix"), - input(:clay_fraction, XY(), bounds = UnitInterval, desc = "Mass fraction of clay in soil matrix"), - input(:thickness, XY(), default = horizon.default_thickness, bounds = Nonnegative, desc = "Thickness of soil horizon"), + input(:sand_fraction, Ground(XY()), default = one(NF), bounds = UnitInterval, desc = "Mass fraction of sand in soil matrix"), + input(:silt_fraction, Ground(XY()), bounds = UnitInterval, desc = "Mass fraction of silt in soil matrix"), + input(:clay_fraction, Ground(XY()), bounds = UnitInterval, desc = "Mass fraction of clay in soil matrix"), + input(:thickness, Ground(XY()), default = horizon.default_thickness, bounds = Nonnegative, desc = "Thickness of soil horizon"), ) @inline function soil_texture(i, j, grid, fields, horizon::PrescribedSoilHorizon{NF}) where {NF} diff --git a/src/processes/soil/stratigraphy/soil_texture.jl b/src/processes/soil/stratigraphy/soil_texture.jl index 029855d51..9743100f8 100644 --- a/src/processes/soil/stratigraphy/soil_texture.jl +++ b/src/processes/soil/stratigraphy/soil_texture.jl @@ -92,7 +92,7 @@ SoilTexture(::Type{NF}, ::Val{:clayloam}) where {NF} = SoilTexture(NF, sand = 0. @kernel function normalize_texture_kernel!(sand, silt, clay, default::SoilTexture) i, j = @index(Global, NTuple) total = sand[i, j] + silt[i, j] + clay[i, j] - sand[i, j, 1] = ifelse(isnan(total), default.sand, sand[i, j] / total) - silt[i, j, 1] = ifelse(isnan(total), default.silt, silt[i, j] / total) - clay[i, j, 1] = ifelse(isnan(total), default.clay, clay[i, j] / total) + sand[i, j, end] = ifelse(isnan(total), default.sand, sand[i, j] / total) + silt[i, j, end] = ifelse(isnan(total), default.silt, silt[i, j] / total) + clay[i, j, end] = ifelse(isnan(total), default.clay, clay[i, j] / total) end diff --git a/src/processes/surface/albedo.jl b/src/processes/surface/albedo.jl index 7f29e022f..781e1bac5 100644 --- a/src/processes/surface/albedo.jl +++ b/src/processes/surface/albedo.jl @@ -9,8 +9,8 @@ $TYPEDFIELDS PrescribedAlbedo(::Type{NF}) where {NF} = PrescribedAlbedo{NF}() variables(::PrescribedAlbedo) = ( - input(:albedo, XY(), bounds = UnitInterval, desc = "Surface albedo, i.e. ratio of outgoing to incoming shortwave radiation [-]"), - input(:emissivity, XY(), bounds = UnitInterval, desc = "Surface emissivity, i.e. efficiency of longwave emission [-]"), + input(:albedo, Surface(XY()), bounds = UnitInterval, desc = "Surface albedo, i.e. ratio of outgoing to incoming shortwave radiation [-]"), + input(:emissivity, Surface(XY()), bounds = UnitInterval, desc = "Surface emissivity, i.e. efficiency of longwave emission [-]"), ) """ @@ -66,8 +66,8 @@ end DiagnosticAlbedo(::Type{NF}; kwargs...) where {NF} = DiagnosticAlbedo{NF}(; kwargs...) variables(::DiagnosticAlbedo) = ( - auxiliary(:albedo, XY(), bounds = UnitInterval, desc = "Diagnosed surface albedo [-]"), - auxiliary(:emissivity, XY(), bounds = UnitInterval, desc = "Diagnosed surface emissivity [-]"), + auxiliary(:albedo, Surface(XY()), bounds = UnitInterval, desc = "Diagnosed surface albedo [-]"), + auxiliary(:emissivity, Surface(XY()), bounds = UnitInterval, desc = "Diagnosed surface emissivity [-]"), ) """ @@ -126,8 +126,8 @@ end args... ) α, ϵ = compute_albedo(i, j, grid, fields, albedo, vegetation, snow) - out.albedo[i, j, 1] = α - out.emissivity[i, j, 1] = ϵ + out.albedo[i, j, end] = α + out.emissivity[i, j, end] = ϵ return nothing end diff --git a/src/processes/surface/canopy_interception/canopy_interception.jl b/src/processes/surface/canopy_interception/canopy_interception.jl index 668b876c9..572c17c67 100644 --- a/src/processes/surface/canopy_interception/canopy_interception.jl +++ b/src/processes/surface/canopy_interception/canopy_interception.jl @@ -9,7 +9,7 @@ struct NoCanopyInterception{NF} <: AbstractCanopyInterception{NF} end NoCanopyInterception(::Type{NF}) where {NF} = NoCanopyInterception{NF}() variables(noop::NoCanopyInterception) = ( - auxiliary(:rainfall_ground, XY(), passthrough_rainfall, noop; desc = "Rainfall rate reaching the ground", units = u"m/s"), + auxiliary(:rainfall_ground, Ground(XY()), passthrough_rainfall, noop; desc = "Rainfall rate reaching the ground", units = u"m/s"), ) passthrough_rainfall(grid, clock, fields, ::NoCanopyInterception) = fields.rainfall # assumes existence of rainfall field @@ -74,13 +74,13 @@ end PALADYNCanopyInterception(::Type{NF}; kwargs...) where {NF} = PALADYNCanopyInterception{NF}(; kwargs...) variables(::PALADYNCanopyInterception) = ( - prognostic(:canopy_water, XY(); desc = "Canopy liquid water", units = u"m", bounds = Nonnegative), - auxiliary(:canopy_water_interception, XY(); desc = "Canopy rain interception rate", units = u"m/s"), - auxiliary(:canopy_water_removal, XY(); desc = "Canopy water removal rate", units = u"m/s"), - auxiliary(:saturation_canopy_water, XY(); desc = "Fraction of the canopy saturated with water"), - auxiliary(:rainfall_ground, XY(); desc = "Rainfall rate reaching the ground", units = u"m/s"), - input(:leaf_area_index, XY(); desc = "Leaf Area Index", units = u"m^2/m^2"), - input(:stem_area_index, XY(); desc = "Stem Area Index", units = u"m^2/m^2"), + prognostic(:canopy_water, Canopy(XY()); desc = "Canopy liquid water", units = u"m", bounds = Nonnegative), + auxiliary(:canopy_water_interception, Canopy(XY()); desc = "Canopy rain interception rate", units = u"m/s"), + auxiliary(:canopy_water_removal, Canopy(XY()); desc = "Canopy water removal rate", units = u"m/s"), + auxiliary(:saturation_canopy_water, Canopy(XY()); desc = "Fraction of the canopy saturated with water"), + auxiliary(:rainfall_ground, Ground(Top()); desc = "Rainfall rate reaching the ground", units = u"m/s"), + input(:leaf_area_index, Canopy(XY()); desc = "Leaf Area Index", units = u"m^2/m^2"), + input(:stem_area_index, Canopy(XY()); desc = "Stem Area Index", units = u"m^2/m^2"), ) @propagate_inbounds canopy_water(i, j, grid, fields, ::PALADYNCanopyInterception) = fields.canopy_water[i, j] @@ -214,10 +214,10 @@ end rainfall_ground = compute_precip_ground(canopy_interception, rain, I_can, R_can) # Store results - out.canopy_water_interception[i, j, 1] = I_can - out.canopy_water_removal[i, j, 1] = R_can - out.saturation_canopy_water[i, j, 1] = f_can - out.rainfall_ground[i, j, 1] = rainfall_ground + out.canopy_water_interception[i, j, end] = I_can + out.canopy_water_removal[i, j, end] = R_can + out.saturation_canopy_water[i, j, end] = f_can + out.rainfall_ground[i, j, end] = rainfall_ground return out end @@ -233,7 +233,7 @@ end R_can = fields.canopy_water_removal[i, j] # Compute canopy water tendency - tendencies.canopy_water[i, j, 1] = compute_canopy_water_tendency(canopy_interception, I_can, E_can, R_can) + tendencies.canopy_water[i, j, end] = compute_canopy_water_tendency(canopy_interception, I_can, E_can, R_can) return tendencies end diff --git a/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl b/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl index 909f334be..4cb3e5a31 100644 --- a/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl +++ b/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl @@ -45,9 +45,9 @@ end variables(::BareGroundEvaporation) = ( # Skin-driven vapor conductance β/rₐ (independent of skin temperature; held fixed during the SEB solve) - auxiliary(:ground_evaporation_conductance, XY(), units = u"m/s", desc = "Ground evaporation vapor conductance"), - auxiliary(:evaporation_ground, XY(), units = u"m/s", desc = "Ground evaporation flux in meters liquid water height"), - input(:skin_temperature, XY(), units = u"°C", desc = "Skin temperature of the surface"), + auxiliary(:ground_evaporation_conductance, Ground(XY()), units = u"m/s", desc = "Ground evaporation vapor conductance"), + auxiliary(:evaporation_ground, Ground(XY()), units = u"m/s", desc = "Ground evaporation flux in meters liquid water height"), + input(:skin_temperature, Surface(XY()), units = u"°C", desc = "Skin temperature of the surface"), ) """ $TYPEDSIGNATURES """ @@ -102,7 +102,7 @@ bare-ground `evaporation` scheme. soil::Optional{AbstractSoil} = nothing ) g_gnd = compute_evapotranspiration_conductances(i, j, grid, fields, evaporation, constants, atmos, soil) - out.ground_evaporation_conductance[i, j, 1] = g_gnd + out.ground_evaporation_conductance[i, j, end] = g_gnd return out end @@ -136,7 +136,7 @@ end ρ_w = constants.material.density_water # Scale Qh_gnd by snow-free fraction and and scale by air-water density ratio to get evaporative flux E_gnd E_gnd = (NF(1) - f_snow) * Qh_gnd * ρ_a / ρ_w - out.evaporation_ground[i, j, 1] = E_gnd + out.evaporation_ground[i, j, end] = E_gnd return out end # Kernels diff --git a/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl b/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl index 8101f24e8..384b41b57 100644 --- a/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl +++ b/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl @@ -76,15 +76,15 @@ end variables(::PALADYNCanopyEvapotranspiration{NF}) where {NF} = ( # Skin-driven vapor conductances (independent of skin temperature; held fixed during the SEB solve) - auxiliary(:ground_evaporation_conductance, XY(); units = u"m/s", desc = "Ground evaporation vapor conductance"), - auxiliary(:canopy_evaporation_conductance, XY(); units = u"m/s", desc = "Canopy evaporation vapor conductance"), - auxiliary(:transpiration_conductance, XY(); units = u"m/s", desc = "Transpiration vapor conductance"), + auxiliary(:ground_evaporation_conductance, Ground(XY()); units = u"m/s", desc = "Ground evaporation vapor conductance"), + auxiliary(:canopy_evaporation_conductance, Canopy(XY()); units = u"m/s", desc = "Canopy evaporation vapor conductance"), + auxiliary(:transpiration_conductance, Canopy(XY()); units = u"m/s", desc = "Transpiration vapor conductance"), # Partitioned humidity fluxes (skin-driven terms refreshed from the converged skin temperature by the finalize pass) - auxiliary(:evaporation_canopy, XY(); desc = "Canopy evaporation flux in meters liquid water height", units = u"m/s"), - auxiliary(:evaporation_ground, XY(), units = u"m/s", desc = "Ground evaporation flux in meters liquid water height"), - auxiliary(:transpiration, XY(), units = u"m/s", desc = "Transpiration evaporation flux in meters liquid water height"), - input(:skin_temperature, XY(); units = u"°C", desc = "Skin temperature"), - input(:ground_temperature, XY(); default = NF(1), units = u"°C", desc = "Ground surface temperature"), + auxiliary(:evaporation_canopy, Canopy(XY()); desc = "Canopy evaporation flux in meters liquid water height", units = u"m/s"), + auxiliary(:evaporation_ground, Ground(XY()), units = u"m/s", desc = "Ground evaporation flux in meters liquid water height"), + auxiliary(:transpiration, Canopy(XY()), units = u"m/s", desc = "Transpiration evaporation flux in meters liquid water height"), + input(:skin_temperature, Surface(XY()); units = u"°C", desc = "Skin temperature"), + input(:ground_temperature, Ground(Top(z = Center())); default = NF(1), units = u"°C", desc = "Ground surface temperature"), ) @propagate_inbounds function ground_evapotranspiration_flux(i, j, grid, fields, ::PALADYNCanopyEvapotranspiration) @@ -184,9 +184,9 @@ Compute and store the skin-driven vapor conductances on `grid` for the given sch g_gnd, g_trp, g_can = compute_evapotranspiration_conductances(i, j, grid, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, args...) # Store skin-driven vapor conductances in corresponding output Fields - out.ground_evaporation_conductance[i, j, 1] = g_gnd - out.canopy_evaporation_conductance[i, j, 1] = g_can - out.transpiration_conductance[i, j, 1] = g_trp + out.ground_evaporation_conductance[i, j, end] = g_gnd + out.canopy_evaporation_conductance[i, j, end] = g_can + out.transpiration_conductance[i, j, end] = g_trp return out end @@ -246,9 +246,9 @@ for the given scheme `evapotranspiration` and process dependencies. # Rescale by snow-covered fraction (if applicable) and convert to liquid water flux f_snow = snow_cover_fraction(i, j, grid, fields, snow) f_bare = NF(1) - f_snow - out.evaporation_ground[i, j, 1] = f_bare * Qh_gnd * r - out.transpiration[i, j, 1] = f_bare * Qh_trp * r - out.evaporation_canopy[i, j, 1] = f_bare * Qh_can * r + out.evaporation_ground[i, j, end] = f_bare * Qh_gnd * r + out.transpiration[i, j, end] = f_bare * Qh_trp * r + out.evaporation_canopy[i, j, end] = f_bare * Qh_can * r return out end diff --git a/src/processes/surface/radiative_fluxes.jl b/src/processes/surface/radiative_fluxes.jl index cc9b7c8eb..05d8db9db 100644 --- a/src/processes/surface/radiative_fluxes.jl +++ b/src/processes/surface/radiative_fluxes.jl @@ -17,9 +17,9 @@ PrescribedRadiativeFluxes(::Type{NF}) where {NF} = PrescribedRadiativeFluxes{NF} ## Top-level interface methods variables(::PrescribedRadiativeFluxes) = ( - input(:surface_shortwave_up, XY(), units = u"W/m^2", desc = "Outgoing (upwelling) shortwave radiation"), - input(:surface_longwave_up, XY(), units = u"W/m^2", desc = "Outgoing (upwelling) longwave radiation"), - auxiliary(:surface_net_radiation, XY(), units = u"W/m^2", desc = "Net outgoing (positive up) radiation"), + input(:surface_shortwave_up, Surface(XY()), units = u"W/m^2", desc = "Outgoing (upwelling) shortwave radiation"), + input(:surface_longwave_up, Surface(XY()), units = u"W/m^2", desc = "Outgoing (upwelling) longwave radiation"), + auxiliary(:surface_net_radiation, Surface(XY()), units = u"W/m^2", desc = "Net outgoing (positive up) radiation"), ) """ $TYPEDSIGNATURES """ @@ -61,7 +61,7 @@ Compute net radiation and store in auxiliary fields at a grid point. args... ) # Compute and store net radiation - out.surface_net_radiation[i, j, 1] = compute_surface_net_radiation(i, j, grid, fields, rad, atmos) + out.surface_net_radiation[i, j, end] = compute_surface_net_radiation(i, j, grid, fields, rad, atmos) return out end @@ -129,9 +129,9 @@ end ## Top-level interface methods variables(::DiagnosedRadiativeFluxes) = ( - auxiliary(:surface_shortwave_up, XY(), units = u"W/m^2", desc = "Outgoing (upwelling) shortwave radiation"), - auxiliary(:surface_longwave_up, XY(), units = u"W/m^2", desc = "Outgoing (upwelling) longwave radiation"), - auxiliary(:surface_net_radiation, XY(), units = u"W/m^2", desc = "Net radiation budget"), + auxiliary(:surface_shortwave_up, Surface(XY()), units = u"W/m^2", desc = "Outgoing (upwelling) shortwave radiation"), + auxiliary(:surface_longwave_up, Surface(XY()), units = u"W/m^2", desc = "Outgoing (upwelling) longwave radiation"), + auxiliary(:surface_net_radiation, Surface(XY()), units = u"W/m^2", desc = "Net radiation budget"), ) """ $TYPEDSIGNATURES """ @@ -191,12 +191,12 @@ end # Compute and store outgoing fluxes outgoing_fluxes = compute_surface_upwelling_radiation(i, j, grid, fields, rad, skinT, abd, consts, atmos) - surface_shortwave_up[i, j, 1] = outgoing_fluxes.surface_shortwave_up - surface_longwave_up[i, j, 1] = outgoing_fluxes.surface_longwave_up + surface_shortwave_up[i, j, end] = outgoing_fluxes.surface_shortwave_up + surface_longwave_up[i, j, end] = outgoing_fluxes.surface_longwave_up # Compute and store net radiation fields = merge(fields, (; surface_shortwave_up, surface_longwave_up)) - out.surface_net_radiation[i, j, 1] = compute_surface_net_radiation(i, j, grid, fields, rad, atmos) + out.surface_net_radiation[i, j, end] = compute_surface_net_radiation(i, j, grid, fields, rad, atmos) return out end diff --git a/src/processes/surface/runoff/direct_surface_runoff.jl b/src/processes/surface/runoff/direct_surface_runoff.jl index 4ce39eba2..6f724a858 100644 --- a/src/processes/surface/runoff/direct_surface_runoff.jl +++ b/src/processes/surface/runoff/direct_surface_runoff.jl @@ -64,9 +64,9 @@ end # Top-level interface methods variables(::DirectSurfaceRunoff) = ( - prognostic(:surface_excess_water, XY(), units = u"m", desc = "Excess water at the soil surface in m³/m²"), - auxiliary(:surface_runoff, XY(), units = u"m/s", desc = "Total surface runoff"), - auxiliary(:infiltration, XY(), units = u"m/s", desc = "Infiltration flux"), + prognostic(:surface_excess_water, Ground(Top()), units = u"m", desc = "Excess water at the soil surface in m³/m²"), + auxiliary(:surface_runoff, Ground(Top()), units = u"m/s", desc = "Total surface runoff"), + auxiliary(:infiltration, Ground(Top()), units = u"m/s", desc = "Infiltration flux"), ) @propagate_inbounds surface_excess_water(i, j, grid, fields, ::AbstractSurfaceRunoff) = fields.surface_excess_water[i, j] @@ -117,7 +117,7 @@ surface hydrology tendencies, so that `surface_excess_water += ∂S∂t * Δt` d ) where {NF} S = surface_excess_water(i, j, grid, fields, runoff) D = compute_surface_drainage(runoff, S) - tendencies.surface_excess_water[i, j, 1] = -min(D, S) + tendencies.surface_excess_water[i, j, end] = -min(D, S) return tendencies end @@ -143,15 +143,15 @@ end # First, compute rate of excess water removal (surface drainage) surface_drainage = compute_surface_drainage(runoff, excess_water) # Calculate infiltration - infil = out.infiltration[i, j, 1] = compute_infiltration(runoff, surface_drainage, sat_top, k_unsat) + infil = out.infiltration[i, j, end] = compute_infiltration(runoff, surface_drainage, sat_top, k_unsat) else # Case 2: No excess water -> rainfall is routed directly to infiltration surface_drainage = zero(NF) - infil = out.infiltration[i, j, 1] = compute_infiltration(runoff, influx, sat_top, k_unsat) + infil = out.infiltration[i, j, end] = compute_infiltration(runoff, influx, sat_top, k_unsat) end # Compute surface runoff - out.surface_runoff[i, j, 1] = compute_surface_runoff(runoff, influx, surface_drainage, infil) + out.surface_runoff[i, j, end] = compute_surface_runoff(runoff, influx, surface_drainage, infil) return out end diff --git a/src/processes/surface/skin_temperature.jl b/src/processes/surface/skin_temperature.jl index d89b6007c..eaefa60bc 100644 --- a/src/processes/surface/skin_temperature.jl +++ b/src/processes/surface/skin_temperature.jl @@ -18,8 +18,8 @@ PrescribedSkinTemperature(::Type{NF}; kwargs...) where {NF} = PrescribedSkinTemp ## Top-level interface methods variables(::PrescribedSkinTemperature) = ( - auxiliary(:ground_heat_flux, XY(), units = u"W/m^2", desc = "Ground heat flux"), - input(:skin_temperature, XY(), units = u"°C", desc = "Longwave emission temperature of the land surface in °C"), + auxiliary(:ground_heat_flux, Ground(Top()), units = u"W/m^2", desc = "Ground heat flux"), + input(:skin_temperature, Surface(XY()), units = u"°C", desc = "Longwave emission temperature of the land surface in °C"), ) @inline compute_auxiliary!(state, grid, ::PrescribedSkinTemperature, args...) = nothing @@ -123,9 +123,9 @@ end ## Top-level interface methods variables(::ImplicitSkinTemperature) = ( - prognostic(:skin_temperature, XY(), units = u"°C", desc = "Longwave emission temperature of the land surface in °C"), - auxiliary(:ground_heat_flux, XY(), units = u"W/m^2", desc = "Ground heat flux"), - input(:ground_temperature, XY(), units = u"°C", desc = "Temperature of the uppermost ground or soil grid cell in °C"), + prognostic(:skin_temperature, Surface(XY()), units = u"°C", desc = "Longwave emission temperature of the land surface in °C"), + auxiliary(:ground_heat_flux, Ground(Top()), units = u"W/m^2", desc = "Ground heat flux"), + input(:ground_temperature, Ground(Top(z = Center())), units = u"°C", desc = "Temperature of the uppermost ground or soil grid cell in °C"), ) """ @@ -227,7 +227,7 @@ Per-cell mutating variant used by the fused surface-energy-balance kernel: store into the auxiliary output field `out`. """ @propagate_inbounds function compute_ground_heat_flux!(out, i, j, grid, fields, skinT::AbstractSkinTemperature, seb::AbstractSurfaceEnergyBalance) - out.ground_heat_flux[i, j, 1] = compute_ground_heat_flux(i, j, grid, fields, skinT, seb) + out.ground_heat_flux[i, j, end] = compute_ground_heat_flux(i, j, grid, fields, skinT, seb) return nothing end @@ -241,9 +241,9 @@ special case (`f_snow = 0`) of the snow-aware method below; it is a separate met """ @inline function compute_skin_temperature(i, j, grid, fields, skinT::ImplicitSkinTemperature{NF}, args...) where {NF} # Get inputs - R_net = fields.surface_net_radiation[i, j, 1] - H_s = fields.sensible_heat_flux[i, j, 1] - H_l = fields.latent_heat_flux[i, j, 1] + R_net = fields.surface_net_radiation[i, j, end] + H_s = fields.sensible_heat_flux[i, j, end] + H_l = fields.latent_heat_flux[i, j, end] G₀ = compute_ground_heat_flux_demand(skinT, R_net, H_s, H_l) Tg, κg, Δzg = ground_thermal_interface(i, j, grid, fields, skinT) Ts = Tg - G₀ * Δzg / (2 * κg) @@ -264,9 +264,9 @@ atmosphere-side demanded flux `G` (`= R_net + H_s + H_l`), by equating `G` to th snow::AbstractSnow ) where {NF} # Get inputs - R_net = fields.surface_net_radiation[i, j, 1] - H_s = fields.sensible_heat_flux[i, j, 1] - H_l = fields.latent_heat_flux[i, j, 1] + R_net = fields.surface_net_radiation[i, j, end] + H_s = fields.sensible_heat_flux[i, j, end] + H_l = fields.latent_heat_flux[i, j, end] G₀ = compute_ground_heat_flux_demand(skinT, R_net, H_s, H_l) Tg, κg, Δzg = ground_thermal_interface(i, j, grid, fields, skinT) Tsnow, κsnow, dsnow = snow_thermal_interface(i, j, grid, fields, snow, constants) @@ -300,7 +300,7 @@ atmosphere-side demanded flux `G_demand = R_net(Ts_prev) + H(Ts_prev) + LE(Ts_pr # into `ground_heat_flux`); `snow` partitions the latent flux by snow-covered fraction compute_surface_energy_fluxes!(out, i, j, grid, fields, seb, constants, atmos, hydrology, snow) Ts_implicit = compute_skin_temperature(i, j, grid, fields, skinT, constants, snow) - Ts_prev = out.skin_temperature[i, j, 1] + Ts_prev = out.skin_temperature[i, j, end] return Ts_prev - Ts_implicit end diff --git a/src/processes/surface/turbulent_fluxes.jl b/src/processes/surface/turbulent_fluxes.jl index 2dcc966ba..81ecc1ce3 100644 --- a/src/processes/surface/turbulent_fluxes.jl +++ b/src/processes/surface/turbulent_fluxes.jl @@ -11,8 +11,8 @@ struct PrescribedTurbulentFluxes{NF} <: AbstractTurbulentFluxes{NF} end PrescribedTurbulentFluxes(::Type{NF}) where {NF} = PrescribedTurbulentFluxes{NF}() variables(::PrescribedTurbulentFluxes) = ( - input(:sensible_heat_flux, XY(), units = u"W/m^2", desc = "Sensible heat flux at the surface [W m⁻²]"), - input(:latent_heat_flux, XY(), units = u"W/m^2", desc = "Latent heat flux at the surface [W m⁻²]"), + input(:sensible_heat_flux, Surface(XY()), units = u"W/m^2", desc = "Sensible heat flux at the surface [W m⁻²]"), + input(:latent_heat_flux, Surface(XY()), units = u"W/m^2", desc = "Latent heat flux at the surface [W m⁻²]"), ) # The turbulent fluxes are prescribed input variables, so there is nothing to diagnose. @@ -88,8 +88,8 @@ end ## Top-level interface methods variables(::DiagnosedTurbulentFluxes) = ( - auxiliary(:sensible_heat_flux, XY(), units = u"W/m^2", desc = "Sensible heat flux at the surface [W m⁻²]"), - auxiliary(:latent_heat_flux, XY(), units = u"W/m^2", desc = "Latent heat flux at the surface [W m⁻²]"), + auxiliary(:sensible_heat_flux, Surface(XY()), units = u"W/m^2", desc = "Sensible heat flux at the surface [W m⁻²]"), + auxiliary(:latent_heat_flux, Surface(XY()), units = u"W/m^2", desc = "Latent heat flux at the surface [W m⁻²]"), ) """ $TYPEDSIGNATURES """ @@ -214,9 +214,9 @@ end @kernel function compute_auxiliary_kernel!(out, grid, fields, tur::DiagnosedTurbulentFluxes, args...) i, j = @index(Global, NTuple) # compute sensible heat flux - out.sensible_heat_flux[i, j, 1] = compute_sensible_heat_flux(i, j, grid, fields, tur, args...) + out.sensible_heat_flux[i, j, end] = compute_sensible_heat_flux(i, j, grid, fields, tur, args...) # compute latent heat flux - pass all args to allow dispatch on evtr presence - out.latent_heat_flux[i, j, 1] = compute_latent_heat_flux(i, j, grid, fields, tur, args...) + out.latent_heat_flux[i, j, end] = compute_latent_heat_flux(i, j, grid, fields, tur, args...) end # Per-process mutating variant used by the fused surface-energy-balance kernel. @@ -228,8 +228,8 @@ Compute the turbulent (sensible and latent) heat fluxes from the current skin te them into the auxiliary output fields `out`. """ @propagate_inbounds function compute_turbulent_fluxes!(out, i, j, grid, fields, tur::DiagnosedTurbulentFluxes, skinT, constants, atmos, hydrology, snow) - out.sensible_heat_flux[i, j, 1] = compute_sensible_heat_flux(i, j, grid, fields, tur, skinT, constants, atmos) - out.latent_heat_flux[i, j, 1] = compute_latent_heat_flux(i, j, grid, fields, tur, skinT, constants, atmos, hydrology, snow) + out.sensible_heat_flux[i, j, end] = compute_sensible_heat_flux(i, j, grid, fields, tur, skinT, constants, atmos) + out.latent_heat_flux[i, j, end] = compute_latent_heat_flux(i, j, grid, fields, tur, skinT, constants, atmos, hydrology, snow) return nothing end diff --git a/src/processes/vegetation/dynamics/carbon_dynamics.jl b/src/processes/vegetation/dynamics/carbon_dynamics.jl index 6ce38923a..f923bf018 100644 --- a/src/processes/vegetation/dynamics/carbon_dynamics.jl +++ b/src/processes/vegetation/dynamics/carbon_dynamics.jl @@ -30,9 +30,9 @@ end PALADYNCarbonDynamics(::Type{NF}; kwargs...) where {NF} = PALADYNCarbonDynamics{NF}(; kwargs...) variables(::PALADYNCarbonDynamics) = ( - prognostic(:carbon_vegetation, XY(), units = u"kg/m^2"), # Vegetation carbon pool [kgC/m²] - auxiliary(:balanced_leaf_area_index, XY()), # Balanced Leaf Area Index [m²/m²] - input(:net_primary_production, XY(), units = u"kg/m^2/s"), # Net Primary Production [kgC/m²/s] + prognostic(:carbon_vegetation, Canopy(XY()), units = u"kg/m^2"), # Vegetation carbon pool [kgC/m²] + auxiliary(:balanced_leaf_area_index, Canopy(XY())), # Balanced Leaf Area Index [m²/m²] + input(:net_primary_production, Canopy(XY()), units = u"kg/m^2/s"), # Net Primary Production [kgC/m²/s] ) """ @@ -162,7 +162,7 @@ Mutating wrapper for [`compute_balanced_leaf_area_index`](@ref) that stores the traits::PlantTraits ) # Compute balanced Leaf Area Index - out.balanced_leaf_area_index[i, j, 1] = compute_balanced_leaf_area_index(vegcarbon_dynamics, traits, fields.carbon_vegetation[i, j]) + out.balanced_leaf_area_index[i, j, end] = compute_balanced_leaf_area_index(vegcarbon_dynamics, traits, fields.carbon_vegetation[i, j]) return nothing end @@ -177,7 +177,7 @@ Calls [`compute_veg_carbon_tendency`](@ref) and stores the result in `out`. traits::PlantTraits ) # Compute and store C_veg tendency - tend.carbon_vegetation[i, j, 1] = compute_veg_carbon_tendency(i, j, grid, fields, vegcarbon_dynamics, traits) + tend.carbon_vegetation[i, j, end] = compute_veg_carbon_tendency(i, j, grid, fields, vegcarbon_dynamics, traits) return nothing end diff --git a/src/processes/vegetation/dynamics/vegetation_dynamics.jl b/src/processes/vegetation/dynamics/vegetation_dynamics.jl index 5f3e6ec8e..4501f1259 100644 --- a/src/processes/vegetation/dynamics/vegetation_dynamics.jl +++ b/src/processes/vegetation/dynamics/vegetation_dynamics.jl @@ -25,8 +25,8 @@ end PALADYNVegetationDynamics(::Type{NF}; kwargs...) where {NF} = PALADYNVegetationDynamics{NF}(; kwargs...) variables(::PALADYNVegetationDynamics) = ( - prognostic(:vegetation_area_fraction, XY()), # PFT fractional area coverage [-] - input(:net_primary_production, XY(), units = u"kg/m^2/s"), + prognostic(:vegetation_area_fraction, Canopy(XY())), # PFT fractional area coverage [-] + input(:net_primary_production, Canopy(XY()), units = u"kg/m^2/s"), ) @propagate_inbounds vegetation_area_fraction(i, j, grid, fields, ::PALADYNVegetationDynamics) = fields.vegetation_area_fraction[i, j] @@ -146,7 +146,7 @@ Mutating wrapper for [`compute_ν_tendency`](@ref) that stores the result in `te vegcarbon_dynamics::PALADYNCarbonDynamics, traits::PlantTraits ) - tend.vegetation_area_fraction[i, j, 1] = compute_ν_tendency(i, j, grid, fields, veg_dynamics, vegcarbon_dynamics, traits) + tend.vegetation_area_fraction[i, j, end] = compute_ν_tendency(i, j, grid, fields, veg_dynamics, vegcarbon_dynamics, traits) return tend end diff --git a/src/processes/vegetation/hydraulics/plant_available_water.jl b/src/processes/vegetation/hydraulics/plant_available_water.jl index 6cf302705..c084079a4 100644 --- a/src/processes/vegetation/hydraulics/plant_available_water.jl +++ b/src/processes/vegetation/hydraulics/plant_available_water.jl @@ -20,9 +20,9 @@ $TYPEDFIELDS FieldCapacityLimitedPAW(::Type{NF} = Float32) where {NF} = FieldCapacityLimitedPAW{NF}() variables(paw::FieldCapacityLimitedPAW{NF}) where {NF} = ( - auxiliary(:plant_available_water, XYZ(), desc = "Fraction of soil water available for plant root water uptake"), - auxiliary(:soil_moisture_limiting_factor, XY(), soil_moisture_limiting_factor, paw), # soil moisture limiting factor - input(:root_fraction, XYZ(), desc = "Fraction of roots in each soil layer"), + auxiliary(:plant_available_water, Ground(XYZ()), desc = "Fraction of soil water available for plant root water uptake"), + auxiliary(:soil_moisture_limiting_factor, Ground(XY()), soil_moisture_limiting_factor, paw), # soil moisture limiting factor + input(:root_fraction, Ground(XYZ()), desc = "Fraction of roots in each soil layer"), ) """ diff --git a/src/processes/vegetation/hydraulics/root_distribution.jl b/src/processes/vegetation/hydraulics/root_distribution.jl index b67b3b46c..08376479e 100644 --- a/src/processes/vegetation/hydraulics/root_distribution.jl +++ b/src/processes/vegetation/hydraulics/root_distribution.jl @@ -41,7 +41,7 @@ Compute the continuous density function of the root distirbution as a function o end variables(rootdist::StaticExponentialRootDistribution) = ( - auxiliary(:root_fraction, XYZ(), root_fraction, rootdist), # Static root fraction defined as function + auxiliary(:root_fraction, Ground(XYZ()), root_fraction, rootdist), # Static root fraction defined as function ) """ @@ -52,7 +52,7 @@ Returns a `FunctionField` that lazily computes the static root distribution on a The `FunctionField` takes `(x, z)` arguments, so this is restricted to land grids with a `Flat` lateral (`y`) dimension, i.e. those built on column discretizations. """ -function root_fraction(grid::AbstractLandGrid{<:Any, <:Any, Flat}, clock, fields, rootdist::StaticExponentialRootDistribution{NF}) where {NF} +function root_fraction(grid::AbstractGrid{<:Any, <:Any, Flat}, clock, fields, rootdist::StaticExponentialRootDistribution{NF}) where {NF} ground_grid = ground_domain(grid) # define pdf of root distribution as a continuous function of depth ∂R∂z = FunctionField{Center, Center, Center}(ground_grid, parameters = rootdist) do x, z, params diff --git a/src/processes/vegetation/phenology/paladyn_phenology.jl b/src/processes/vegetation/phenology/paladyn_phenology.jl index ce58a20f0..da5b92b18 100644 --- a/src/processes/vegetation/phenology/paladyn_phenology.jl +++ b/src/processes/vegetation/phenology/paladyn_phenology.jl @@ -54,9 +54,9 @@ end PALADYNPhenology(::Type{NF}; kwargs...) where {NF} = PALADYNPhenology{NF}(; kwargs...) variables(::PALADYNPhenology) = ( - prognostic(:growing_degree_days, XY(), units = u"K*d"), # Growing degree days [K⋅day] - auxiliary(:phenology_factor, XY()), # Phenology factor [-] - auxiliary(:leaf_area_index, XY()), # Leaf Area Index [m²/m²] + prognostic(:growing_degree_days, Canopy(XY()), units = u"K*d"), # Growing degree days [K⋅day] + auxiliary(:phenology_factor, Canopy(XY())), # Phenology factor [-] + auxiliary(:leaf_area_index, Canopy(XY())), # Leaf Area Index [m²/m²] ) """ @@ -157,8 +157,8 @@ Mutating wrapper for [`compute_phenology`](@ref) that stores the result in `out` """ @propagate_inbounds function compute_phenology!(out, i, j, grid, fields, phenol::PALADYNPhenology, atmos::AbstractAtmosphere) ϕ, LAI = compute_phenology(i, j, grid, fields, phenol, atmos) - out.phenology_factor[i, j, 1] = ϕ - out.leaf_area_index[i, j, 1] = LAI + out.phenology_factor[i, j, end] = ϕ + out.leaf_area_index[i, j, end] = LAI return out end @@ -170,7 +170,7 @@ Mutating wrapper for [`compute_gdd_tendency`](@ref) that stores the growing-degr @propagate_inbounds function compute_gdd_tendency!(tend, i, j, grid, fields, phenol::PALADYNPhenology, atmos::AbstractAtmosphere) gdd = fields.growing_degree_days[i, j] T_air = air_temperature(i, j, grid, fields, atmos) - tend.growing_degree_days[i, j, 1] = compute_gdd_tendency(phenol, gdd, T_air) + tend.growing_degree_days[i, j, end] = compute_gdd_tendency(phenol, gdd, T_air) return tend end diff --git a/src/processes/vegetation/phenology/prescribed_phenology.jl b/src/processes/vegetation/phenology/prescribed_phenology.jl index d778021e7..8e47fb4b6 100644 --- a/src/processes/vegetation/phenology/prescribed_phenology.jl +++ b/src/processes/vegetation/phenology/prescribed_phenology.jl @@ -11,12 +11,12 @@ $TYPEDFIELDS PrescribedPhenology(::Type{NF}) where {NF} = PrescribedPhenology{NF}() variables(::PrescribedPhenology) = ( - input(:leaf_area_index, XY()), # Leaf Area Index [m²/m²] + input(:leaf_area_index, Canopy(XY())), # Leaf Area Index [m²/m²] ) # if PlantTraits are given, also compute phenology factor variables(phenol::PrescribedPhenology, traits::PlantTraits) = ( - auxiliary(:phenology_factor, XY(), kernel(phenology_factor, phenol, traits)), + auxiliary(:phenology_factor, Canopy(XY()), kernel(phenology_factor, phenol, traits)), variables(phenol)..., ) diff --git a/src/processes/vegetation/photosynthesis/lue_photosynthesis.jl b/src/processes/vegetation/photosynthesis/lue_photosynthesis.jl index 8fe78e327..17b46d2ec 100644 --- a/src/processes/vegetation/photosynthesis/lue_photosynthesis.jl +++ b/src/processes/vegetation/photosynthesis/lue_photosynthesis.jl @@ -65,11 +65,11 @@ end LUEPhotosynthesis(::Type{NF}; kwargs...) where {NF} = LUEPhotosynthesis{NF}(; kwargs...) variables(::LUEPhotosynthesis{NF}) where {NF} = ( - auxiliary(:net_assimilation, XY(), units = u"g/m^2/s"), # Net photosynthesis rate [gC/m²/s] - auxiliary(:leaf_respiration, XY(), units = u"g/m^2/s"), # Leaf respiration rate [gC/m²/s] - auxiliary(:gross_primary_production, XY(), units = u"kg/m^2/s"), # Gross primary production rate [kgC/m²/s] - input(:soil_moisture_limiting_factor, XY(), default = NF(1)), # soil moisture limiting factor with default value of 1 - input(:leaf_area_index, XY()), # Leaf Area Index [m²/m²] + auxiliary(:net_assimilation, Canopy(XY()), units = u"g/m^2/s"), # Net photosynthesis rate [gC/m²/s] + auxiliary(:leaf_respiration, Canopy(XY()), units = u"g/m^2/s"), # Leaf respiration rate [gC/m²/s] + auxiliary(:gross_primary_production, Canopy(XY()), units = u"kg/m^2/s"), # Gross primary production rate [kgC/m²/s] + input(:soil_moisture_limiting_factor, Ground(XY()), default = NF(1)), # soil moisture limiting factor with default value of 1 + input(:leaf_area_index, Canopy(XY())), # Leaf Area Index [m²/m²] ) """ @@ -427,9 +427,9 @@ Calls [`compute_photosynthesis`](@ref) and stores the results in `out`. atmos::AbstractAtmosphere ) Rd, An, GPP = compute_photosynthesis(i, j, grid, fields, photo, stomcond, traits, constants, atmos) - out.leaf_respiration[i, j, 1] = Rd - out.net_assimilation[i, j, 1] = An - out.gross_primary_production[i, j, 1] = GPP + out.leaf_respiration[i, j, end] = Rd + out.net_assimilation[i, j, end] = An + out.gross_primary_production[i, j, end] = GPP return out end diff --git a/src/processes/vegetation/respiration/autotrophic_respiration.jl b/src/processes/vegetation/respiration/autotrophic_respiration.jl index d35bd1e5e..02719fe3a 100644 --- a/src/processes/vegetation/respiration/autotrophic_respiration.jl +++ b/src/processes/vegetation/respiration/autotrophic_respiration.jl @@ -30,11 +30,11 @@ end PALADYNAutotrophicRespiration(::Type{NF}; kwargs...) where {NF} = PALADYNAutotrophicRespiration{NF}(; kwargs...) variables(::PALADYNAutotrophicRespiration) = ( - auxiliary(:autotrophic_respiration, XY(), units = u"kg/m^2/s"), # Autotrophic respiration [kgC/m²/s] - auxiliary(:net_primary_production, XY(), units = u"kg/m^2/s"), # Net Primary Production [kgC/m²/s] - input(:gross_primary_production, XY(), units = u"kg/m^2/s"), # Gross Primary Production [kgC/m²/s] - input(:daily_leaf_respiration, XY(), units = u"g/m^2/s"), # Daily leaf respiration [gC/m²/s] - input(:ground_temperature, XY(), default = 10.0, units = u"°C"), # Ground surface temperature [°C] + auxiliary(:autotrophic_respiration, Canopy(XY()), units = u"kg/m^2/s"), # Autotrophic respiration [kgC/m²/s] + auxiliary(:net_primary_production, Canopy(XY()), units = u"kg/m^2/s"), # Net Primary Production [kgC/m²/s] + input(:gross_primary_production, Canopy(XY()), units = u"kg/m^2/s"), # Gross Primary Production [kgC/m²/s] + input(:daily_leaf_respiration, Canopy(XY()), units = u"g/m^2/s"), # Daily leaf respiration [gC/m²/s] + input(:ground_temperature, Ground(Top(z = Center())), default = 10.0, units = u"°C"), # Ground surface temperature [°C] ) """ @@ -201,8 +201,8 @@ Mutating wrapper for [`compute_autotrophic_respiration`](@ref) that stores the r @propagate_inbounds function compute_autotrophic_respiration!(out, i, j, grid, fields, autoresp::AbstractAutotrophicRespiration, args...) # Compute and store results Ra, NPP = compute_autotrophic_respiration(i, j, grid, fields, autoresp, args...) - out.autotrophic_respiration[i, j, 1] = Ra - out.net_primary_production[i, j, 1] = NPP + out.autotrophic_respiration[i, j, end] = Ra + out.net_primary_production[i, j, end] = NPP return out end diff --git a/src/processes/vegetation/stomatal_conductance/medlyn_stomatal_conductance.jl b/src/processes/vegetation/stomatal_conductance/medlyn_stomatal_conductance.jl index 3475705d8..62424c36b 100644 --- a/src/processes/vegetation/stomatal_conductance/medlyn_stomatal_conductance.jl +++ b/src/processes/vegetation/stomatal_conductance/medlyn_stomatal_conductance.jl @@ -36,10 +36,10 @@ end MedlynStomatalConductance(::Type{NF}; kwargs...) where {NF} = MedlynStomatalConductance{NF}(; kwargs...) variables(::MedlynStomatalConductance) = ( - auxiliary(:canopy_water_conductance, XY(), units = u"m/s"), # Canopy conducatance for water vapor - input(:leaf_area_index, XY(), units = u"m^2/m^2"), - input(:net_assimilation, XY(), units = u"g/m^2/s"), # Net photosynthesis rate [gC/m²/s] - input(:soil_moisture_limiting_factor, XY()), + auxiliary(:canopy_water_conductance, Canopy(XY()), units = u"m/s"), # Canopy conducatance for water vapor + input(:leaf_area_index, Canopy(XY()), units = u"m^2/m^2"), + input(:net_assimilation, Canopy(XY()), units = u"g/m^2/s"), # Net photosynthesis rate [gC/m²/s] + input(:soil_moisture_limiting_factor, Ground(XY())), ) @inline @propagate_inbounds stomatal_conductance(i, j, grid, fields, ::MedlynStomatalConductance) = fields.canopy_water_conductance[i, j] @@ -178,7 +178,7 @@ Calls [`compute_stomatal_conductance`](@ref) and stores the result in `out`. args... ) where {NF} g_stm = compute_stomatal_conductance(i, j, grid, fields, stomcond, traits, constants, atmos, args...) - out.canopy_water_conductance[i, j, 1] = g_stm + out.canopy_water_conductance[i, j, end] = g_stm return out end diff --git a/src/processes/vegetation/vegetation_carbon_cycle.jl b/src/processes/vegetation/vegetation_carbon_cycle.jl index 960d5ee5c..f5998444f 100644 --- a/src/processes/vegetation/vegetation_carbon_cycle.jl +++ b/src/processes/vegetation/vegetation_carbon_cycle.jl @@ -137,7 +137,7 @@ end @propagate_inbounds function vegetation_area_fraction(i, j, grid, fields, veg::VegetationCarbonCycle) if isnothing(veg.vegetation_dynamics) - LAI = fields.leaf_area_index[i, j, 1] + LAI = fields.leaf_area_index[i, j, end] LAI_max = maximum_leaf_area_index(i, j, grid, fields, veg.traits) f_veg = LAI / LAI_max return f_veg diff --git a/src/state_variables.jl b/src/state_variables.jl index 9fe3587eb..4b197c0b4 100644 --- a/src/state_variables.jl +++ b/src/state_variables.jl @@ -320,8 +320,9 @@ Initialize a `StateVariables` data structure containing `Field`s defined on the for all variables defined by `process`. Any predefined `boundary_conditions` and `fields` will be passed through to `initialize` for each variable. -The `grid` may be either a land grid or an ordinary spatial discretization, which is converted to -a land grid via [`create_land_grid`](@ref). +The `grid` may be either a land grid or an ordinary spatial discretization; fields are allocated on +whichever is given. [`AbstractLandGrid`](@ref)s resolve variables' [`VarDomain`](@ref)s to a +their respective vertical discretization; on an ordinary grid the domain is ignored. """ function StateVariables( process::AbstractProcess{NF}, @@ -350,8 +351,9 @@ for all variables in `vars`. Any predefined `boundary_conditions` and `fields` w through to `initialize` for each variable. The `timestepper`'s cache is allocated via `initialize(timestepper, state, progvars)`. -The `grid` may be either a land grid or an ordinary spatial discretization, which is converted to -a land grid via [`create_land_grid`](@ref); state variables always live on a land grid. +The `grid` may be either a land grid or an ordinary spatial discretization; fields are allocated on +whichever is given. Only an [`AbstractLandGrid`](@ref) resolves a variable's [`VarDomain`](@ref) to a +per-domain discretization; on an ordinary grid the domain is ignored. """ function StateVariables( vars::Variables, @@ -363,20 +365,17 @@ function StateVariables( initializers = (;), fields = (;) ) where {NF} - # State variables always live on a land grid, so an ordinary spatial discretization is converted - # to one here exactly as the model constructors do; this is a no-op for a land grid. - land_grid = create_land_grid(grid) # Initialize Fields for each variable group, if they are not already given in the user defined `fields`. fields_dict = OrderedDict{Symbol, AbstractField}(pairs(fields)) - input_fields_dict = initialize(vars.inputs, land_grid, clock, fields_dict, boundary_conditions) - tendency_fields_dict = initialize(vars.tendencies, land_grid, clock, fields_dict, boundary_conditions) - prognostic_fields_dict = initialize(vars.prognostic, land_grid, clock, merge(fields_dict, input_fields_dict), boundary_conditions) - auxiliary_fields_dict = initialize(vars.auxiliary, land_grid, clock, merge(fields_dict, input_fields_dict, prognostic_fields_dict), boundary_conditions) + input_fields_dict = initialize(vars.inputs, grid, clock, fields_dict, boundary_conditions) + tendency_fields_dict = initialize(vars.tendencies, grid, clock, fields_dict, boundary_conditions) + prognostic_fields_dict = initialize(vars.prognostic, grid, clock, merge(fields_dict, input_fields_dict), boundary_conditions) + auxiliary_fields_dict = initialize(vars.auxiliary, grid, clock, merge(fields_dict, input_fields_dict, prognostic_fields_dict), boundary_conditions) # recursively initialize state variables for each namespace namespaces = map(values(vars.namespaces)) do ns ns_bcs = get(boundary_conditions, varname(ns), (;)) ns_fields = get(fields, varname(ns), (;)) - varname(ns) => StateVariables(variables(ns), land_grid; clock, boundary_conditions = ns_bcs, fields = ns_fields) + varname(ns) => StateVariables(variables(ns), grid; clock, boundary_conditions = ns_bcs, fields = ns_fields) end # get closure variable names closurenames = map(varname, closure_variables(values(vars.prognostic))) @@ -427,7 +426,7 @@ for each variable. """ function initialize( vars::OrderedDict{Symbol, <:AbstractVariable}, - grid::AbstractLandGrid, + grid::AbstractGrid, clock::Clock, fields::OrderedDict{Symbol, AbstractField}, boundary_conditions::NamedTuple, @@ -448,7 +447,7 @@ end # Convenience dispatch that accepts `fields` as a NamedTuple and converts to OrderedDict initialize( var::AbstractVariable, - grid::AbstractLandGrid, + grid::AbstractGrid, clock::Clock, fields::NamedTuple, boundary_conditions::NamedTuple, @@ -464,7 +463,7 @@ Otherwise, the new `Field` is constructed using the given `boundary_conditions`. """ function initialize( var::AbstractVariable, - grid::AbstractLandGrid, + grid::AbstractGrid, clock::Clock, fields::OrderedDict{Symbol, AbstractField}, boundary_conditions::NamedTuple, @@ -474,7 +473,7 @@ function initialize( return fields[name] else bcs = get(boundary_conditions, name, nothing) - field = Field(grid, vardims(var), bcs) + field = Field(grid, varloc(var), bcs) # if field is an input variable and has a default value/initializer, call set! on it if isa(var, InputVariable) && !isnothing(var.default) set!(field, var.default) @@ -491,7 +490,7 @@ Initialize a `Field` on `grid` for the given [`AuxiliaryVariable`](@ref). """ function initialize( var::AuxiliaryVariable, - grid::AbstractLandGrid, + grid::AbstractGrid, clock::Clock, fields::OrderedDict{Symbol, AbstractField}, boundary_conditions::NamedTuple @@ -502,7 +501,7 @@ function initialize( elseif isnothing(var.ctor) # retrieve boundary condition (if any) and create Field bcs = get(boundary_conditions, name, nothing) - return Field(grid, vardims(var), bcs) + return Field(grid, varloc(var), bcs) else # invoke field constructor if specified return var.ctor(var, grid, clock, NamedTuple(fields)) diff --git a/src/timesteppers/abstract_timestepper.jl b/src/timesteppers/abstract_timestepper.jl index f5eab8a17..ef13df930 100644 --- a/src/timesteppers/abstract_timestepper.jl +++ b/src/timesteppers/abstract_timestepper.jl @@ -182,6 +182,32 @@ function explicit_step!( return nothing end +""" +Alias for a `Field` restricted to a single vertical index, i.e. one declared at the [`Top`](@ref) or +[`Bottom`](@ref) of a domain. Such a field carries a `Center` or `Face` vertical location, so it is +not distinguishable from a fully resolved (`XYZ`) field by its location parameters alone; the +`indices` type parameter is what separates the two. +""" +const VerticallySlicedField{LX, LY, LZ} = Field{LX, LY, LZ, <:Any, <:Any, Tuple{Colon, Colon, UnitRange{Int}}} + +# A vertically sliced field occupies exactly one `k`, so it is stepped by a 2D kernel which resolves +# that index with `end`. Launching the 3D kernel over the full column would index outside the slice. +function explicit_step!( + field::VerticallySlicedField{LX, LY, LZ}, + tendency::VerticallySlicedField{LX, LY, LZ}, + grid::AbstractGrid{NF}, + timestepper::AbstractTimeStepper, + Δt, + args... + ) where {LX, LY, LZ, NF} + Δt = convert_dt(NF, Δt) + launch!( + grid, XY, explicit_step_slab_kernel!, + field, tendency, timestepper, Δt, args... + ) + return nothing +end + function explicit_step!( field::AbstractField{LX, LY, Nothing}, tendency::AbstractField{LX, LY, Nothing}, @@ -223,7 +249,22 @@ end u = field ∂u∂t = tendency @inbounds let Δt = convert(eltype(tendency), Δt) - u[i, j, 1] += ∂u∂t[i, j] * Δt + u[i, j, end] += ∂u∂t[i, j] * Δt + end +end + +@kernel function explicit_step_slab_kernel!( + field, + grid, + tendency, + ::AbstractTimeStepper, + Δt + ) + i, j = @index(Global, NTuple) + u = field + ∂u∂t = tendency + @inbounds let Δt = convert(eltype(tendency), Δt) + u[i, j, end] += ∂u∂t[i, j, end] * Δt end end diff --git a/src/timesteppers/model_integrator.jl b/src/timesteppers/model_integrator.jl index 44a736738..870b6b6c2 100644 --- a/src/timesteppers/model_integrator.jl +++ b/src/timesteppers/model_integrator.jl @@ -10,9 +10,9 @@ treated as a "model" in `Oceananigans` `Simulation`s and output reading/writing struct ModelIntegrator{ NF, Arch <: AbstractArchitecture, - Grid <: AbstractLandGrid{NF}, - TimeStepper <: AbstractTimeStepper{NF}, - Model <: AbstractModel{NF, Grid}, + Grid <: AbstractGrid, + TimeStepper <: AbstractTimeStepper, + Model <: AbstractModel, StateVars <: AbstractStateVariables, ClockType <: Clock, Inits <: NamedTuple, diff --git a/src/utils/kernel_utils.jl b/src/utils/kernel_utils.jl index 6e19e8a25..6c5694cd9 100644 --- a/src/utils/kernel_utils.jl +++ b/src/utils/kernel_utils.jl @@ -79,7 +79,7 @@ This is intended to be used as a constructor for `AuxiliaryVariable`s: ```julia myvar(i, j, k, grid, fields) = clamp(fields.x[i, j, k], zero(eltype(grid)), one(eltype(grid))) -auxvar = auxiliary(:myvar, XYZ(), kernel(myvar)) +auxvar = auxiliary(:myvar, Ground(XYZ()), kernel(myvar)) ``` """ function kernel(func, args...; clock = false) diff --git a/test/abstract_variables.jl b/test/abstract_variables.jl index 513b59fda..157def8f3 100644 --- a/test/abstract_variables.jl +++ b/test/abstract_variables.jl @@ -296,3 +296,39 @@ end @test Terrarium.varname(ns2) == :boundary @test Terrarium.varname(first(ns2.vars)) == :temperature end + +@testset "Variable domains" begin + using Terrarium: Ground, Snow, Surface, Atmosphere, Top, vardomain + + # A declaration may name a domain or omit it entirely. + @test vardomain(Terrarium.var(:x, Ground(XYZ()))) == Ground() + @test isnothing(vardomain(Terrarium.var(:x, XYZ()))) + + # `Surface` and `Atmosphere` have no vertical discretization, so they are 2D only. + @test_throws ErrorException Surface(XYZ()) + @test_throws ErrorException Atmosphere(XYZ()) + @test vardomain(Terrarium.var(:x, Atmosphere(XY()))) == Atmosphere() + + # A domainless declaration makes no claim, so it unifies with a stated domain rather than + # conflicting with it. This is what lets an `InputSource` stay domain-agnostic. + vars = Variables(input(:x, XY(); units = u"K"), auxiliary(:x, Surface(XY()); units = u"K")) + @test vardomain(vars.auxiliary[:x]) == Surface() + # ... in either declaration order + vars = Variables(auxiliary(:x, Surface(XY()); units = u"K"), input(:x, XY(); units = u"K")) + @test vardomain(vars.auxiliary[:x]) == Surface() + # ... and two domainless declarations stay domainless + @test isnothing(vardomain(Variables(input(:x, XY()), auxiliary(:x, XY())).auxiliary[:x])) + + # Two *different* stated domains remain an error. Note that these must be declarations of the + # same kind: a prognostic/auxiliary pair of the same name trips the cross-group duplicate check + # instead, which is a different code path with its own message. + @test_throws ErrorException Variables(auxiliary(:x, Ground(XY())), auxiliary(:x, Snow(XY()))) + + # The conflict message names the two domains rather than printing both declarations. + err = try + Variables(auxiliary(:x, Ground(XY())), auxiliary(:x, Snow(XY()))) + catch e + sprint(showerror, e) + end + @test occursin("ground", err) && occursin("snow", err) +end diff --git a/test/coupled_models/land_model_tests.jl b/test/coupled_models/land_model_tests.jl index 659488ab9..b54d6bef8 100644 --- a/test/coupled_models/land_model_tests.jl +++ b/test/coupled_models/land_model_tests.jl @@ -239,10 +239,10 @@ end Δt = 60.0 S₀ = 0.1 set!(state.surface_excess_water, S₀) - pools = Float64[Array(state.surface_excess_water)[1, 1, 1]] + pools = Float64[Array(interior(state.surface_excess_water))[1, 1, 1]] for _ in 1:5 timestep!(integrator, Δt) - push!(pools, Array(state.surface_excess_water)[1, 1, 1]) + push!(pools, Array(interior(state.surface_excess_water))[1, 1, 1]) end # The pool is monotonically drawn down (never grows) ... @test all(diff(pools) .< 0) diff --git a/test/grids.jl b/test/grids.jl index b64ae1ba9..ee10be78c 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -306,7 +306,7 @@ end @test Terrarium.num_layers(grid) == 5 @test Oceananigans.Grids.topology(grid) == Oceananigans.Grids.topology(column_grid) @test architecture(grid) == architecture(column_grid) - @test size(Field(grid, XYZ())) == size(grid) + @test size(Field(grid, Terrarium.Ground(XYZ()))) == size(grid) # Grids which are not land grids are their own ground discretization. @test ground_domain(column_grid) === column_grid @@ -374,16 +374,24 @@ end @testset "Model construction from spatial discretizations" begin column_grid = ColumnGrid(UniformSpacing(Δz = 0.1f0, N = 5), 2) - # Models accept an ordinary spatial discretization and build their land grid internally. + # Single-domain models keep whatever discretization they are given: an ordinary spatial + # discretization is no longer wrapped in a land grid, so a model which resolves only the ground + # can be built on a plain Oceananigans grid. model = SoilModel(column_grid) - @test get_grid(model) isa LandGrid - @test ground_domain(get_grid(model)) === column_grid + @test get_grid(model) === column_grid - # A pre-built land grid is stored as-is. + # A pre-built land grid is likewise stored as-is. land_grid = LandGrid(column_grid) @test get_grid(SoilModel(land_grid)) === land_grid + @test ground_domain(get_grid(SoilModel(land_grid))) === column_grid - # Models can also be built directly on a plain Oceananigans grid. + # ... including a bare `RectilinearGrid`. rect_grid = RectilinearGrid(Float32, size = (2, 1, 5), x = (0, 1), y = (0, 1), z = (-1, 0)) - @test ground_domain(get_grid(SoilModel(rect_grid))) === rect_grid + @test get_grid(SoilModel(rect_grid)) === rect_grid + + # `LandModel` couples several vertical domains, so it is the exception: it builds a land grid + # from an ordinary discretization and stores a pre-built one unchanged. + @test get_grid(LandModel(column_grid; vegetation = nothing)) isa LandGrid + @test ground_domain(get_grid(LandModel(column_grid; vegetation = nothing))) === column_grid + @test get_grid(LandModel(land_grid; vegetation = nothing)) === land_grid end diff --git a/test/inputs/input_forcing.jl b/test/inputs/input_forcing.jl index 11238f5fb..c72093783 100644 --- a/test/inputs/input_forcing.jl +++ b/test/inputs/input_forcing.jl @@ -1,12 +1,12 @@ module ForcingInputTest using Terrarium -using Terrarium: AbstractLandGrid, prognostic, input +using Terrarium: AbstractGrid, XY, XYZ, prognostic, input using Test DEFAULT_NF = Float32 -@kwdef struct TestModel{NF, Grid <: AbstractLandGrid{NF}} <: Terrarium.AbstractModel{NF, Grid} +@kwdef struct TestModel{NF, Grid <: AbstractGrid{NF}} <: Terrarium.AbstractModel{NF, Grid} grid::Grid initializer = DefaultInitializer(eltype(grid)) timestepper = ForwardEuler(eltype(grid)) diff --git a/test/inputs/inputs.jl b/test/inputs/inputs.jl index cb0a89a31..59e0ec941 100644 --- a/test/inputs/inputs.jl +++ b/test/inputs/inputs.jl @@ -1,5 +1,6 @@ using Terrarium -using Terrarium: FieldInputSource, FieldTimeSeriesInputSource, Variables, initialize!, interior, InputSources, varname +using Terrarium: InputSources, FieldInputSource, FieldTimeSeriesInputSource, Variables +using Terrarium: initialize!, interior, varname using Test using Unitful @@ -10,7 +11,7 @@ using Unitful field_input = InputSource(grid, X1; name = :X1) @test isa(field_input, FieldInputSource) ## check that dimensions and name were inferred correctly - @test field_input.dims == XY() + @test Terrarium.vardims(field_input) == XY() @test varname(field_input) == :X1 @test variables(field_input) == (Terrarium.input(:X1, XY()),) ## check state variable is allocated and initialize! copies data @@ -35,7 +36,7 @@ using Unitful fts_input = InputSource(S1; name = :S1) @test isa(fts_input, FieldTimeSeriesInputSource) ## check that dimensions and name were inferred correctly - @test fts_input.dims == XY() + @test Terrarium.vardims(fts_input) == XY() @test varname(fts_input) == :S1 @test fts_input.fts === S1 # populate S1 with random data and check update_inputs! @@ -74,7 +75,7 @@ end # Test single field source = InputSource(grid, ring_field1; name = :temperature) @test isa(source, FieldInputSource) - @test source.dims == XY() + @test Terrarium.vardims(source) == XY() @test varname(source) == :temperature @test variables(source) == (Terrarium.input(:temperature, XY()),) @@ -98,5 +99,5 @@ end ring_field_3d = rand(ring_grid, 5) # 5 vertical levels source_3d = InputSource(grid, ring_field_3d; name = :field3d) @test isa(source_3d, FieldInputSource) - @test source_3d.dims == XYZ() + @test Terrarium.vardims(source_3d) == XYZ() end diff --git a/test/inputs/namespaced_inputs.jl b/test/inputs/namespaced_inputs.jl index db5318452..bc82a3d4d 100644 --- a/test/inputs/namespaced_inputs.jl +++ b/test/inputs/namespaced_inputs.jl @@ -1,5 +1,5 @@ using Terrarium -using Terrarium: Variables, Namespace, initialize!, interior, varname, matches_scope, with_scope +using Terrarium: Variables, Namespace, XY, XYZ, initialize!, interior, varname, matches_scope, with_scope using Test @testset "Namespaced input sources" begin diff --git a/test/inputs/raster_inputs.jl b/test/inputs/raster_inputs.jl index 2fec40b72..a076deb85 100644 --- a/test/inputs/raster_inputs.jl +++ b/test/inputs/raster_inputs.jl @@ -48,7 +48,7 @@ const RasterInputSource = TerrariumRastersExt.RasterInputSource # A static (time-invariant) raster reduces to a plain FieldInputSource source = InputSource(grid, raster) @test isa(source, Terrarium.FieldInputSource) - @test source.dims == XY() + @test Terrarium.vardims(source) == XY() @test varname(source) == :temperature # Check variables are correctly inferred diff --git a/test/reactant/correctness.jl b/test/reactant/correctness.jl index f217ba7ba..867be15e7 100644 --- a/test/reactant/correctness.jl +++ b/test/reactant/correctness.jl @@ -144,7 +144,7 @@ function test_model(config::Symbol; nsteps = NSTEPS, rtol = RTOL, atol = ATOL, N @testset "$name" begin # Regression guard: Test that the Reactant model and grid types are correct @testset "Reactant model and grid types" begin - @test typeof(Terrarium.get_grid(rea.model)) <: ReactantExt.ReactantLandGrid + @test typeof(Terrarium.get_grid(rea.model)) <: ReactantExt.ReactantGrid @test typeof(rea.model) <: ReactantExt.ReactantModel @test typeof(rea) <: ReactantExt.ReactantIntegrator end diff --git a/test/soil/soil_hydrology_tests.jl b/test/soil/soil_hydrology_tests.jl index 553f69dd3..33493bb6f 100644 --- a/test/soil/soil_hydrology_tests.jl +++ b/test/soil/soil_hydrology_tests.jl @@ -226,7 +226,7 @@ end strat = HomogeneousSoilStratigraphy(eltype(grid); porosity = ConstantSoilPorosity(eltype(grid); mineral_porosity)) bgc = ConstantSoilCarbonDensity(eltype(grid)) hydrology = SoilHydrology(eltype(grid), RichardsEq()) - infiltration_var = auxiliary(:infiltration, XY(), units = u"m/s", desc = "Infiltration flux") + infiltration_var = auxiliary(:infiltration, Terrarium.Ground(Terrarium.Top()), units = u"m/s", desc = "Infiltration flux") # `strat`'s own variables (the per-horizon namespace `porosity_top` reaches into) must be part # of the state too, not just `hydrology`'s. vars = merge(Variables(hydrology), Variables(strat), Variables((infiltration_var,))) diff --git a/test/soil/soil_stratigrapy_tests.jl b/test/soil/soil_stratigrapy_tests.jl index 474077657..ad1b23870 100644 --- a/test/soil/soil_stratigrapy_tests.jl +++ b/test/soil/soil_stratigrapy_tests.jl @@ -70,11 +70,11 @@ end @testset "Texture normalization" begin grid = ColumnGrid(CPU(), Float64, UniformSpacing(Δz = 0.1, N = 10), 10) - sand = Field(grid, XY()) + sand = Field(grid, Terrarium.Ground(XY())) set!(sand, 0.5) - silt = Field(grid, XY()) + silt = Field(grid, Terrarium.Ground(XY())) set!(silt, 0.4) - clay = Field(grid, XY()) + clay = Field(grid, Terrarium.Ground(XY())) set!(clay, 0.2) # violate bounds normalize_texture!(sand, silt, clay) # all entries sum to unity diff --git a/test/state_variables.jl b/test/state_variables.jl index f6e508174..39d32203f 100644 --- a/test/state_variables.jl +++ b/test/state_variables.jl @@ -1,18 +1,18 @@ using Terrarium using Test -using Terrarium: AbstractLandGrid, VarDims, XY, XYZ, prognostic, auxiliary, input, namespace +using Terrarium: AbstractGrid, VarDims, XY, XYZ, prognostic, auxiliary, input, namespace DEFAULT_NF = Float32 module StateVariablesTestTypes using Terrarium - using Terrarium: AbstractLandGrid, VarDims, XY, XYZ, prognostic, auxiliary, input, namespace + using Terrarium: AbstractGrid, VarDims, XY, XYZ, prognostic, auxiliary, input, namespace using Test - @kwdef struct SubModel{NF, Grid <: AbstractLandGrid{NF}} <: Terrarium.AbstractModel{NF, Grid} + @kwdef struct SubModel{NF, Grid <: AbstractGrid{NF}} <: Terrarium.AbstractModel{NF, Grid} grid::Grid initializer = DefaultInitializer(eltype(grid)) timestepper = ForwardEuler(eltype(grid)) @@ -25,7 +25,7 @@ module StateVariablesTestTypes input(:forcing, XY()), ) - @kwdef struct TestModel{NF, Grid <: AbstractLandGrid{NF}, Sub} <: Terrarium.AbstractModel{NF, Grid} + @kwdef struct TestModel{NF, Grid <: AbstractGrid{NF}, Sub} <: Terrarium.AbstractModel{NF, Grid} grid::Grid submodel::Sub = SubModel(; grid) initializer = DefaultInitializer(eltype(grid)) diff --git a/test/surface/albedo.jl b/test/surface/albedo.jl index 181be21a8..1460ea3a6 100644 --- a/test/surface/albedo.jl +++ b/test/surface/albedo.jl @@ -16,8 +16,8 @@ end albd = PrescribedAlbedo(eltype(grid)) state = ( inputs = ( - albedo = set!(Field(grid, XY()), 0.4), - emissivity = set!(Field(grid, XY()), 0.8), + albedo = set!(Field(grid, Terrarium.Ground(XY())), 0.4), + emissivity = set!(Field(grid, Terrarium.Ground(XY())), 0.8), ) ) @test albedo(1, 1, grid, state, albd) == 0.4 diff --git a/test/surface/hydrology/surface_runoff_tests.jl b/test/surface/hydrology/surface_runoff_tests.jl index b67161bc3..b78b944b1 100644 --- a/test/surface/hydrology/surface_runoff_tests.jl +++ b/test/surface/hydrology/surface_runoff_tests.jl @@ -67,7 +67,7 @@ end S = 0.1 set!(state.surface_excess_water, S) Terrarium.compute_tendencies!(state, grid, runoff) - ∂S∂t = Array(state.tendencies.surface_excess_water)[1, 1, 1] + ∂S∂t = Array(interior(state.tendencies.surface_excess_water))[1, 1, 1] # The tendency is a negative removal rate, -min(D, S), so the pool is drawn down @test ∂S∂t ≈ -min(S / runoff.τ_r, S) @test ∂S∂t < 0 @@ -77,7 +77,7 @@ end state = StateVariables(runoff, grid) set!(state.surface_excess_water, S) Terrarium.compute_tendencies!(state, grid, runoff) - ∂S∂t = Array(state.tendencies.surface_excess_water)[1, 1, 1] + ∂S∂t = Array(interior(state.tendencies.surface_excess_water))[1, 1, 1] @test ∂S∂t ≈ -S # An empty pool has no tendency diff --git a/test/timestepping/explicit_step.jl b/test/timestepping/explicit_step.jl index e18171fe8..9717855ab 100644 --- a/test/timestepping/explicit_step.jl +++ b/test/timestepping/explicit_step.jl @@ -59,3 +59,40 @@ Terrarium.variables(closure::TestClosure) = ( # check that z was not changed (inverse closure not evaluated) @test all(iszero.(state.auxiliary.z)) end + +@testset "Forward Euler on vertically sliced fields" begin + # A variable declared at the `Top` or `Bottom` of a domain occupies a single vertical index but + # still carries a `Center`/`Face` vertical location, so it is indistinguishable from a fully + # resolved variable by its location parameters alone. `explicit_step!` must dispatch on the + # field's `indices` instead and step only the index the slice occupies. + Δt = 10.0 + Nz = 10 + grid = ColumnGrid(CPU(), Float64, ExponentialSpacing(N = Nz)) + clock = Clock(time = 0.0) + + for (name, loc, k) in ( + ("Top(Face)", Terrarium.Top(), Nz + 1), + ("Top(Center)", Terrarium.Top(z = Center()), Nz), + ("Bottom(Face)", Terrarium.Bottom(), 1), + ) + @testset "$name" begin + state = ( + prognostic = (x = Field(grid, loc),), + auxiliary = (;), + tendencies = (x = Field(grid, loc),), + namespaces = (;), + clock = clock, + ) + dxdt = 0.1 + set!(state.tendencies.x, dxdt) + + Terrarium.explicit_step!(state, grid, ForwardEuler(; Δt), Δt, (:x,)) + + # the slice is stepped, whichever vertical index it occupies + @test axes(state.prognostic.x, 3) == k:k + @test all(interior(state.prognostic.x) .≈ Δt * dxdt) + # and nothing outside the slice is touched, i.e. the step did not walk the whole column + @test count(!iszero, parent(state.prognostic.x)) == length(interior(state.prognostic.x)) + end + end +end diff --git a/test/timestepping/heun.jl b/test/timestepping/heun.jl index 9e274caab..129502abf 100644 --- a/test/timestepping/heun.jl +++ b/test/timestepping/heun.jl @@ -3,7 +3,7 @@ using Test # mock a simple model with exponential dynamics (and a constant offset) to test time steppers -@kwdef struct ExpModel{NF, Grid <: Terrarium.AbstractLandGrid{NF}, I, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} +@kwdef struct ExpModel{NF, Grid <: Terrarium.AbstractGrid{NF}, I, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} grid::Grid initializer::I = DefaultInitializer(eltype(grid)) timestepper::TS = ForwardEuler(eltype(grid)) @@ -72,7 +72,7 @@ end # variable (and its auxiliary offset and an input) living inside a namespace `:inner`. # This exercises time stepping of prognostic and input variables defined in namespaces and # also tests that actually the correct timestepper is used in the namespace as well. -@kwdef struct NamespacedExpModel{NF, Grid <: Terrarium.AbstractLandGrid{NF}, I, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} +@kwdef struct NamespacedExpModel{NF, Grid <: Terrarium.AbstractGrid{NF}, I, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} grid::Grid initializer::I = DefaultInitializer(eltype(grid)) timestepper::TS = ForwardEuler(eltype(grid)) diff --git a/test/timestepping/imex.jl b/test/timestepping/imex.jl index 68ae4f77d..9045dd42b 100644 --- a/test/timestepping/imex.jl +++ b/test/timestepping/imex.jl @@ -1,12 +1,12 @@ using Terrarium using Test -using Terrarium: AbstractLandGrid, AbstractIMEX, AbstractTimeStepper, AbstractVariable, EmptyCache, Explicit, Implicit, prognostic, XY +using Terrarium: AbstractGrid, AbstractIMEX, AbstractTimeStepper, AbstractVariable, EmptyCache, Explicit, Implicit, prognostic, XY module IMEXTestTypes using Terrarium - using Terrarium: AbstractLandGrid, AbstractIMEX, AbstractTimeStepper, AbstractVariable, Implicit, prognostic, XY + using Terrarium: AbstractGrid, AbstractIMEX, AbstractTimeStepper, AbstractVariable, Implicit, prognostic, XY # A mock implicit timestepper used only to verify IMEX routing. Its update is deliberately distinct # from forward Euler (u += 2·∂u∂t·Δt) so we can tell which sub-stepper integrated which variable. @@ -32,13 +32,16 @@ module IMEXTestTypes # Minimal two-variable model with constant unit tendencies. Both variables default to the `Explicit` # timestepping class; specific routing under an IMEX timestepper is declared via `timestepping` - @kwdef struct TwoVarModel{NF, Grid <: AbstractLandGrid{NF}, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} + @kwdef struct TwoVarModel{NF, Grid <: AbstractGrid{NF}, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} grid::Grid initializer = DefaultInitializer(eltype(grid)) timestepper::TS = ForwardEuler(eltype(grid)) end - Terrarium.variables(::TwoVarModel) = (prognostic(:a, XY()), prognostic(:b, XY())) + Terrarium.variables(::TwoVarModel) = ( + prognostic(:a, XY()), + prognostic(:b, XY()), + ) Terrarium.compute_auxiliary!(state, ::TwoVarModel) = nothing function Terrarium.compute_tendencies!(state, ::TwoVarModel) set!(state.tendencies.a, 1.0) @@ -51,13 +54,16 @@ module IMEXTestTypes # A second model identical to `TwoVarModel` but with the routing flipped, used to exercise the # per-model nature of `timestepping`: here `:a` is integrated implicitly and `:b` explicitly. - @kwdef struct FlippedModel{NF, Grid <: AbstractLandGrid{NF}, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} + @kwdef struct FlippedModel{NF, Grid <: AbstractGrid{NF}, TS <: Terrarium.AbstractTimeStepper} <: Terrarium.AbstractModel{NF, Grid} grid::Grid initializer = DefaultInitializer(eltype(grid)) timestepper::TS = ForwardEuler(eltype(grid)) end - Terrarium.variables(::FlippedModel) = (prognostic(:a, XY()), prognostic(:b, XY())) + Terrarium.variables(::FlippedModel) = ( + prognostic(:a, XY()), + prognostic(:b, XY()), + ) Terrarium.compute_auxiliary!(state, ::FlippedModel) = nothing function Terrarium.compute_tendencies!(state, ::FlippedModel) set!(state.tendencies.a, 1.0)