diff --git a/Project.toml b/Project.toml index eb46b53ba..3c8b05ad5 100644 --- a/Project.toml +++ b/Project.toml @@ -43,6 +43,6 @@ IntervalSets = "0.7" KernelAbstractions = "0.9" Oceananigans = "0.100, 0.101, 0.102, 0.103, 0.104, 0.105, 0.106" RingGrids = "0.1" -SpeedyWeatherInternals = "0.1" +SpeedyWeatherInternals = "0.2" Unitful = "1" julia = "1.10" diff --git a/docs/src/processes/surface_energy/skin_temperature.md b/docs/src/processes/surface_energy/skin_temperature.md index 06536182e..8667299f1 100644 --- a/docs/src/processes/surface_energy/skin_temperature.md +++ b/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 diff --git a/docs/src/processes/surface_energy/surface_energy_balance.md b/docs/src/processes/surface_energy/surface_energy_balance.md index 8ae2f8976..bb6fa35d9 100644 --- a/docs/src/processes/surface_energy/surface_energy_balance.md +++ b/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. diff --git a/examples/Project.toml b/examples/Project.toml index ee395dac8..0506e53c2 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -15,7 +15,6 @@ 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" diff --git a/examples/simulations/land_column.jl b/examples/simulations/land_column.jl index 3a8cc2b40..e8df98cb9 100644 --- a/examples/simulations/land_column.jl +++ b/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] diff --git a/examples/simulations/speedy_downscaling.jl b/examples/simulations/speedy_downscaling.jl new file mode 100644 index 000000000..18de35fb3 --- /dev/null +++ b/examples/simulations/speedy_downscaling.jl @@ -0,0 +1,374 @@ +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 +terrarium_arch = CUDA.functional() ? GPU() : CPU() +speedy_arch = CUDA.functional() ? Speedy.GPU() : Speedy.CPU() + +""" +Naive implementation of a SpeedyWeather "wet" land model based on Terrarium. +Operates with two separate ring grids: a lower-resolution grid for Speedy and a +higher-resolution grid for Terrarium, with `RingGrids.interpolate!` used to couple them. +""" +struct TerrariumWetLand{ + NF, + LG <: Speedy.LandGeometry, + TM <: Terrarium.ModelIntegrator{NF}, + IU <: RingGrids.AbstractInterpolator, + ID <: RingGrids.AbstractInterpolator, + FT <: RingGrids.AbstractField, + FS <: RingGrids.AbstractField, + } <: Speedy.AbstractWetLand + "Speedy spectral grid (low resolution)" + spectral_grid::Speedy.SpectralGrid + + "Speedy land model geometry" + geometry::LG + + "Initialized Terrarium model integrator (high resolution)" + integrator::TM + + "Pre-computed interpolator: Speedy ring grid (low-res) → Terrarium ring grid (high-res)" + interp_up::IU + + "Pre-computed interpolator: Terrarium ring grid (high-res) → Speedy ring grid (low-res)" + interp_down::ID + + "Reusable 2D buffer Field on Terrarium ring grid (for upscaling Speedy → Terrarium)" + buffer_terrarium::FT + + "Reusable 2D buffer Field on Speedy ring grid (for downscaling Terrarium → Speedy)" + buffer_speedy::FS + + function TerrariumWetLand( + integrator::Terrarium.ModelIntegrator{NF, Arch, Grid}, + speedy_ring_grid::RingGrids.AbstractGrid, + terrarium_ring_grid::RingGrids.AbstractGrid; + spectral_grid_kwargs... + ) where {NF, Arch, Grid <: ColumnRingGrid} + spectral_grid = Speedy.SpectralGrid(speedy_ring_grid; NF, spectral_grid_kwargs...) + interp_up = RingGrids.interpolator(terrarium_ring_grid, speedy_ring_grid) + interp_down = RingGrids.interpolator(speedy_ring_grid, terrarium_ring_grid) + buffer_terrarium = zeros(NF, terrarium_ring_grid) + buffer_speedy = zeros(NF, speedy_ring_grid) + land_grid = integrator.grid + Δz = on_architecture(CPU(), land_grid.z.Δᵃᵃᶜ) + geometry = Speedy.LandGeometry(1, Δz[end]) + return new{NF, typeof(geometry), typeof(integrator), typeof(interp_up), typeof(interp_down), typeof(buffer_terrarium), typeof(buffer_speedy)}( + spectral_grid, geometry, integrator, interp_up, interp_down, buffer_terrarium, buffer_speedy + ) + 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!( + vars::Speedy.Variables, + land::TerrariumWetLand{NF}, + model::Speedy.PrimitiveEquation, + ) where {NF} + Terrarium.initialize!(land.integrator) + state = land.integrator.state + grid = land.integrator.model.grid + # Downscale Terrarium initial state to Speedy resolution + scatter_land!(land.buffer_terrarium, interior(state.temperature)[:, 1, end], grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.prognostic.land.soil_temperature .= land.buffer_speedy .+ NF(273.15) + scatter_land!(land.buffer_terrarium, interior(state.saturation_water_ice)[:, 1, end], grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.prognostic.land.soil_moisture .= land.buffer_speedy + return nothing +end + +function Speedy.timestep!( + vars::Speedy.Variables, + land::TerrariumWetLand{NF}, + model::Speedy.PrimitiveEquation, + ) where {NF} + speedy_timestep!(vars, land) + return nothing +end + +## Scatter land-only Oceananigans interior data (length Nh) back to a full-resolution RingGrids Field2D. +## Sea grid points are set to zero. +function scatter_land!(out::RingGrids.AbstractField, data::AbstractArray, grid::ColumnRingGrid) + fill!(out, zero(eltype(out))) + out.data[grid.mask.data] .= data + return out +end + +function speedy_timestep!( + vars::Speedy.Variables, + land::TerrariumWetLand{NF}, + ) where {NF} + # land constants + consts = land.integrator.model.constants + state = land.integrator.state + terrarium_grid = land.integrator.model.grid + mask = terrarium_grid.mask.data + # Upscale Speedy fields to Terrarium resolution and update inputs + ## For each field: interpolate to full Terrarium ring grid, then gather only the land-masked points + RingGrids.interpolate!(land.buffer_terrarium, vars.grid.temperature[:, end], land.interp_up) + land.buffer_terrarium.data .-= NF(273.15) ## convert K → °C + interior(state.inputs.air_temperature)[:, 1, 1] .= land.buffer_terrarium.data[mask] + RingGrids.interpolate!(land.buffer_terrarium, vars.grid.pressure, land.interp_up) + land.buffer_terrarium.data .= exp.(land.buffer_terrarium.data) ## convert log(Pa) → Pa + interior(state.inputs.air_pressure)[:, 1, 1] .= land.buffer_terrarium.data[mask] + RingGrids.interpolate!(land.buffer_terrarium, vars.grid.humidity[:, end], land.interp_up) + interior(state.inputs.specific_humidity)[:, 1, 1] .= land.buffer_terrarium.data[mask] + RingGrids.interpolate!(land.buffer_terrarium, vars.parameterizations.rain_rate, land.interp_up) + interior(state.inputs.rainfall)[:, 1, 1] .= land.buffer_terrarium.data[mask] + RingGrids.interpolate!(land.buffer_terrarium, vars.parameterizations.snow_rate, land.interp_up) + interior(state.inputs.snowfall)[:, 1, 1] .= land.buffer_terrarium.data[mask] + RingGrids.interpolate!(land.buffer_terrarium, vars.parameterizations.surface_wind_speed, land.interp_up) + interior(state.inputs.windspeed)[:, 1, 1] .= land.buffer_terrarium.data[mask] + RingGrids.interpolate!(land.buffer_terrarium, vars.parameterizations.surface_shortwave_down, land.interp_up) + interior(state.inputs.surface_shortwave_down)[:, 1, 1] .= land.buffer_terrarium.data[mask] + RingGrids.interpolate!(land.buffer_terrarium, vars.parameterizations.surface_longwave_down, land.interp_up) + interior(state.inputs.surface_longwave_down)[:, 1, 1] .= land.buffer_terrarium.data[mask] + # run land forward over speedy timestep interval; + # we use a smaller actual timestep to ensure stability + Terrarium.run!(land.integrator, period = vars.prognostic.clock.Δt, Δt = 300.0) + # Downscale Terrarium output fields to Speedy resolution + scatter_land!(land.buffer_terrarium, interior(state.skin_temperature)[:, 1, 1], terrarium_grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.prognostic.land.soil_temperature .= land.buffer_speedy .+ NF(273.15) + scatter_land!(land.buffer_terrarium, interior(state.saturation_water_ice)[:, 1, end], terrarium_grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.prognostic.land.soil_moisture .= land.buffer_speedy + scatter_land!(land.buffer_terrarium, interior(state.sensible_heat_flux)[:, 1, 1], terrarium_grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.prognostic.land.sensible_heat_flux .= land.buffer_speedy + scatter_land!(land.buffer_terrarium, interior(state.latent_heat_flux)[:, 1, 1], terrarium_grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.prognostic.land.surface_humidity_flux .= land.buffer_speedy ./ consts.Llg + scatter_land!(land.buffer_terrarium, interior(state.surface_longwave_up)[:, 1, 1], terrarium_grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.parameterizations.land.surface_longwave_up .= land.buffer_speedy + scatter_land!(land.buffer_terrarium, interior(state.surface_shortwave_up)[:, 1, 1], terrarium_grid) + RingGrids.interpolate!(land.buffer_speedy, land.buffer_terrarium, land.interp_down) + vars.parameterizations.land.surface_shortwave_up .= land.buffer_speedy + return nothing +end + +# quick test of default Speedy PrimitiveWetModel on GPU +speedy_ring_grid = RingGrids.FullGaussianGrid(24, speedy_arch) +speedy_ring_grid_cpu = on_architecture(CPU(), speedy_ring_grid) +spectral_grid = Speedy.SpectralGrid(speedy_ring_grid) +orography = Speedy.EarthOrography(spectral_grid, smoothing = false) +primitive_wet = Speedy.PrimitiveWetModel(spectral_grid; orography) +sim = Speedy.initialize!(primitive_wet) +Speedy.run!(sim, period = Day(1)) + +# Higher-resolution ring grid on Speedy's architecture for use in coupling buffers/interpolators +terrarium_ring_grid = RingGrids.FullGaussianGrid(72 * 2, speedy_arch) +terrarium_ring_grid_cpu = on_architecture(CPU(), terrarium_ring_grid) + +# Load SpeedyWeather land-sea mask at Terrarium resolution; land = 1, sea = 0 (fractional) +## Threshold at 0.5 to get a Bool field suitable for ColumnRingGrid masking +land_sea_mask_raw = RingGrids.get_asset( + "data/boundary_conditions/land-sea_mask.nc", + from_assets = true, name = "lsm", + ArrayType = RingGrids.FullClenshawField, FileFormat = NCDataset, +) +land_sea_mask_terrarium = clamp.(RingGrids.grid_cell_average!(RingGrids.Field(terrarium_ring_grid_cpu), land_sea_mask_raw), 0, 1) +land_sea_mask_atmos = clamp.(RingGrids.grid_cell_average!(RingGrids.Field(speedy_ring_grid_cpu), land_sea_mask_raw), 0, 1) + +Nz = 30 +Δz_min = 0.05 +grid = ColumnRingGrid(terrarium_arch, Float32, ExponentialSpacing(; N = Nz, Δz_min), terrarium_ring_grid_cpu, land_sea_mask_terrarium .> 0.0f0) +# 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 using separate low-res (Speedy) and high-res (Terrarium) ring grids +land = TerrariumWetLand(integrator, speedy_ring_grid, terrarium_ring_grid) +# Set up coupled model +land_sea_mask = Speedy.LandSeaMask(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.PrimitiveWetModel, 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 (kept for reference) — we'll also run a stepwise loop below to build an animation +period = Day(10) +@info "Running simulation for $period" +@time Speedy.run!(sim_coupled; period) +Terrarium.checkfinite!(integrator.state.prognostic) + +heatmap(sim_coupled.variables.grid.temperature[:, 1]) + +# --- 2×2 Animation: land and atmosphere state evolution during the coupled run --- +# Re-initialize the simulation to start from the beginning for the animation. +Nt = 240 +Speedy.initialize!(sim_coupled; steps = Nt, output = false) +# Spin-up +Speedy.run!(sim_coupled; period = Day(10)) + +# Coordinate axes for terrarium (high-res) and Speedy (low-res) grids, computed once. +# We replicate what RingGridsMakieExt does internally for heatmap (no mutating variant exists) +# so that we can drive heatmap!(ax, lond, latd, obs::Observable{Matrix}) in the record loop. +lond_hi = Float32.(RingGrids.get_lond(terrarium_ring_grid_cpu)) +latd_hi = Float32.(RingGrids.get_latd(terrarium_ring_grid_cpu)) +lond_lo = Float32.(RingGrids.get_lond(speedy_ring_grid_cpu)) +latd_lo = Float32.(RingGrids.get_latd(speedy_ring_grid_cpu)) + +# CPU buffer for scatter_land! → Matrix conversion during each animation frame. +buffer_anim = zeros(Float32, terrarium_ring_grid_cpu) + +# Helper: scatter land-only Oceananigans field data at `layer` into `buffer_anim` and return +# as a (nlon_hi × nlat_hi) Matrix suitable for heatmap!. +# Sea grid cells are set to NaN so they render as masked (transparent) in the animation. +function terrarium_frame(oc_field, layer) + data_cpu = Array(interior(oc_field)[:, 1, layer]) + fill!(buffer_anim, NaN32) + buffer_anim.data[grid.mask.data] .= data_cpu + return Matrix(buffer_anim) +end + +# Helper: extract a (nlon_lo × nlat_lo) Matrix from a Speedy 2D ring-grid field (moves to CPU). +# An optional element-wise transform is applied to the flat data before reshaping. +function speedy_frame(rg_field_2d; transform = identity) + data_cpu = transform(Array(rg_field_2d.data)) + return reshape(data_cpu, length(lond_lo), length(latd_lo)) +end + +# Initial matrices for all four panels. +mat_tsoil = terrarium_frame(land.integrator.state.temperature, Nz) +mat_evap = terrarium_frame(land.integrator.state.evaporation_ground, 1) +mat_tair = speedy_frame(sim_coupled.variables.grid.temperature[:, end]; transform = x -> x .- 273.15f0) +mat_humid = speedy_frame(sim_coupled.variables.grid.humidity[:, end]) + +# Build the 2×2 Figure layout. +# Axes occupy columns 1 and 3; their matching Colorbars occupy columns 2 and 4. +function heatmap_axis(fig, row, col, title) + return Axis( + fig[row, col]; + title, aspect = 2, titlesize = 12, + xticks = 0:60:360, yticks = -60:30:60, + xticklabelsize = 9, yticklabelsize = 9, + xtickformat = vs -> ["$(round(Int, v))˚E" for v in vs], + ytickformat = vs -> ["$(round(Int, v))˚N" for v in vs], + ) +end + + +# Fixed colorbar ranges — standardised so colors are consistent across all frames. +# Adjust these if your simulation uses very different conditions. +crange_tsoil = (-20.0f0, 20.0f0) # °C — typical land surface temperature range +crange_evap = (0.0f0, 1.0f-5) # m s⁻¹ — bare-soil evaporation +crange_tair = (-20.0f0, 20.0f0) # °C — tropospheric temperature at lowest model level +crange_humid = (0.0f0, 0.01) # [-] — specificy humidity + +Makie.with_theme(fontsize = 16) do + fig_anim = Figure(size = (1400, 700)) + + ax_tsoil = heatmap_axis(fig_anim, 1, 1, "Soil temperature (°C)") + ax_evap = heatmap_axis(fig_anim, 1, 3, "Evaporation (m s⁻¹)") + ax_tair = heatmap_axis(fig_anim, 2, 1, "Air temperature (°C)") + ax_humid = heatmap_axis(fig_anim, 2, 3, "Specific humidity") + + obs_tsoil = Observable(mat_tsoil) + obs_evap = Observable(mat_evap) + obs_tair = Observable(mat_tair) + obs_humid = Observable(mat_humid) + + hm_tsoil = heatmap!(ax_tsoil, lond_hi, latd_hi, obs_tsoil; colormap = :temperaturemap, colorrange = crange_tsoil) + hm_evap = heatmap!(ax_evap, lond_hi, latd_hi, obs_evap; colormap = :viridis, colorrange = crange_evap) + hm_tair = heatmap!(ax_tair, lond_lo, latd_lo, obs_tair; colormap = :temperaturemap, colorrange = crange_tair) + hm_humid = heatmap!(ax_humid, lond_lo, latd_lo, obs_humid; colormap = :Blues, colorrange = crange_humid) + + Colorbar(fig_anim[1, 2], hm_tsoil; label = "°C", ticklabelsize = 8) + Colorbar(fig_anim[1, 4], hm_evap; label = "m s⁻¹", ticklabelsize = 8) + Colorbar(fig_anim[2, 2], hm_tair; label = "°C", ticklabelsize = 8) + Colorbar(fig_anim[2, 4], hm_humid; label = "", ticklabelsize = 8) + + clock_anim = sim_coupled.variables.prognostic.clock + time_label = Observable(string(clock_anim.time)) + Label(fig_anim[0, :], time_label; fontsize = 14) + + mkpath("plots") + Makie.record(fig_anim, "plots/speedy_terrarium_bare_soil_atmosphere_10days.mp4", 1:Nt; framerate = 12) do _ + Speedy.run!(sim_coupled, period = Hour(1)) + obs_tsoil[] = terrarium_frame(land.integrator.state.temperature, Nz) + obs_evap[] = terrarium_frame(land.integrator.state.evaporation_ground, 1) + obs_tair[] = speedy_frame(sim_coupled.variables.grid.temperature[:, end]; transform = x -> x .- 273.15f0) + obs_humid[] = speedy_frame(sim_coupled.variables.grid.humidity[:, end]) + time_label[] = string(clock_anim.time) + end +end + +## Plot a RingGrids field: converts to CPU, draws heatmap, and labels the colorbar as "name / units". +function land_heatmap(field, grid::ColumnRingGrid, label; kwargs...) + fig = heatmap(RingGrids.Field(field, grid); kwargs...) + filter(x -> x isa Makie.Colorbar, fig.content)[1].label = label + return fig +end + +function atmos_heatmap(field, ring_grid, label; kwargs...) + fig = heatmap(RingGrids.Field(field, ring_grid); kwargs...) + filter(x -> x isa Makie.Colorbar, fig.content)[1].label = label + return fig +end + +# Land variables — convert masked Oceananigans field to full CPU ring grid, then plot +Tsoil_fig = land_heatmap(RingGrids.Field(on_architecture(CPU(), integrator.state.temperature), grid)[:, end - 5], terrarium_ring_grid_cpu, "Temperature / °C", title = "Soil temperature", size = (800, 400)) +save("plots/terrarium_speedy_Tsoil_hires.png", Tsoil_fig) +Tsurf_fig = land_heatmap(RingGrids.Field(on_architecture(CPU(), integrator.state.skin_temperature), grid)[:, 1], terrarium_ring_grid_cpu, "Temperature / °C", title = "Skin temperature", size = (800, 400)) +save("plots/terrarium_speedy_Tsurf_hires.png", Tsurf_fig) +Hs_fig = land_heatmap(RingGrids.Field(on_architecture(CPU(), integrator.state.sensible_heat_flux), grid)[:, 1], terrarium_ring_grid_cpu, "Sensible heat flux / W m⁻²", title = "Sensible heat flux", size = (800, 400)) +save("plots/terrarium_speedy_Hs_hires.png", Hs_fig) +Hl_fig = land_heatmap(RingGrids.Field(on_architecture(CPU(), integrator.state.latent_heat_flux), grid)[:, 1], terrarium_ring_grid_cpu, "Latent heat flux / W m⁻²", title = "Latent heat flux", size = (800, 400)) +save("plots/terrarium_speedy_Hl_hires.png", Hl_fig) +E_fig = land_heatmap(RingGrids.Field(on_architecture(CPU(), integrator.state.evaporation_ground), grid)[:, 1], terrarium_ring_grid_cpu, "Evaporation / m s⁻¹", title = "Evaporation", size = (800, 400)) +save("plots/terrarium_speedy_E_hires.png", E_fig) +# sat_fig = land_heatmap(RingGrids.Field(on_architecture(CPU(), integrator.state.saturation_water_ice), grid)[:, end], terrarium_ring_grid_cpu, "Saturation / -", title = "Surface saturation", size = (800, 400)) + +# Atmosphere variables (Speedy v0.19 API) — already on FullGaussianGrid(16) +Tair_fig = atmos_heatmap(sim_coupled.variables.grid.temperature[:, 8].data .- 273.15f0, speedy_ring_grid_cpu, "Temperature / °C", title = "Air temperature", size = (800, 400)) +pres_fig = atmos_heatmap(exp.(sim_coupled.variables.grid.pressure.data), speedy_ring_grid_cpu, "Pressure / Pa", title = "Surface pressure", size = (800, 400)) +srad_fig = atmos_heatmap(sim_coupled.variables.parameterizations.surface_shortwave_down.data, speedy_ring_grid_cpu, "Shortwave down / W m⁻²", title = "Surface shortwave down", size = (800, 400)) + +# Pick a point somewhere in the mid-latitudes and plot vertical profiles +T = on_architecture(CPU(), interior(integrator.state.temperature)[8000, 1, :]) +sat = on_architecture(CPU(), interior(integrator.state.saturation_water_ice)[8000, 1, :]) +f = on_architecture(CPU(), interior(integrator.state.liquid_water_fraction)[8000, 1, :]) +zs = on_architecture(CPU(), znodes(integrator.state.temperature)) + +# Plot temperature, saturation, and liquid fraction vertical profiles +Makie.scatterlines(T, zs) +Makie.scatterlines(sat, zs) +Makie.scatterlines(f, zs) diff --git a/examples/simulations/speedy_dry_land.jl b/examples/simulations/speedy_dry_land.jl index 871353059..a33831ef6 100644 --- a/examples/simulations/speedy_dry_land.jl +++ b/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,16 +98,15 @@ 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)) @@ -103,13 +114,10 @@ Tsoil_fig = heatmap(RingGrids.Field(interior(integrator.state.temperature)[:, 1, 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, :] diff --git a/examples/simulations/speedy_wet_land.jl b/examples/simulations/speedy_wet_land.jl new file mode 100644 index 000000000..5da4491a8 --- /dev/null +++ b/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) diff --git a/src/diagnostics/debugging.jl b/src/diagnostics/debugging.jl index 62d3b9bc1..86e21e64e 100644 --- a/src/diagnostics/debugging.jl +++ b/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 diff --git a/src/grids/column_ring_grid.jl b/src/grids/column_ring_grid.jl index 28046fd67..ccd0547d5 100644 --- a/src/grids/column_ring_grid.jl +++ b/src/grids/column_ring_grid.jl @@ -100,17 +100,17 @@ get_field_grid(grid::ColumnRingGrid) = grid.grid Converts the given Oceananigans `Field` to a `RingGrids.Field` with a ring grid matching that of the given `ColumnRingGrid`. """ RingGrids.Field(field::Field{LX, LY, LZ}, grid::ColumnRingGrid; fill_value = NaN) where {LX, LY, LZ} = RingGrids.Field(architecture(field), dropdims(interior(field), dims = 2), grid; fill_value) +RingGrids.Field(arch::AbstractArchitecture, field::Field{LX, LY, LZ}, grid::ColumnRingGrid; fill_value = NaN) where {LX, LY, LZ} = RingGrids.Field(arch, dropdims(interior(field), dims = 2), grid; fill_value) RingGrids.Field(field::AbstractArray, grid::ColumnRingGrid; fill_value = NaN) = RingGrids.Field(architecture(grid), field, grid; fill_value) function RingGrids.Field(arch::AbstractArchitecture, field::AbstractArray, grid::ColumnRingGrid; fill_value = NaN) # need to be on CPU to do the copying grid = on_architecture(arch, grid) field = on_architecture(arch, field) # create new RingGrids field initialized with fill_value - ring_field = RingGrids.Field(grid.rings, size(field)[2:end]...) + ring_field = RingGrids.Field(grid.rings, size(field, 2)) fill!(ring_field, fill_value) # need to access underlying data arrays directly to avoid scalar indexing - colons = (Colon() for d in size(field)[2:end]) - ring_field.data[grid.mask.data, colons...] .= field + ring_field.data[grid.mask.data, :] .= field return ring_field end diff --git a/src/grids/grid_utils.jl b/src/grids/grid_utils.jl index ab6fd8c96..50f3ea535 100644 --- a/src/grids/grid_utils.jl +++ b/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 """ diff --git a/src/models/soil/soil_model_init.jl b/src/models/soil/soil_model_init.jl index 63e496413..880edbb0a 100644 --- a/src/models/soil/soil_model_init.jl +++ b/src/models/soil/soil_model_init.jl @@ -140,13 +140,13 @@ Properties: $TYPEDFIELDS """ @kwdef struct SaturationWaterTable{NF} <: AbstractInitializer{NF} - vadose_zone_saturation::NF = 0.5 + vadose_zone_saturation::NF = 0.75 water_table_depth::NF = 5.0 end SaturationWaterTable(::Type{NF}; kwargs...) where {NF} = SaturationWaterTable{NF}(; kwargs...) function initialize!(state, ::AbstractModel, init::SaturationWaterTable{NF}) where {NF} - set!(state.saturation_water_ice, (x, z) -> z <= init.water_table_depth ? one(NF) : init.vadose_zone_saturation) + set!(state.saturation_water_ice, (x, z) -> z <= -init.water_table_depth ? one(NF) : init.vadose_zone_saturation) return nothing end diff --git a/src/models/surface/surface_energy_model.jl b/src/models/surface/surface_energy_model.jl index bdebdc24f..12974851f 100644 --- a/src/models/surface/surface_energy_model.jl +++ b/src/models/surface/surface_energy_model.jl @@ -41,12 +41,11 @@ function SurfaceEnergyModel( end function compute_auxiliary!(state, model::SurfaceEnergyModel) - compute_auxiliary!(state, model, model.atmosphere) - compute_auxiliary!(state, model, model.surface_energy_balance) + compute_auxiliary!(state, model.grid, model.atmosphere) + compute_auxiliary!(state, model.grid, model.surface_energy_balance, model.constants, model.atmosphere) return nothing end function compute_tendencies!(state, ::SurfaceEnergyModel) - compute_tendencies!(state, model, model.surface_energy_balance) return nothing end diff --git a/src/processes/atmosphere/prescribed_atmosphere.jl b/src/processes/atmosphere/prescribed_atmosphere.jl index 42ac0cef6..e5d4eb141 100644 --- a/src/processes/atmosphere/prescribed_atmosphere.jl +++ b/src/processes/atmosphere/prescribed_atmosphere.jl @@ -1,23 +1,24 @@ """ Generic type representing the concentration of a particular tracer gas in the atmosphere. """ -@kwdef struct TracerGas{name} - TracerGas(name::Symbol) = new{name}() +@kwdef struct TracerGas{NF, name} + "Default concentration of the gas" + concentration::NF + + TracerGas(name::Symbol, concentration::NF) where {NF} = new{NF, name}(concentration) end -Base.nameof(::TracerGas{name}) where {name} = name +Base.nameof(::TracerGas{NF, name}) where {NF, name} = name -variables(::TracerGas{name}) where {name} = ( - input(name, XY(), default = default_tracer_conc(Val(name)), units = u"ppm", desc = "Ambient atmospheric $(name) concentration in ppm"), +variables(gas::TracerGas{NF, name}) where {NF, name} = ( + input(name, XY(), default = gas.concentration, units = u"ppm", desc = "Ambient atmospheric $(name) concentration in ppm"), ) -default_tracer_conc(::Val{:CO2}) = 380 # in ppm - """ Creates a `TracerGas` for ambient CO2 with concentration prescribed by an input variable with the given name. """ -AmbientCO2(name = :CO2) = TracerGas(name) +AmbientCO2(::Type{NF}, name = :CO2) where {NF} = TracerGas(name, NF(380)) """ Creates a `NamedTuple` from the given tracer gas types. @@ -49,7 +50,7 @@ struct PrescribedAtmosphere{ IncomingRad <: AbstractIncomingRadiation, Humidity <: AbstractHumidity, Aerodynamics <: AbstractAerodynamics, - Tracers <: NamedTuple{tracernames, <:Tuple{Vararg{TracerGas}}}, + Gases <: Tuple{Vararg{TracerGas{NF}}}, } <: AbstractAtmosphere{NF, Precip, IncomingRad, Humidity, Aerodynamics} "Surface-relative altitude in meters at which the atmospheric forcings are assumed to be applied" altitude::NF @@ -70,7 +71,7 @@ struct PrescribedAtmosphere{ aerodynamics::Aerodynamics "Atmospheric tracer gases" - tracers::Tracers + tracers::NamedTuple{tracernames, Gases} end function PrescribedAtmosphere( @@ -81,7 +82,7 @@ function PrescribedAtmosphere( radiation::AbstractIncomingRadiation = LongShortWaveRadiation(), humidity::AbstractHumidity = SpecificHumidity(), aerodynamics::AbstractAerodynamics = ConstantAerodynamics(NF), - tracers::NamedTuple = TracerGases(AmbientCO2()), + tracers::NamedTuple = TracerGases(AmbientCO2(NF)), ) where {NF} return PrescribedAtmosphere(altitude, min_windspeed, precip, radiation, humidity, aerodynamics, tracers) end @@ -225,8 +226,8 @@ length [hr]. struct LongShortWaveRadiation <: AbstractIncomingRadiation end variables(::LongShortWaveRadiation) = ( - input(:surface_shortwave_down, XY(), default = 300, units = u"W/m^2", desc = "Incoming (downwelling) shortwave solar radiation"), - input(:surface_longwave_down, XY(), default = 50, units = u"W/m^2", desc = "Incoming (downwelling) longwave thermal radiation"), + 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"), ) diff --git a/src/processes/processes.jl b/src/processes/processes.jl index 21fb4770a..88111abb5 100644 --- a/src/processes/processes.jl +++ b/src/processes/processes.jl @@ -124,5 +124,5 @@ export SurfaceHydrology include("surface_hydrology/surface_hydrology.jl") # Default debug hooks -@inline debughook!(::typeof(compute_auxiliary_kernel!), out, args...) = nancheck!(out) -@inline debughook!(::typeof(compute_tendencies_kernel!), out, args...) = nancheck!(out) +@inline debughook!(::typeof(compute_auxiliary_kernel!), out, args...) = checkfinite!(out) +@inline debughook!(::typeof(compute_tendencies_kernel!), out, args...) = checkfinite!(out) diff --git a/src/processes/soil/hydrology/soil_hydraulic_closures.jl b/src/processes/soil/hydrology/soil_hydraulic_closures.jl index 9ec5b19ad..88ee53c86 100644 --- a/src/processes/soil/hydrology/soil_hydraulic_closures.jl +++ b/src/processes/soil/hydrology/soil_hydraulic_closures.jl @@ -114,8 +114,10 @@ end # get inverse of SWRC inv_swrc = inv(get_swrc(hydrology)) por = porosity(i, j, k, grid, fields, strat, bgc) + sat_res = residual_saturation(get_hydraulic_properties(hydrology)) # compute matric pressure head - ψm = inv_swrc(sat * por; θsat = por) + ψm = inv_swrc(max(sat * por, sat_res); θsat = por) + @assert isfinite(ψm) "NaN/infinite matric potential ψm = $ψm with sat=$sat por=$por resid=$sat_res" # compute elevation pressure head ψz = z - z_ref # compute hydrostatic pressure head assuming impermeable lower boundary @@ -147,3 +149,7 @@ end i, j, k = @index(Global, NTuple) saturation_to_pressure!(out.pressure_head, i, j, k, grid, fields, closure, args...) end + +# Default debug hooks +@inline debughook!(::typeof(pressure_to_saturation_kernel!), out, args...) = checkfinite!(out) +@inline debughook!(::typeof(saturation_to_pressure_kernel!), out, args...) = checkfinite!(out) diff --git a/src/processes/soil/hydrology/soil_hydraulic_properties.jl b/src/processes/soil/hydrology/soil_hydraulic_properties.jl index b47cfc78f..e3ddce5ed 100644 --- a/src/processes/soil/hydrology/soil_hydraulic_properties.jl +++ b/src/processes/soil/hydrology/soil_hydraulic_properties.jl @@ -51,6 +51,13 @@ Compute the empirical field capacity of the soil. """ function field_capacity end +""" + $SIGNATURES + +Compute the (numerical) residual saturation level of the soil. +""" +function residual_saturation end + # Hydraulics parameterizations """ @@ -73,11 +80,14 @@ $TYPEDFIELDS "Hydraulic conductivity at saturation [m/s]" sat_hydraulic_cond::NF = 1.0e-5 - "Prescribed field capacity [-]" + "Constant field capacity [-]" field_capacity::NF = 0.25 - "Prescribed wilting point [-]" + "Constant wilting point [-]" wilting_point::NF = 0.05 + + "Residual (minimum) saturation level [-]" + residual::NF = 0.01 end function ConstantSoilHydraulics( @@ -96,6 +106,8 @@ end @inline field_capacity(hydraulics::ConstantSoilHydraulics, args...) = hydraulics.field_capacity +@inline residual_saturation(hydraulics::ConstantSoilHydraulics, args...) = hydraulics.residual + """ $TYPEDEF @@ -127,6 +139,9 @@ $TYPEDFIELDS "Exponent of field capacity adjustment due to clay content [-]" field_capacity_exp::NF = 0.35 + + "Residual (minimum) saturation level [-]" + residual::NF = 0.01 end function SoilHydraulicsSURFEX( @@ -142,6 +157,8 @@ end # TODO: this is not quite correct, SURFEX uses a hydraulic conductivity function that decreases exponentially with depth @inline saturated_hydraulic_conductivity(hydraulics::SoilHydraulicsSURFEX, args...) = hydraulics.sat_hydraulic_cond +@inline residual_saturation(hydraulics::SoilHydraulicsSURFEX, args...) = hydraulics.residual + @inline function wilting_point(hydraulics::SoilHydraulicsSURFEX, texture::SoilTexture) β_w = hydraulics.wilting_point_coef wp = β_w * sqrt(texture.clay * 100) @@ -171,10 +188,10 @@ function hydraulic_conductivity( hydraulics::AbstractSoilHydraulics{NF, RC, UnsatKLinear{NF}}, soil::SoilVolume ) where {NF, RC} - let fracs = volumetric_fractions(soil), - θw = fracs.water, # unfrozen water content - θsat = fracs.water + fracs.ice + fracs.air, # water + ice content at saturation (porosity) - K_sat = saturated_hydraulic_conductivity(hydraulics, soil.solid.texture) + let fracs = volumetric_fractions(soil) + θw = fracs.water # unfrozen water content + θsat = fracs.water + fracs.ice + fracs.air # water + ice content at saturation (porosity) + K_sat = saturated_hydraulic_conductivity(hydraulics, soil.solid.texture) K = K_sat * θw / θsat return K end @@ -205,15 +222,13 @@ function hydraulic_conductivity( soil::SoilVolume ) where {NF} # TODO: The SWRC parameters will need to also be spatially varying at some point - let n = hydraulics.swrc.n, # van Genuchten parameter `n` - fracs = volumetric_fractions(soil), - θw = fracs.water, # unfrozen water content - θwi = fracs.water + fracs.ice, # total water + ice content - θsat = porosity(soil), # porosity - f = liquid_fraction(soil), - Ω = hydraulics.unsat_hydraulic_cond.impedance, # scaling parameter for ice impedance - I_ice = 10^(-Ω * (1 - f)), # ice impedance factor - K_sat = saturated_hydraulic_conductivity(hydraulics, soil.solid.texture) + let n = hydraulics.swrc.n # van Genuchten parameter `n` + θw = water(soil) # volumetric content of unfrozen water + θsat = porosity(soil) # porosity + f = liquid_fraction(soil) # fraction of pore water that is unfrozen (equiv. θw / θwi) + Ω = hydraulics.unsat_hydraulic_cond.impedance # scaling parameter for ice impedance + I_ice = 10^(-Ω * (1 - f)) # ice impedance factor + K_sat = saturated_hydraulic_conductivity(hydraulics, soil.solid.texture) # We use `complex` types here to permit illegal state values which may occur when using adaptive time steppers. K = abs(K_sat * I_ice * sqrt(complex(θw / θsat)) * (1 - complex(1 - complex(θw / θsat)^(n / (n + 1)))^((n - 1) / n))^2) return K diff --git a/src/processes/soil/hydrology/soil_hydrology.jl b/src/processes/soil/hydrology/soil_hydrology.jl index b4709add7..82046c27a 100644 --- a/src/processes/soil/hydrology/soil_hydrology.jl +++ b/src/processes/soil/hydrology/soil_hydrology.jl @@ -140,27 +140,9 @@ end """ $TYPEDSIGNATURES -Kernel function that computes soil hydraulics and unsaturated hydraulic conductivity. +Kernel function that computes dynamic soil hydraulic properties. """ -@propagate_inbounds function compute_hydraulics!( - out, i, j, k, grid, fields, - hydrology::SoilHydrology, - strat::AbstractStratigraphy, - bgc::AbstractSoilBiogeochemistry - ) - # Get underlying grid - fgrid = get_field_grid(grid) - # compute hydraulic conductivity - @inbounds if k <= 1 - out.hydraulic_conductivity[i, j, k] = hydraulic_conductivity(i, j, 1, fgrid, fields, hydrology, strat, bgc) - elseif k >= fgrid.Nz - out.hydraulic_conductivity[i, j, k] = hydraulic_conductivity(i, j, fgrid.Nz, fgrid, fields, hydrology, strat, bgc) - out.hydraulic_conductivity[i, j, k + 1] = out.hydraulic_conductivity[i, j, k] - else - out.hydraulic_conductivity[i, j, k] = min_zᵃᵃᶠ(i, j, k, fgrid, hydraulic_conductivity, fields, hydrology, strat, bgc) - end - return nothing -end +compute_hydraulics!(out, i, j, k, grid, fields, hydrology::SoilHydrology, args...) = nothing """ $TYPEDSIGNATURES @@ -182,13 +164,16 @@ arising due to numerical error. This implementation scans over the saturation pr grid cell and redistributes excess water upward layer-by-layer until reaching the topmost layer, where any remaining excess water is added to the `surface_excess_water` pool. """ -@propagate_inbounds function adjust_saturation_profile!(out, i, j, grid, ::SoilHydrology{NF}) where {NF} - sat = out.saturation_water_ice - surface_excess_water = out.surface_excess_water +@propagate_inbounds function adjust_saturation_profile!(out, i, j, grid, hydrology::SoilHydrology{NF}) where {NF} + props = get_hydraulic_properties(hydrology) + sat_min = residual_saturation(props) field_grid = get_field_grid(grid) N = field_grid.Nz + sat = out.saturation_water_ice + surface_excess_water = out.surface_excess_water - # First iterate over soil layers from bottom to top + # First iterate over soil layers from bottom to top, transferring water from + # overfilled layers to the layer above for k in 1:(N - 1) # calculate excess saturation excess_sat = max(sat[i, j, k] - one(NF), zero(NF)) @@ -198,10 +183,10 @@ any remaining excess water is added to the `surface_excess_water` pool. sat[i, j, k + 1] += excess_sat * Δzᵃᵃᶜ(i, j, k, field_grid) / Δzᵃᵃᶜ(i, j, k + 1, field_grid) end - # then from top to bottom + # then from top to bottom, extracting water for underfilled cells from layers below for k in N:-1:2 - # calculate saturation deficit - deficit_sat = max(-sat[i, j, k], zero(NF)) + # calculate saturation deficit from residual saturation level + deficit_sat = max(-sat[i, j, k] + sat_min, zero(NF)) # add back saturation deficit and subtract from layer below sat[i, j, k] += deficit_sat sat[i, j, k - 1] -= deficit_sat * Δzᵃᵃᶜ(i, j, k, field_grid) / Δzᵃᵃᶜ(i, j, k - 1, field_grid) @@ -212,9 +197,9 @@ any remaining excess water is added to the `surface_excess_water` pool. sat[i, j, N] -= excess_sat surface_excess_water[i, j, 1] += excess_sat * Δzᵃᵃᶜ(i, j, N, field_grid) - # If the lowermost (bottom) layer has a deficit, just set to zero. + # If the lowermost (bottom) layer has a deficit, just set to the residual saturation level. # This constitutes a mass balance violation but should not happen under realistic conditions. - sat[i, j, 1] = max(sat[i, j, 1], zero(NF)) + sat[i, j, 1] = max(sat[i, j, 1], sat_min) return nothing end @@ -249,10 +234,9 @@ the porosity and is thus not the same as the saturation tendency. evtr::Optional{AbstractEvapotranspiration} ) where {NF} # Compute divergence of water fluxes due to forcings only - ∂θ∂t = ( - + forcing(i, j, k, grid, clock, fields, evtr, hydrology, constants) # ET forcing - + forcing(i, j, k, grid, clock, fields, hydrology.vwc_forcing, hydrology) # generic user-defined forcing - ) + ET_loss = forcing(i, j, k, grid, clock, fields, evtr, hydrology, constants) # ET forcing + F = forcing(i, j, k, grid, clock, fields, hydrology.vwc_forcing, hydrology) # generic user-defined forcing + ∂θ∂t = ET_loss + F return ∂θ∂t end @@ -298,3 +282,8 @@ end i, j, k = @index(Global, NTuple) compute_hydraulics!(out, i, j, k, grid, fields, hydrology, args...) end + +# Default debug hooks +@inline debughook!(::typeof(compute_water_table_kernel!), out, args...) = checkfinite!(out) +@inline debughook!(::typeof(adjust_saturation_profile_kernel!), out, args...) = checkfinite!(out) +@inline debughook!(::typeof(compute_hydraulics_kernel!), out, args...) = checkfinite!(out) diff --git a/src/processes/soil/hydrology/soil_hydrology_rre.jl b/src/processes/soil/hydrology/soil_hydrology_rre.jl index 021f31250..0cd3fe48f 100644 --- a/src/processes/soil/hydrology/soil_hydrology_rre.jl +++ b/src/processes/soil/hydrology/soil_hydrology_rre.jl @@ -145,6 +145,31 @@ Compute the hydraulic conductivity at the center of the grid cell `i, j, k`. return hydraulic_conductivity(hydrology.hydraulic_properties, soil) end +""" + $TYPEDSIGNATURES + +Computes the unsaturated hydraulic conductivity for `RichardsEq` configurations of `SoilHydrology`. +""" +@propagate_inbounds function compute_hydraulics!( + out, i, j, k, grid, fields, + hydrology::SoilHydrology{NF, RichardsEq}, + strat::AbstractStratigraphy, + bgc::AbstractSoilBiogeochemistry + ) where {NF} + # Get underlying grid + fgrid = get_field_grid(grid) + # compute hydraulic conductivity + @inbounds if k <= 1 + out.hydraulic_conductivity[i, j, k] = hydraulic_conductivity(i, j, 1, fgrid, fields, hydrology, strat, bgc) + elseif k >= fgrid.Nz + out.hydraulic_conductivity[i, j, k] = hydraulic_conductivity(i, j, fgrid.Nz, fgrid, fields, hydrology, strat, bgc) + out.hydraulic_conductivity[i, j, k + 1] = out.hydraulic_conductivity[i, j, k] + else + out.hydraulic_conductivity[i, j, k] = min_zᵃᵃᶠ(i, j, k, fgrid, hydraulic_conductivity, fields, hydrology, strat, bgc) + end + return nothing +end + # Kernels @kernel inbounds = true function compute_tendencies_kernel!( diff --git a/src/processes/soil/stratigraphy/soil_volume.jl b/src/processes/soil/stratigraphy/soil_volume.jl index 03f7c7e47..0590fd8a6 100644 --- a/src/processes/soil/stratigraphy/soil_volume.jl +++ b/src/processes/soil/stratigraphy/soil_volume.jl @@ -30,17 +30,19 @@ $FIELDS end end -porosity(soil::SoilVolume) = soil.porosity +@inline porosity(soil::SoilVolume) = soil.porosity -saturation(soil::SoilVolume) = soil.saturation +@inline saturation(soil::SoilVolume) = soil.saturation -liquid_fraction(soil::SoilVolume) = soil.liquid +@inline liquid_fraction(soil::SoilVolume) = soil.liquid -water_ice(soil::SoilVolume) = soil.porosity * soil.saturation +@inline water_ice(soil::SoilVolume) = soil.porosity * soil.saturation -organic_fraction(soil::SoilVolume) = organic_fraction(soil.solid) +@inline water(soil::SoilVolume) = soil.liquid * soil.porosity * soil.saturation -mineral_texture(soil::SoilVolume) = mineral_texture(soil.solid) +@inline organic_fraction(soil::SoilVolume) = organic_fraction(soil.solid) + +@inline mineral_texture(soil::SoilVolume) = mineral_texture(soil.solid) """ $TYPEDSIGNATURES @@ -92,6 +94,8 @@ Alias for `SoilVolume{T, MineralOrganic{T}}` """ const MineralOrganicSoil{NF} = SoilVolume{NF, MineralOrganic{NF}} +@inline mineral_texture(solid::MineralOrganic) = solid.texture + @inline organic_fraction(solid::MineralOrganic) = solid.organic """ diff --git a/src/processes/surface_energy/radiative_fluxes.jl b/src/processes/surface_energy/radiative_fluxes.jl index 348db8b8b..6f5f44df9 100644 --- a/src/processes/surface_energy/radiative_fluxes.jl +++ b/src/processes/surface_energy/radiative_fluxes.jl @@ -78,27 +78,46 @@ struct DiagnosedRadiativeFluxes{NF} <: AbstractRadiativeFluxes{NF} end DiagnosedRadiativeFluxes(::Type{NF}) where {NF} = DiagnosedRadiativeFluxes{NF}() """ - compute_shortwave_up(::DiagnosedRadiativeFluxes, surface_shortwave_down, α) + compute_shortwave_up(::DiagnosedRadiativeFluxes, S_down, α) -Compute outgoing shortwave radiation from the incoming `surface_shortwave_down` and albedo `α`. +Compute outgoing shortwave radiation from the incoming shortwave radiation `S_down` and albedo `α`. """ -@inline function compute_shortwave_up(::DiagnosedRadiativeFluxes, surface_shortwave_down, α) - surface_shortwave_up = α * surface_shortwave_down - return surface_shortwave_up +@inline function compute_shortwave_up(::DiagnosedRadiativeFluxes, S_down, α) + S_up = α * S_down + return S_up end """ - compute_longwave_up(::DiagnosedRadiativeFluxes, constants::PhysicalConstants, surface_longwave_down, Ts, ϵ) + compute_longwave_up(::DiagnosedRadiativeFluxes, constants::PhysicalConstants, L_down, Ts, ϵ) -Compute outgoing longwave radiation from incoming `surface_longwave_down`, surface temperature `Ts`, and emissivity `ϵ`. +Compute outgoing longwave radiation from incoming longwave radiation `L_down`, surface temperature `Ts`, and emissivity `ϵ`. """ -@inline function compute_longwave_up(::DiagnosedRadiativeFluxes, constants::PhysicalConstants, surface_longwave_down, Ts, ϵ) +@inline function compute_longwave_up(::DiagnosedRadiativeFluxes, constants::PhysicalConstants, L_down, Ts, ϵ) T = celsius_to_kelvin(constants, Ts) + L_emit = stefan_boltzmann(constants, T, ϵ) # outgoing LW radiation is the sum of the radiant emittance and the residual incoming radiation - surface_longwave_up = stefan_boltzmann(constants, T, ϵ) + (1 - ϵ) * surface_longwave_down + surface_longwave_up = L_emit + (1 - ϵ) * L_down return surface_longwave_up end +""" + $TYPEDEF + +Compute the net radiation budget given incoming and outgoing shortwave and longwave radiation. +""" +@inline function compute_surface_net_radiation( + ::AbstractRadiativeFluxes, + SW_up, + SW_down, + LW_up, + LW_down + ) + # Sum up radiation fluxes; note that by convention fluxes are positive upward + rad_net = SW_up - SW_down + LW_up - LW_down + return rad_net +end + + ## Top-level interface methods variables(::DiagnosedRadiativeFluxes) = ( @@ -189,25 +208,6 @@ end return surface_net_radiation end -# Common implementations - -""" - $TYPEDEF - -Compute the net radiation budget given incoming and outgoing shortwave and longwave radiation. -""" -@inline function compute_surface_net_radiation( - ::AbstractRadiativeFluxes, - SW_up, - SW_down, - LW_up, - LW_down - ) - # Sum up radiation fluxes; note that by convention fluxes are positive upward - rad_net = SW_up - SW_down + LW_up - LW_down - return rad_net -end - ## Kernels @kernel inbounds = true function compute_auxiliary_kernel!(out, grid, fields, rad::AbstractRadiativeFluxes, args...) diff --git a/src/processes/surface_energy/skin_temperature.jl b/src/processes/surface_energy/skin_temperature.jl index 69e38f01b..b9db87a0b 100644 --- a/src/processes/surface_energy/skin_temperature.jl +++ b/src/processes/surface_energy/skin_temperature.jl @@ -75,7 +75,7 @@ sensible `H_s` and latent `H_l` heat flux. """ @inline function compute_ground_heat_flux(::AbstractSkinTemperature, R_net, H_s, H_l) # Compute ground heat flux as the residual of R_net and turbulent fluxes - G = R_net - H_s - H_l + G = R_net + H_s + H_l return G end diff --git a/src/processes/surface_energy/turbulent_fluxes.jl b/src/processes/surface_energy/turbulent_fluxes.jl index 0c8f44ecb..5c49d7a07 100644 --- a/src/processes/surface_energy/turbulent_fluxes.jl +++ b/src/processes/surface_energy/turbulent_fluxes.jl @@ -42,10 +42,10 @@ end $TYPEDSIGNATURES Compute the latent heat flux as a function of the humidity flux `Q_h` [m/s], the density `ρₐ` [kg/m³] of air, -and the specific latent heat of fusion `Lsl` [J/kg]. +and the specific latent heat of vaporization or sublimation `L` [J/kg]. """ -function compute_latent_heat_flux(::DiagnosedTurbulentFluxes, Q_h, ρₐ, Lsl) - Hₗ = Lsl * ρₐ * Q_h +function compute_latent_heat_flux(::DiagnosedTurbulentFluxes, Q_h, ρₐ, L) + Hₗ = L * ρₐ * Q_h return Hₗ end @@ -139,12 +139,12 @@ Compute the sensible heat flux at `i, j` based on the current skin temperature a constants::PhysicalConstants, atmos::AbstractAtmosphere ) - let ρₐ = constants.ρₐ, # density of air - cₐ = constants.cₐ, # specific heat capacity of air - rₐ = aerodynamic_resistance(i, j, grid, fields, atmos), # aerodynamic resistance - Tₛ = skin_temperature(i, j, grid, fields, skinT), # skin temperature - Tₐ = air_temperature(i, j, grid, fields, atmos), # air temperature - Q_T = (Tₛ - Tₐ) / rₐ # bulk aerodynamic temperature-gradient + let ρₐ = constants.ρₐ # density of air + cₐ = constants.cₐ # specific heat capacity of air + rₐ = aerodynamic_resistance(i, j, grid, fields, atmos) # aerodynamic resistance + Tₛ = skin_temperature(i, j, grid, fields, skinT) # skin temperature + Tₐ = air_temperature(i, j, grid, fields, atmos) # air temperature + Q_T = (Tₛ - Tₐ) / rₐ # bulk aerodynamic temperature-gradient # Calculate sensible heat flux (positive upwards) Hₛ = compute_sensible_heat_flux(tur, Q_T, ρₐ, cₐ) return Hₛ @@ -165,12 +165,12 @@ to the latent heat flux. constants::PhysicalConstants, atmos::AbstractAtmosphere ) - let L = constants.Llg, # specific latent heat of vaporization of water - ρₐ = constants.ρₐ, # density of air - Tₛ = skin_temperature(i, j, grid, fields, skinT), - rₐ = aerodynamic_resistance(i, j, grid, fields, atmos), # aerodynamic resistance - Δq = compute_specific_humidity_difference(i, j, grid, fields, atmos, constants, Tₛ), - Q_h = Δq / rₐ # humidity flux + let L = constants.Llg # specific latent heat of vaporization of water + ρₐ = constants.ρₐ # density of air + Tₛ = skin_temperature(i, j, grid, fields, skinT) + rₐ = aerodynamic_resistance(i, j, grid, fields, atmos) # aerodynamic resistance + Δq = compute_specific_humidity_difference(i, j, grid, fields, atmos, constants, Tₛ) + Q_h = Δq / rₐ # humidity flux # Calculate latent heat flux (positive upwards) Hₗ = compute_latent_heat_flux(tur, Q_h, ρₐ, L) return Hₗ @@ -190,9 +190,9 @@ defined by `evtr` which is assumed to be already computed. evtr::AbstractEvapotranspiration, constants::PhysicalConstants ) - let L = constants.Llg, # specific latent heat of vaporization of water - ρₐ = constants.ρₐ, # density of air - Q_h = surface_humidity_flux(i, j, grid, fields, evtr) # humidity flux + let L = constants.Llg # specific latent heat of vaporization of water + ρₐ = constants.ρₐ # density of air + Q_h = surface_humidity_flux(i, j, grid, fields, evtr) # humidity flux # Calculate latent heat flux (positive upwards) Hₗ = compute_latent_heat_flux(tur, Q_h, ρₐ, L) return Hₗ @@ -201,32 +201,10 @@ end # Kernels -@kernel function compute_auxiliary_kernel!( - out, grid, fields, - tur::DiagnosedTurbulentFluxes, - skinT::AbstractSkinTemperature, - constants::PhysicalConstants, - atmos::AbstractAtmosphere - ) - 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, skinT, constants, atmos) - # compute latent heat flux - out.latent_heat_flux[i, j, 1] = compute_latent_heat_flux(i, j, grid, fields, tur, skinT, constants, atmos) -end - -# TODO: Can these dispatches be standardized to reduce redundancy? -@kernel function compute_auxiliary_kernel!( - out, grid, fields, - tur::DiagnosedTurbulentFluxes, - skinT::AbstractSkinTemperature, - constants::PhysicalConstants, - atmos::AbstractAtmosphere, - evtr::AbstractEvapotranspiration - ) +@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, skinT, constants, atmos) + out.sensible_heat_flux[i, j, 1] = compute_sensible_heat_flux(i, j, grid, fields, tur, args...) # compute latent heat flux - out.latent_heat_flux[i, j, 1] = compute_latent_heat_flux(i, j, grid, fields, tur, evtr, constants) + out.latent_heat_flux[i, j, 1] = compute_latent_heat_flux(i, j, grid, fields, tur, args...) end diff --git a/src/processes/surface_hydrology/canopy_interception/canopy_interception.jl b/src/processes/surface_hydrology/canopy_interception/canopy_interception.jl index e5d274094..c6e9643ed 100644 --- a/src/processes/surface_hydrology/canopy_interception/canopy_interception.jl +++ b/src/processes/surface_hydrology/canopy_interception/canopy_interception.jl @@ -61,7 +61,7 @@ variables(::PALADYNCanopyInterception) = ( 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(:SAI, XY(); desc = "Stem Area Index", units = u"m^2/m^2"), + input(:stem_area_index, 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] @@ -181,7 +181,7 @@ end # Get inputs rain = rainfall(i, j, grid, fields, atmos) LAI = fields.leaf_area_index[i, j] - SAI = fields.SAI[i, j] + SAI = fields.stem_area_index[i, j] w_can = fields.canopy_water[i, j] # Compute canopy saturation faction diff --git a/src/processes/surface_hydrology/evapotranspiration/bare_ground_evaporation.jl b/src/processes/surface_hydrology/evapotranspiration/bare_ground_evaporation.jl index 9eaf3bae4..658d65bfd 100644 --- a/src/processes/surface_hydrology/evapotranspiration/bare_ground_evaporation.jl +++ b/src/processes/surface_hydrology/evapotranspiration/bare_ground_evaporation.jl @@ -16,7 +16,7 @@ end BareGroundEvaporation( ::Type{NF}; - ground_resistance::GR = ConstantEvaporationResistanceFactor(one(NF)) + ground_resistance::GR = SoilMoistureResistanceFactor(NF) ) where {NF, GR} = BareGroundEvaporation{NF, GR}(ground_resistance) @propagate_inbounds surface_humidity_flux(i, j, grid, fields, evaporation::BareGroundEvaporation, args...) = fields.evaporation_ground[i, j] @@ -58,5 +58,6 @@ end β = ground_evaporation_resistance_factor(i, j, grid, fields, evaporation.ground_resistance, soil) Δq = compute_specific_humidity_difference(i, j, grid, fields, atmos, constants, Ts) # Calculate water evaporation flux in m/s (positive upwards) - return out.evaporation_ground[i, j, 1] = β * Δq / rₐ + E_g = β * Δq / rₐ + return out.evaporation_ground[i, j, 1] = E_g end diff --git a/src/processes/surface_hydrology/evapotranspiration/canopy_evapotranspiration.jl b/src/processes/surface_hydrology/evapotranspiration/canopy_evapotranspiration.jl index 18f17f7fc..67205ec2c 100644 --- a/src/processes/surface_hydrology/evapotranspiration/canopy_evapotranspiration.jl +++ b/src/processes/surface_hydrology/evapotranspiration/canopy_evapotranspiration.jl @@ -138,7 +138,7 @@ for the given scheme `evapotranspiration` and process dependencies. ) # Get inputs Ts = fields.skin_temperature[i, j] # skin temperature (top of canopy) - Tg = fields.ground_temperature[i, j] # ground temeprature (top snow/soil layer) + Tg = fields.ground_temperature[i, j] # ground temperature (top snow/soil layer) gw_can = fields.canopy_water_conductance[i, j] # stomatal conductance (assumed to be defined by vegetation) # Compute VPD and resistance terms @@ -164,11 +164,11 @@ Compute the aerodynamic resistance between the ground and canopy as a function o @inline function aerodynamic_resistance( i, j, grid, fields, atmos::AbstractAtmosphere, - evapotranspiration::PALADYNCanopyEvapotranspiration, + evapotranspiration::PALADYNCanopyEvapotranspiration{NF}, ::AbstractVegetation # included just to make explicit the dependence on vegetation fields - ) - @inbounds let LAI = fields.leaf_area_index[i, j], - SAI = fields.SAI[i, j], + ) where {NF} + @inbounds let LAI = max(fields.leaf_area_index[i, j], zero(NF)), + SAI = max(fields.stem_area_index[i, j], zero(NF)), Vₐ = windspeed(i, j, grid, fields, atmos), C = evapotranspiration.C_can # drag coefficient for the canopy rₙ = (1 - exp(-LAI - SAI)) / (C * Vₐ) diff --git a/src/processes/surface_hydrology/evapotranspiration/evapotranspiration_base.jl b/src/processes/surface_hydrology/evapotranspiration/evapotranspiration_base.jl index 1683b1ab7..ac5da9037 100644 --- a/src/processes/surface_hydrology/evapotranspiration/evapotranspiration_base.jl +++ b/src/processes/surface_hydrology/evapotranspiration/evapotranspiration_base.jl @@ -7,8 +7,8 @@ Compute and return the evapotranspiration forcing for soil moisture at the given The ET forcing is just the `surface_humidity_flux` rescaled by the thickness of layer `k`. """ @inline function forcing(i, j, k, grid, clock, fields, evapotranspiration::AbstractEvapotranspiration, ::AbstractSoilHydrology) - let Δz = Δzᵃᵃᶜ(i, j, k, grid), - Qh = surface_humidity_flux(i, j, grid, fields, evapotranspiration) + let Δz = Δzᵃᵃᶜ(i, j, k, grid) + Qh = surface_humidity_flux(i, j, grid, fields, evapotranspiration) ∂θ∂t = -Qh / Δz # rescale by layer thickness to get water content flux return ∂θ∂t * (k == grid.Nz) end diff --git a/src/processes/surface_hydrology/evapotranspiration/ground_resistance_factor.jl b/src/processes/surface_hydrology/evapotranspiration/ground_resistance_factor.jl index 8659c6c1f..22f4538d5 100644 --- a/src/processes/surface_hydrology/evapotranspiration/ground_resistance_factor.jl +++ b/src/processes/surface_hydrology/evapotranspiration/ground_resistance_factor.jl @@ -38,18 +38,27 @@ ground_evaporation_resistance_factor(i, j, grid, fields, ::SoilMoistureResistanc @inline function ground_evaporation_resistance_factor( i, j, grid, fields, - ::SoilMoistureResistanceFactor{NF}, + res::SoilMoistureResistanceFactor{NF}, soil::AbstractSoil ) where {NF} fgrid = get_field_grid(grid) strat = get_stratigraphy(soil) hydrology = get_hydrology(soil) bgc = get_biogeochemistry(soil) + props = get_hydraulic_properties(hydrology) soil = soil_volume(i, j, fgrid.Nz, grid, fields, strat, hydrology, bgc) - fc = field_capacity(get_hydraulic_properties(hydrology), soil) + texture = mineral_texture(soil) fracs = volumetric_fractions(soil) - if fracs.water < fc - β = (1 - cos(π * fracs.water / fc))^2 / 4 + # Get field capacity, water content, and residual water content + θfc = field_capacity(get_hydraulic_properties(hydrology), texture) + θw = fracs.water + θres = props.residual * soil.saturation + return ground_evaporation_resistance_factor(res, θw, θfc, θres) +end + +@inline function ground_evaporation_resistance_factor(::SoilMoistureResistanceFactor{NF}, θw, θfc, θres) where {NF} + if θw < θfc + β = (1 - cos(π * (θw - θres) / (θfc - θres)))^2 / 4 else β = NF(1) end diff --git a/src/state_variables.jl b/src/state_variables.jl index 6677acd24..f59540159 100644 --- a/src/state_variables.jl +++ b/src/state_variables.jl @@ -74,9 +74,10 @@ function Oceananigans.TimeSteppers.update_state!(state::StateVariables, model::A update_inputs!(state, inputs) fill_halo_regions!(state) compute_auxiliary!(state, model) - return if compute_tendencies + if compute_tendencies compute_tendencies!(state, model) end + return nothing end """ diff --git a/src/timesteppers/abstract_timestepper.jl b/src/timesteppers/abstract_timestepper.jl index 6613bacea..c1d11058f 100644 --- a/src/timesteppers/abstract_timestepper.jl +++ b/src/timesteppers/abstract_timestepper.jl @@ -67,8 +67,12 @@ function explicit_step!(state, grid::AbstractLandGrid, timestepper::AbstractTime fastiterate(keys(state.prognostic)) do name # apply flux BCs, if present compute_z_bcs!(state.tendencies[name], state.prognostic[name], grid, state) + # debug site post-BC + debugsite!(explicit_step!, state.tendencies[name], name) # update prognostic state variable explicit_step!(state.prognostic[name], state.tendencies[name], grid, timestepper, Δt) + # debug site post-step + debugsite!(explicit_step!, state.prognostic[name], name) end fastiterate(state.namespaces) do ns explicit_step!(ns, grid, timestepper, Δt) @@ -139,3 +143,6 @@ end u[i, j, 1] += ∂u∂t[i, j] * Δt end end + +# Default debug hooks +@inline debughook!(::typeof(explicit_step!), field, name) = checkfinite!(field, name) diff --git a/test/benchmarks/gpu/soil_heat_hydrology_global.jl b/test/benchmarks/gpu/soil_heat_hydrology_global.jl index 844f006ef..d8e8d20b3 100644 --- a/test/benchmarks/gpu/soil_heat_hydrology_global.jl +++ b/test/benchmarks/gpu/soil_heat_hydrology_global.jl @@ -83,23 +83,54 @@ exit() # Interactive plotting using DataFrames, CSV -using GLMakie, GeoMakie +using CairoMakie, GeoMakie +using CairoMakie.GeometryBasics +using RingGrids -GLMakie.activate!(inline = true) +# GLMakie.activate!(inline = true) cpu_data = DataFrame(CSV.File("outputs/benchmarks/soil_heat_hydrology_benchmark_cpu_nthreads=32.csv")) gpu_data = DataFrame(CSV.File("outputs/benchmarks/soil_heat_hydrology_benchmark_gpu_nthreads=32.csv")) -let fig = Figure(size = (800, 400)), - ax = Axis(fig[1, 1], title = "Terrarium.jl GPU vs. CPU scaling", ylabel = "log simulated years per day (SYPD)", xlabel = "log number of grid cells (Nₕ)", xscale = log10, yscale = log10) - # data is in ms / sim hr x 1 d / (1000*24*3600 ms) x 24 sim hr / sim day -> - cpu_sypd = 1000 * 24 * 3600 ./ (24 * cpu_data.mid_time) - gpu_sypd = 1000 * 24 * 3600 ./ (24 * gpu_data.mid_time) - scatterlines!(ax, cpu_data.npoints, cpu_sypd, label = "CPU") - scatterlines!(ax, gpu_data.npoints, gpu_sypd, label = "GPU") - axislegend(ax) - Makie.save("plots/terrarium_cpu_vs_gpu_scaling.svg", fig) - fig +ring_grid_5deg = FullGaussianGrid(14) +ring_grid_3deg = FullGaussianGrid(24) +ring_grid_1deg = FullGaussianGrid(72) +ring_grid_halfdeg = FullGaussianGrid(72 * 2) +ring_grid_qtrdeg = FullGaussianGrid(72 * 4) +ring_grid_10km = FullGaussianGrid(72 * 10) +ring_grid_1km = FullGaussianGrid(72 * 100) +ref_grids = [ring_grid_5deg, ring_grid_1deg, ring_grid_halfdeg, ring_grid_qtrdeg, ring_grid_10km] +ref_npoints = [get_npoints(rg) for rg in ref_grids] + +Makie.with_theme(fontsize = 18) do + let fig = Figure(size = (800, 400)) + ax = Axis( + fig[1, 1], + title = "Terrarium.jl GPU vs. CPU scaling", + ylabel = "Simulated years per day (SYPD)", + xlabel = "Number of grid cells", + xscale = log10, + yscale = log10 + ) + # data is in ms / sim hr x 1 d / (1000*24*3600 ms) x 24 sim hr / sim day -> + cpu_sypd = 1000 * 24 * 3600 ./ (24 * cpu_data.mid_time) + gpu_sypd = 1000 * 24 * 3600 ./ (24 * gpu_data.mid_time) + scatterlines!(ax, cpu_data.npoints, cpu_sypd, label = "CPU") + scatterlines!(ax, gpu_data.npoints, gpu_sypd, label = "GPU") + # Draw vertical reference lines and place a short annotation next to each one. + # Choose human-friendly labels for the reference grids. + ref_labels = ["5°", "1°", "0.5°", "0.25°", "0.1°"] + ys_max = maximum(vcat(cpu_sypd, gpu_sypd)) + for (x, lbl) in zip(ref_npoints, ref_labels) + vlines!(ax, [x], linestyle = :dash, color = :black) + # place label slightly to the right of the line at the top of the plotted range + text!(ax, GeometryBasics.Point2(x * 1.2, 200); text = lbl, align = (:left, :top), fontsize = 18) + end + axislegend(ax) + Makie.save("plots/terrarium_cpu_vs_gpu_scaling.svg", fig) + Makie.save("plots/terrarium_cpu_vs_gpu_scaling.pdf", fig) + fig + end end using SpeedyWeather, CUDA diff --git a/test/soil/soil_hydrology_tests.jl b/test/soil/soil_hydrology_tests.jl index 0455ad0c0..a11b49f69 100644 --- a/test/soil/soil_hydrology_tests.jl +++ b/test/soil/soil_hydrology_tests.jl @@ -6,7 +6,7 @@ using FreezeCurves using Oceananigans @testset "Hydraulic properties (constant)" begin - # For prescribed hyraulic properties, just check that the returned values + # For prescribed hydraulic properties, just check that the returned values # match what was set. hydraulic_props = ConstantSoilHydraulics( Float64; @@ -113,13 +113,13 @@ end compute!(∫sat_deficit) Terrarium.adjust_saturation_profile!(state, grid, hydrology) ∫sat₁ = compute!(∫sat)[1, 1, 1] - @test all(state.saturation_water_ice .>= 0) - @test all(∫sat₁ - ∫sat₀ .≈ 0) + @test all(state.saturation_water_ice .>= hydraulic_properties.residual) + @test ∫sat₁ ≈ ∫sat₀ # Case 3: Completely dry with negative saturation near surface set!(state.saturation_water_ice, (x, z) -> min(-0.1 - z, 0.0)) Terrarium.adjust_saturation_profile!(state, grid, hydrology) - @test all(state.saturation_water_ice .≈ 0) + @test all(state.saturation_water_ice .≈ hydraulic_properties.residual) end @testset "SoilHydrology: Richardson-Richards' equation" begin diff --git a/test/surface_energy/skin_temperature.jl b/test/surface_energy/skin_temperature.jl index 79fbe1f20..4822273b5 100644 --- a/test/surface_energy/skin_temperature.jl +++ b/test/surface_energy/skin_temperature.jl @@ -27,21 +27,51 @@ end set!(state.specific_humidity, 0.002) # dry conditions set!(state.air_pressure, 101_325) # standard pressure set!(state.air_temperature, 10.0) # 10 °C - set!(state.ground_temperature, 2.0) # 2 °C + set!(state.ground_temperature, 5.0) # 5 °C set!(state.windspeed, 1.0) # 1 m/s - compute_auxiliary!(state, grid, skin_temperature, seb) + compute_auxiliary!(state, model) @test all(isfinite.(state.skin_temperature)) + @test all(state.sensible_heat_flux .< 0) # check that skin temperature converges for a large number of iterations Tskin_old = deepcopy(state.skin_temperature) resid = nothing - for i in 1:5 + balance = nothing + for i in 1:10 # compute fluxes Terrarium.compute_surface_energy_fluxes!(state, grid, seb, model.constants, model.atmosphere) # diagnose skin temperature Terrarium.update_skin_temperature!(state, grid, seb.skin_temperature) resid = maximum(abs.(state.skin_temperature - Tskin_old)) + balance = state.surface_net_radiation + state.latent_heat_flux + state.sensible_heat_flux - state.ground_heat_flux Tskin_old = deepcopy(state.skin_temperature) - println("skin temperature residual at iteration $i: $resid") + println("skin temperature at iteration $i: $(state.skin_temperature[1, 1]) residual: $resid energy balance: $(balance[1, 1, 1])") end @test all(resid .< sqrt(eps())) + @test all(abs.(balance) .< sqrt(eps())) + # Cloudy and wet + set!(state.surface_shortwave_down, 150.0) + set!(state.surface_longwave_down, 50.0) + set!(state.specific_humidity, 0.01) + set!(state.air_pressure, 101_325) # standard pressure + set!(state.air_temperature, 5.0) # 5 °C + set!(state.ground_temperature, 10.0) # 10 °C + set!(state.windspeed, 10.0) # 10 m/s + set!(state.skin_temperature, 10.0) # initial skin temperature = ground temperature + compute_auxiliary!(state, model) + @test all(isfinite.(state.skin_temperature)) + @test all(state.sensible_heat_flux .> 0) + # check that skin temperature converges for a large number of iterations + Tskin_old = deepcopy(state.skin_temperature) + for i in 1:10 + # compute fluxes + Terrarium.compute_surface_energy_fluxes!(state, grid, seb, model.constants, model.atmosphere) + # diagnose skin temperature + Terrarium.update_skin_temperature!(state, grid, seb.skin_temperature) + resid = maximum(abs.(state.skin_temperature - Tskin_old)) + balance = state.surface_net_radiation + state.latent_heat_flux + state.sensible_heat_flux - state.ground_heat_flux + Tskin_old = deepcopy(state.skin_temperature) + println("skin temperature at iteration $i: $(state.skin_temperature[1, 1]) residual: $resid energy balance: $(balance[1, 1, 1])") + end + @test all(resid .< sqrt(eps())) + @test all(abs.(balance) .< sqrt(eps())) end