Coupling LandModel with full surface energy balance to SpeedyWeather

docs/src/processes/surface_energy/skin_temperature.md

@@ -21,7 +21,7 @@ ImplicitSkinTemperature
 
 The implicit approach requires solving a nonlinear equation at each grid point and time step. In order to avoid iteration, the current approach implementation of [`ImplicitSkinTemperature`](@ref) in Terrarium approximately solves for the skin temperature using a fixed point iteration where the ground heat flux is computed as the residual energy flux given the current prognostic state of the `skin_temperature` $T_s$,
 ```math
-G^\star = R_{\text{net}}(T_s) - H_s(T_s) - H_l(T_s)
+G^\star = R_{\text{net}}(T_s) + H_s(T_s) + H_l(T_s)
 ```
 and the new skin temperature $T_s^\star$ is determined by setting this heat flux equal to gradient between the skin and the ground surface temperature (uppermost soil layer) $T_g$,
 ```math

docs/src/processes/surface_energy/surface_energy_balance.md

@@ -16,10 +16,10 @@ The surface energy balance (SEB) describes how solar radiation, thermal radiatio
 
 ```math
 \begin{equation}
-R_{\text{net}} = H_s + H_l + G\,.
+R_{\text{net}} + H_s + H_l - G = 0\,.
 \end{equation}
 ```
-Following the standard convention of Terrarium and Oceananigans, all surface energy fluxes are defined **positive upward** (away from surface).
+Following the standard convention of Terrarium and Oceananigans, all surface energy fluxes are defined **positive upward**. The negative sign in front of $G$ reflects that heat flows *towards* the surface from the uppermost ground layer.
 
 The [`SurfaceEnergyBalance`](@ref) process is responsible for computing all of the above flux terms and thus closing the energy balance between the atmosphere and land surface. Implementations of [`AbstractSurfaceEnergyBalance`](@ref) should generally include, at minimum, representations of each of the four SEB components:
 - An implementation of [`AbstractSkinTemperature`](@ref) that defines and updates both the skin temperature $T_s$ and the ground heat flux $G$. The skin temperature $T_s$ is the effective radiative temperature of the land surface. For an implicit approach, $T_s$ self-consistently satisfies the energy balance at each time step. For a prescribed approach, $T_s$ is given as input. The ground heat flux at the surface is derived either directly or as a residual from the energy balance. See [Skin temperature and ground heat flux](@ref) for further details.

examples/Project.toml

@@ -15,10 +15,12 @@ Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 Rasters = "a3a2b9e3-a471-40c9-b274-f788e487c689"
 RingGrids = "d1845624-ad4f-453b-8ff4-a8db365bf3a7"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
-SpeedyWeatherInternals = "34489162-d270-4603-8b96-37b04f830d73"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Terrarium = "80418d68-07fa-499d-ae2b-c12e531f5cd8"
 
+[compat]
+SpeedyWeather = "0.18.1"
+
 [sources.Terrarium]
 path = ".."

examples/simulations/land_column.jl

@@ -14,9 +14,14 @@ vegetation = VegetationCarbon(eltype(grid))
 land = LandModel(grid; soil, vegetation)
 # Variably saturated with water table at roughly 5 m depth
 initializers = (
+    temperature = 15.0,
     saturation_water_ice = (x, z) -> min(1, 0.5 - 0.1 * z),
-    carbon_vegetation = 0.1,
+    carbon_vegetation = 0.5,
 )
 integrator = @time initialize(land, ForwardEuler(); initializers);
+# manually set atmospheric inputs to different values
+set!(integrator.state.windspeed, 1.0) # 1 m/s
+set!(integrator.state.specific_humidity, 1.0e-4) # kg/kg
 Δt = 60.0
 @time timestep!(integrator, Δt)
+@show integrator.state.latent_heat_flux[1, 1, 1]

examples/simulations/speedy_dry_land.jl

@@ -14,38 +14,52 @@ import SpeedyWeather as Speedy
 Speedy.SpeedyTransforms.FFTW.set_num_threads(1)
 
 """
-Naive implementation of a SpeedyWeather "dry" land model that couples to a Terrarium model.
+Naive implementation of a SpeedyWeather "dry" land model based on Terrarium.
 """
-struct TerrariumDryLand{NF, TMI <: Terrarium.ModelIntegrator{NF}} <: Speedy.AbstractDryLand
+struct TerrariumDryLand{
+        NF,
+        LG <: Speedy.LandGeometry,
+        TM <: Terrarium.ModelIntegrator{NF},
+    } <: Speedy.AbstractDryLand
     "Speedy spectral grid"
     spectral_grid::Speedy.SpectralGrid
 
+    "Speedy land model geometry"
+    geometry::LG
+
     "Initialized Terrarium model integrator"
-    integrator::TMI
+    integrator::TM
 
-    function TerrariumDryLand(itnegrator::Terrarium.ModelIntegrator{NF, Arch, Grid}; spectral_grid_kwargs...) where {NF, Arch, Grid <: ColumnRingGrid}
-        spectral_grid = Speedy.SpectralGrid(integrator.model.grid.rings; NF, nlayers_soil = 1, spectral_grid_kwargs...)
-        return new{eltype(integrator), typeof(integrator)}(spectral_grid, integrator)
+    function TerrariumDryLand(integrator::Terrarium.ModelIntegrator{NF, Arch, Grid}; spectral_grid_kwargs...) where {NF, Arch, Grid <: ColumnRingGrid}
+        spectral_grid = Speedy.SpectralGrid(integrator.model.grid.rings; NF, spectral_grid_kwargs...)
+        land_grid = integrator.grid
+        Δz = on_architecture(CPU(), land_grid.z.Δᵃᵃᶜ)
+        geometry = Speedy.LandGeometry(1, Δz[end])
+        return new{eltype(integrator), typeof(geometry), typeof(integrator)}(spectral_grid, geometry, integrator)
     end
 end
 
+Speedy.variables(land::TerrariumDryLand) = (
+    Speedy.PrognosticVariable(name = :soil_temperature, dims = Speedy.Grid3D(), namespace = :land),
+)
+
 function Speedy.initialize!(
-        progn_land::Speedy.PrognosticVariablesLand,  # for dispatch
-        progn::Speedy.PrognosticVariables{NF},
-        diagn::Speedy.DiagnosticVariables{NF},
-        land::TerrariumDryLand,
+        ::Any,
+        progn::Speedy.PrognosticVariables,
+        diagn::Speedy.DiagnosticVariables,
+        land::TerrariumDryLand{NF},
         model::Speedy.PrimitiveEquation,
     ) where {NF}
-    Tsoil = interior(land.integrator.state.temperature)[:, 1, end]
-    progn_land.soil_temperature .= Tsoil .+ NF(273.15)
+    Tsoil = interior(land.integrator.state.temperature)[:, 1, end] .+ 273.15
+    progn.land.soil_temperature .= Tsoil
     Terrarium.initialize!(land.integrator)
     return nothing
 end
 
 function Speedy.timestep!(
-        progn::Speedy.PrognosticVariables{NF},
-        diagn::Speedy.DiagnosticVariables{NF},
-        land::TerrariumDryLand,
+        progn::Speedy.PrognosticVariables,
+        diagn::Speedy.DiagnosticVariables,
+        land::TerrariumDryLand{NF},
         model::Speedy.PrimitiveEquation,
     ) where {NF}
     # get speedy state variables
@@ -71,11 +85,9 @@ ring_grid = RingGrids.FullGaussianGrid(24)
 Nz = 30
 Δz_min = 0.05 # currently the coupling is only stable with a large surface layer
 grid = ColumnRingGrid(CPU(), Float32, ExponentialSpacing(; N = Nz, Δz_min), ring_grid)
-# grid = ColumnGrid(CPU(), Float32, ExponentialSpacing(N=30))
 # Initial conditions
 soil_initializer = SoilInitializer(eltype(grid))
-
-# Soil model with prescribed surface temperautre BC
+# Soil model with prescribed surface temperature BC
 model = SoilModel(grid, initializer = soil_initializer)
 air_temperature = Field(grid, XY())
 Tair_input = InputSource(grid, air_temperature; name = :air_temperature)
@@ -86,30 +98,26 @@ integrator = initialize(model, ForwardEuler(eltype(grid)), Tair_input, boundary_
 land = TerrariumDryLand(integrator)
 # Set up coupled model
 land_sea_mask = Speedy.RockyPlanetMask(land.spectral_grid)
-output = Speedy.NetCDFOutput(land.spectral_grid, Speedy.PrimitiveDryModel, path = "outputs/")
+output = Speedy.NetCDFOutput(land.spectral_grid, Speedy.PrimitiveDryModel, land.geometry, path = "outputs/")
 time_stepping = Speedy.Leapfrog(land.spectral_grid, Δt_at_T31 = Minute(15))
 primitive_dry_coupled = Speedy.PrimitiveDryModel(land.spectral_grid; land, land_sea_mask, time_stepping, output)
 # add soil temperature as output variable for Speedy simulation
 Speedy.add!(primitive_dry_coupled.output, Speedy.SoilTemperatureOutput())
 # initialize coupled simulation
 sim_coupled = Speedy.initialize!(primitive_dry_coupled)
-# Speedy.timestep!(sim)
 # run it
-Speedy.run!(sim_coupled, period = Month(1), output = true)
+Speedy.run!(sim_coupled, period = Day(2), output = true)
 
 # Soil temperature in the 5th layer (~0.54 m)
 Tsoil_fig = heatmap(RingGrids.Field(interior(integrator.state.temperature)[:, 1, end - 4], grid), title = "", size = (800, 400))
 # Atmosphere variables
 Tair_fig = heatmap(sim_coupled.diagnostic_variables.grid.temp_grid[:, 8] .- 273.15, title = "Air temperature", size = (800, 400))
 pres_fig = heatmap(exp.(sim_coupled.diagnostic_variables.grid.pres_grid), title = "Surface pressure", size = (800, 400))
 srad_fig = heatmap(exp.(sim_coupled.diagnostic_variables.physics.surface_shortwave_down), title = "Surface shortwave down", size = (800, 400))
-# Tskin_fig = heatmap(RingGrids.Field(interior(integrator.state.skin_temperature)[:,1,end], grid), title="", size=(800,400))
-# save("plots/speedy_primitive_dry_coupled_tair_R48.png", Tair_fig, px_per_unit=1)
-# save("plots/speedy_primitive_dry_coupled_tsoil_R48.png", Tsoil_fig, px_per_unit=1)
 
 # animate surface air and soil temperatures
-Speedy.animate(sim, variable = "temp", coastlines = false, level = spectral_grid.nlayers, output_file = "plots/speedy_terrarium_dry_air_temperature.mp4")
-Speedy.animate(sim, variable = "st", coastlines = false, level = 1, output_file = "plots/speedy_terrarium_dry_soil_temperature.mp4")
+Speedy.animate(sim_coupled, variable = "temp", coastlines = false, level = spectral_grid.nlayers, output_file = "plots/speedy_terrarium_dry_air_temperature.mp4")
+# Speedy.animate(sim_coupled, variable = "st", coastlines = false, level = 1, output_file = "plots/speedy_terrarium_dry_soil_temperature.mp4") # TODO: this is broken now for some reason?
 
 # pick a point somewhere in the mid-lattitudes
 T = interior(integrator.state.temperature)[2000, 1, :]

examples/simulations/speedy_wet_land.jl

@@ -0,0 +1,184 @@
+using Terrarium
+
+using CUDA
+using Dates
+using Rasters, NCDatasets
+using Statistics
+
+using CairoMakie, GeoMakie
+
+import RingGrids
+import SpeedyWeather as Speedy
+
+# Choose architecture based on available hardware
+arch = CUDA.functional() ? GPU() : CPU()
+
+"""
+Naive implementation of a SpeedyWeather "wet" land model based on Terrarium.
+"""
+struct TerrariumWetLand{
+        NF,
+        LG <: Speedy.LandGeometry,
+        TM <: Terrarium.ModelIntegrator{NF},
+    } <: Speedy.AbstractWetLand
+    "Speedy spectral grid"
+    spectral_grid::Speedy.SpectralGrid
+
+    "Speedy land model geometry"
+    geometry::LG
+
+    "Initialized Terrarium model integrator"
+    integrator::TM
+
+    function TerrariumWetLand(integrator::Terrarium.ModelIntegrator{NF, Arch, Grid}; spectral_grid_kwargs...) where {NF, Arch, Grid <: ColumnRingGrid}
+        spectral_grid = Speedy.SpectralGrid(integrator.model.grid.rings; NF, spectral_grid_kwargs...)
+        land_grid = integrator.grid
+        Δz = on_architecture(CPU(), land_grid.z.Δᵃᵃᶜ)
+        geometry = Speedy.LandGeometry(1, Δz[end])
+        return new{eltype(integrator), typeof(geometry), typeof(integrator)}(spectral_grid, geometry, integrator)
+    end
+end
+
+Speedy.variables(land::TerrariumWetLand) = (
+    Speedy.PrognosticVariable(name = :soil_temperature, dims = Speedy.Grid3D(), namespace = :land),
+    Speedy.PrognosticVariable(name = :soil_moisture, dims = Speedy.Grid3D(), namespace = :land),
+)
+
+function Speedy.initialize!(
+        ::Any,
+        progn::Speedy.PrognosticVariables,
+        diagn::Speedy.DiagnosticVariables,
+        land::TerrariumWetLand{NF},
+        model::Speedy.PrimitiveEquation,
+    ) where {NF}
+    Terrarium.initialize!(land.integrator)
+    Tsoil = interior(land.integrator.state.temperature)[:, 1, end] .+ 273.15
+    sat = interior(land.integrator.state.saturation_water_ice)[:, 1, end]
+    progn.land.soil_temperature .= Tsoil
+    progn.land.soil_moisture .= sat
+    return nothing
+end
+
+function Speedy.timestep!(
+        progn::Speedy.PrognosticVariables,
+        diagn::Speedy.DiagnosticVariables,
+        land::TerrariumWetLand{NF},
+        model::Speedy.PrimitiveEquation,
+    ) where {NF}
+    speedy_timestep!(progn, diagn, land)
+    return nothing
+end
+
+function speedy_timestep!(
+        progn::Speedy.PrognosticVariables,
+        diagn::Speedy.DiagnosticVariables,
+        land::TerrariumWetLand{NF},
+    ) where {NF}
+    # land constants
+    consts = land.integrator.model.constants
+    # get speedy state variables
+    Tair = @view diagn.grid.temp_grid[:, end]
+    humid = @view diagn.grid.humid_grid[:, end]
+    pres = diagn.grid.pres_grid
+    wind = diagn.physics.surface_wind_speed
+    rain = diagn.physics.rain_rate
+    snow = diagn.physics.snow_rate
+    SwIn = diagn.physics.surface_shortwave_down
+    LwIn = diagn.physics.surface_longwave_down
+    # update Terrarium input fields
+    state = land.integrator.state
+    set!(state.inputs.air_temperature, Tair) # first set directly to avoid allocating new fields
+    set!(state.inputs.air_temperature, state.inputs.air_temperature - 273.15) # then convert to celsius
+    set!(state.inputs.air_pressure, pres) # similar with pressure
+    set!(state.inputs.air_pressure, exp(state.inputs.air_pressure)) # take exp to get pressure in Pa
+    set!(state.inputs.specific_humidity, humid)
+    set!(state.inputs.rainfall, rain)
+    set!(state.inputs.snowfall, snow)
+    set!(state.inputs.windspeed, wind)
+    set!(state.inputs.surface_shortwave_down, SwIn)
+    set!(state.inputs.surface_longwave_down, LwIn)
+    # run land forward over speedy timestep interval;
+    # we use a smaller actual timestep to ensure stability
+    Terrarium.run!(land.integrator, period = progn.clock.Δt, Δt = 300.0)
+    # Update speedy variables
+    progn.land.soil_temperature .= state.skin_temperature .+ NF(273.15)
+    progn.land.soil_moisture .= interior(state.saturation_water_ice)[:, 1, end]
+    progn.land.sensible_heat_flux .= state.sensible_heat_flux
+    progn.land.surface_humidity_flux .= state.latent_heat_flux ./ consts.Llg
+    diagn.physics.surface_longwave_up .= state.surface_longwave_up
+    diagn.physics.surface_shortwave_up .= state.surface_shortwave_up
+    return nothing
+end
+
+# quick test of default Speedy PrimitiveWetModel
+ring_grid = RingGrids.FullGaussianGrid(24)
+spectral_grid = Speedy.SpectralGrid(ring_grid)
+primitive_wet = Speedy.PrimitiveWetModel(spectral_grid)
+sim = Speedy.initialize!(primitive_wet)
+Speedy.run!(sim, period = Day(1))
+
+Nz = 30
+Δz_min = 0.05
+grid = ColumnRingGrid(CPU(), Float32, ExponentialSpacing(; N = Nz, Δz_min), ring_grid)
+# Initial conditions
+soil_initializer = SoilInitializer(eltype(grid))
+soil = SoilEnergyWaterCarbon(eltype(grid), hydrology = SoilHydrology(eltype(grid)))
+# Land model with "prescribed" atmosphere (from the perspective of the land model at least...)
+# vegetation = PrescribedVegetationCarbon(eltype(grid))
+model = LandModel(grid; initializer = soil_initializer, vegetation = nothing, soil)
+initializers = (;)
+integrator = initialize(model, ForwardEuler(eltype(grid)); initializers)
+# check if land model works standalone (with default atmospheric state)
+timestep!(integrator, 60.0) # one step
+run!(integrator, period = Hour(1), Δt = 300.0) # one hour
+Terrarium.initialize!(integrator) # reinitialize before setting up atmosphere
+
+# Initialize Terrarium-Speedy land model
+land = TerrariumWetLand(integrator)
+# Set up coupled model
+land_sea_mask = Speedy.RockyPlanetMask(land.spectral_grid)
+surface_heat_flux = Speedy.SurfaceHeatFlux(land.spectral_grid, land = Speedy.PrescribedLandHeatFlux())
+surface_humidity_flux = Speedy.SurfaceHumidityFlux(land.spectral_grid, land = Speedy.PrescribedLandHumidityFlux())
+output = Speedy.NetCDFOutput(land.spectral_grid, Speedy.PrimitiveDryModel, land.geometry, path = "outputs/")
+time_stepping = Speedy.Leapfrog(land.spectral_grid, Δt_at_T31 = Minute(15))
+primitive_wet_coupled = Speedy.PrimitiveWetModel(
+    land.spectral_grid;
+    land,
+    surface_heat_flux,
+    surface_humidity_flux,
+    land_sea_mask,
+    time_stepping,
+    output
+)
+# add soil temperature as output variable for Speedy simulation
+Speedy.add!(primitive_wet_coupled.output, Speedy.SoilTemperatureOutput())
+# initialize coupled simulation
+sim_coupled = Speedy.initialize!(primitive_wet_coupled)
+# run it
+period = Day(1)
+@info "Running simulation for $period"
+@time Speedy.run!(sim_coupled, period = period)
+Terrarium.checkfinite!(integrator.state.prognostic)
+
+# Land variables
+Tsoil_fig = heatmap(RingGrids.Field(interior(integrator.state.temperature)[:, 1, end - 2], grid), title = "", size = (800, 400))
+Tsurf_fig = heatmap(RingGrids.Field(interior(integrator.state.skin_temperature)[:, 1], grid), title = "", size = (800, 400))
+Hs_fig = heatmap(RingGrids.Field(interior(integrator.state.sensible_heat_flux)[:, 1], grid), title = "", size = (800, 400))
+Hl_fig = heatmap(RingGrids.Field(interior(integrator.state.latent_heat_flux)[:, 1], grid), title = "", size = (800, 400))
+E_fig = heatmap(RingGrids.Field(interior(integrator.state.evaporation_ground)[:, 1], grid), title = "", size = (800, 400))
+sat_fig = heatmap(RingGrids.Field(interior(integrator.state.saturation_water_ice)[:, 1, end], grid), title = "", size = (800, 400))
+# Atmosphere variables
+Tair_fig = heatmap(sim_coupled.diagnostic_variables.grid.temp_grid[:, 8] .- 273.15, title = "Air temperature", size = (800, 400))
+pres_fig = heatmap(exp.(sim_coupled.diagnostic_variables.grid.pres_grid), title = "Surface pressure", size = (800, 400))
+srad_fig = heatmap(sim.diagnostic_variables.physics.surface_shortwave_down, title = "Surface shortwave down", size = (800, 400))
+
+# pick a point somewhere in the mid-lattitudes
+T = interior(integrator.state.temperature)[2000, 1, :]
+sat = interior(integrator.state.saturation_water_ice)[2000, 1, :]
+f = interior(integrator.state.liquid_water_fraction)[2000, 1, :]
+zs = znodes(integrator.state.temperature)
+
+# Plot temperature and liquid fraction profiles in upper 15 layers
+Makie.scatterlines(T, zs)
+Makie.scatterlines(sat, zs)
+Makie.scatterlines(f, zs)

src/diagnostics/debugging.jl

@@ -14,25 +14,25 @@ end
 """
     $SIGNATURES
 
-Check whether the given `field` has any `NaN` values using `Diagnostics.hasnan` and raise an error if `NaN`s are detected.
+Check whether the given `field` has any `NaN` or `Inf` values and raise an error if `NaN`s are detected.
 """
-nancheck!(field::AbstractField, name = nothing) = Diagnostics.hasnan(field) && error("Found NaNs in Field $name: $field")
-function nancheck!(nt::NamedTuple)
+checkfinite!(field::AbstractField, name = nothing) = any(!isfinite, parent(field)) && error("Found NaN/Inf values in Field $name: $field")
+function checkfinite!(nt::NamedTuple)
     for key in keys(nt)
-        nancheck!(nt[key], key)
+        checkfinite!(nt[key], key)
     end
-    return
+    return nothing
 end
 
 """
     $SIGNATURES
 
 Provides a "hook" for handling debug calls from relevant callsites. Default implementations for
-`Field` and `NamedTuple` (assumed to be of `Field`s) simply forward to [`nancheck!`](@ref).
+`Field` and `NamedTuple` (assumed to be of `Field`s) simply forward to [`checkfinite!`](@ref).
 """
 @inline debughook!(args...) = nothing
-@inline debughook!(field::AbstractField) = nancheck!(field)
-@inline debughook!(nt::NamedTuple) = nancheck!(nt)
+@inline debughook!(field::AbstractField) = checkfinite!(field)
+@inline debughook!(nt::NamedTuple) = checkfinite!(nt)
 
 """
     $SIGNATURES

src/grids/grid_utils.jl

@@ -2,7 +2,8 @@
 function Oceananigans.launch!(grid::AbstractLandGrid, workspec, kernel!::Function, first_arg, other_args...; kwargs...)
     fgrid = get_field_grid(grid)
     launch!(fgrid.architecture, fgrid, get_workspec(workspec), kernel!, first_arg, grid, other_args...; kwargs...)
-    return debugsite!(kernel!, first_arg, other_args...)
+    debugsite!(kernel!, first_arg, grid, other_args...)
+    return nothing
 end
 
 """