Update speedy primitive dry coupling script

Cleaned up to only use surface temperature coupling (radiative fluxes are ignored)

Author: bgroenks96

examples/Project.toml

@@ -2,7 +2,6 @@
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
-GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
@@ -16,8 +15,8 @@ SpeedyWeatherInternals = "34489162-d270-4603-8b96-37b04f830d73"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Terrarium = "80418d68-07fa-499d-ae2b-c12e531f5cd8"
 
-[sources]
-Terrarium = {path = ".."}
-
 [compat]
 SpeedyWeatherInternals = "0.1.2"
+
+[sources.Terrarium]
+path = ".."

examples/speedy_dry_land.jl

@@ -10,19 +10,22 @@ using CairoMakie, GeoMakie
 import RingGrids
 import SpeedyWeather as Speedy
 
+# Temporary workaround for SpeedyWeather.jl#919
+Speedy.SpeedyTransforms.FFTW.set_num_threads(1)
+
 """
 Naive implementation of a SpeedyWeather "dry" land model that couples to a Terrarium model.
 """
-struct TerrariumDryLand{NF, TM<:ModelState{NF}} <: Speedy.AbstractDryLand
+struct TerrariumDryLand{NF, TMI<:Terrarium.ModelIntegrator{NF}} <: Speedy.AbstractDryLand
     "Speedy spectral grid"
     spectral_grid::Speedy.SpectralGrid
 
-    "Initialized Terrarium model"
-    model::TM
+    "Initialized Terrarium model integrator"
+    integrator::TMI
 
-    function TerrariumDryLand(model::ModelState{NF, Arch, Grid}; spectral_grid_kwargs...) where {NF, Arch, Grid<:ColumnRingGrid}
-        spectral_grid = Speedy.SpectralGrid(model.grid.rings; NF, nlayers_soil=1, spectral_grid_kwargs...)
-        return new{eltype(model), typeof(model)}(spectral_grid, model)
+    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)
     end
 end
 
@@ -33,9 +36,9 @@ function Speedy.initialize!(
     land::TerrariumDryLand,
     model::Speedy.PrimitiveEquation,
 ) where {NF}
-    Tsoil = interior(land.model.state.temperature)[:, 1, end]
+    Tsoil = interior(land.integrator.state.temperature)[:, 1, end]
     progn_land.soil_temperature .= Tsoil .+ NF(273.15)
-    Terrarium.initialize!(land.model)
+    Terrarium.initialize!(land.integrator)
     return nothing
 end
 
@@ -47,130 +50,75 @@ function Speedy.timestep!(
 ) where {NF}
     # get speedy state variables
     Tair = @view diagn.grid.temp_grid[:, end]
-    pres = diagn.grid.pres_grid
-    SwIn = diagn.physics.surface_shortwave_down
-    LwIn = diagn.physics.surface_longwave_down
-    SwOut = diagn.physics.surface_shortwave_up
-    LwOut = diagn.physics.surface_longwave_up
-    H = diagn.physics.sensible_heat_flux
-    # update Terrarium input fields
-    inputs = land.model.state.inputs
-    set!(inputs.air_temperature, Tair) # first set directly to avoid allocating new fields
-    set!(inputs.air_temperature, inputs.air_temperature - 273.15) # then convert to celsius
-    # set!(inputs.air_pressure, pres) # similar with pressure
-    # set!(inputs.air_pressure, exp(inputs.air_pressure)) # take exp to get pressure in Pa
-    # set!(inputs.surface_shortwave_down, SwIn)
-    # set!(inputs.surface_longwave_down, LwIn)
-    # set!(inputs.surface_shortwave_up, SwOut)
-    # set!(inputs.surface_longwave_up, LwOut)
-    # set!(inputs.sensible_heat_flux, H)
-    # do land timestep
-    Terrarium.timestep!(land.model, Dates.second(progn.clock.Δt))
-    # SEB coupling:
-    # Ts = land.model.state.skin_temperature
-    # progn.land.soil_temperature .= Ts .+ NF(273.15)
-    # progn.land.sensible_heat_flux .= land.model.state.sensible_heat_flux
-    # Soil surface temperature coupling:
-    Tsoil = land.model.state.temperature
+    # terrarium state
+    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
+    # 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)
+    # Get soil temperatures
+    Tsoil = state.temperature
+    # Get surface temperature (last z-layer in Oceananigans grids)
     Nx, Nz = Tsoil.grid.Nx, Tsoil.grid.Nz
     Tsurf = @view Tsoil[1:Nx, 1, Nz:Nz]
+    # Update speedy soil/skin temperature
     progn.land.soil_temperature .= Tsurf .+ NF(273.15)
     return nothing
 end
 
 ring_grid = RingGrids.FullGaussianGrid(24)
 Nz = 30
-Δz_min = 0.1 # currently the coupling is only stable with a large surface layer
-grid = ColumnRingGrid(CPU(), Float32, ExponentialSpacing(;N=Nz, Δz_min), ring_grid)
+Δ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 = FieldInitializers(
     # steady-ish state initial condition for soil temperature
-    temperature = (x,z) -> 10 - 0.02f0*z,
+    temperature = (x,z) -> 0 - 0.02f0*z,
     # fully saturated soil
     saturation_water_ice = 1.0f0,
 )
 
-# Coupled model with prescribed surface energy balance
-# currently not working :(
-# surface_energy_balance = SurfaceEnergyBalance(
-#     eltype(grid);
-#     radiative_fluxes = PrescribedRadiativeFluxes(),
-#     turbulent_fluxes = PrescribedTurbulentFluxes(),
-# )
-# soil_model = SoilModel(grid, initializer=soil_initializer)
-# model = CoupledSoilAtmosphereModel(soil_model; surface_energy_balance)
-# # Quick check to make sure everything works OK
-# land_state = initialize(model)
-# set!(land_state.state.surface_shortwave_down, 200.0)
-# set!(land_state.state.surface_longwave_down, 100.0)
-# set!(land_state.state.air_temperature, 11.0)
-# set!(land_state.state.air_pressure, 101_325)
-# set!(land_state.state.specific_humidity, 0.001)
-# timestep!(land_state)
-# run!(land_state, period=Day(1))
-# @assert all(isfinite.(land_state.state.skin_temperature))
-# @assert all(isfinite.(land_state.state.latent_heat_flux))
-# @assert all(isfinite.(land_state.state.temperature))
-
 # Soil model with prescribed surface temperautre BC
-soil_bcs = SoilBoundaryConditions(
-    eltype(grid);
-    top=PrescribedValue(:temperature, Input(:air_temperature, units=u"°C")),
-    bottom=DefaultBoundaryConditions()
-)
-model = SoilModel(grid, initializer=soil_initializer, boundary_conditions=soil_bcs)
-
-# Initialize Terrarium model
-land_state = initialize(model)
-land = TerrariumDryLand(land_state)
+model = SoilModel(grid, initializer=soil_initializer)
+Tair_input = InputSource(eltype(grid), :air_temperature)
+bcs = PrescribedSurfaceTemperature(:air_temperature)
+integrator = initialize(model, ForwardEuler(eltype(grid)), boundary_conditions=bcs, inputs=InputSources(Tair_input))
 
+# Initialize Terrarium-Speedy land model
+land = TerrariumDryLand(integrator)
 # Set up coupled model
 land_sea_mask = Speedy.RockyPlanetMask(land.spectral_grid)
-# surface_heat_flux = Speedy.SurfaceHeatFlux(land=Speedy.PrescribedLandHeatFlux())
 output = Speedy.NetCDFOutput(land.spectral_grid, Speedy.PrimitiveDryModel, path="outputs/")
-primitive_dry_coupled = Speedy.PrimitiveDryModel(land.spectral_grid; land, land_sea_mask, output)
+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)
-# spin up
-Speedy.run!(sim_coupled, period=Day(2), output=true)
-# run for a few more days
-Speedy.run!(sim_coupled, period=Day(5), output=true)
-
-with_theme(fontsize=18) do
-    Tair_fig = heatmap(sim_coupled.diagnostic_variables.grid.temp_grid[:,8] .- 273.15, title="", size=(800,400))
-    Tsoil_fig = heatmap(RingGrids.Field(interior(land_state.state.temperature)[:,1,end-4], 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)
-    Tair_fig
-end
-
-# animate surface layer air temperature
+# run it
+Speedy.run!(sim_coupled, period=Month(1), 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")
 
-@show land_state.state.air_temperature
-@show land_state.state.skin_temperature
-@show sim.prognostic_variables.land.soil_temperature
-
-heatmap(sim.diagnostic_variables.grid.temp_grid[:,end])
-heatmap(sim.prognostic_variables.land.soil_temperature[:,end])
-
-# Speedy standalone
-
-spectral_grid2 = Speedy.SpectralGrid(trunc=31)
-primitive_dry = Speedy.PrimitiveDryModel(spectral_grid2)
-sim2 = Speedy.initialize!(primitive_dry)
-Speedy.run!(sim2)
-
-using BenchmarkTools
-
-@profview @btime Speedy.timestep!($sim)
-
-default_spectral_grid = Speedy.SpectralGrid(grid.rings)
-default_primitive_dry = Speedy.PrimitiveDryModel(default_spectral_grid)
-default_sim = Speedy.initialize!(default_primitive_dry)
-Speedy.run!(default_sim, period=Day(1))
+# pick a point somewhere in the mid-lattitudes
+T = interior(integrator.state.temperature)[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[end-15:end], zs[end-15:end])
+Makie.scatterlines(f, zs)

src/models/models.jl

@@ -11,6 +11,11 @@ export SoilInitializer
 export ConstantInitialSoilTemperature, QuasiThermalSteadyState, PiecewiseLinearInitialSoilTemperature
 include("soil/soil_model_init.jl")
 
+# Soil-atmosphere
+
+export CoupledSoilAtmosphereModel
+include("soil/soil_atmosphere_model.jl")
+
 # Vegetation
 
 export VegetationModel

src/models/soil/soil_atmosphere_model.jl

@@ -3,10 +3,9 @@ struct CoupledSoilAtmosphereModel{
     GridType<:AbstractLandGrid{NF},
     Atmosphere<:AbstractAtmosphere,
     SurfaceEnergy<:AbstractSurfaceEnergyBalance,
-    TimeStepper<:AbstractTimeStepper,
-    SoilModel<:AbstractSoilModel{NF, GridType, TimeStepper},
+    SoilModel<:AbstractSoilModel{NF, GridType},
     Initializer<:AbstractInitializer,
-} <: AbstractSoilModel{NF, GridType, TimeStepper}
+} <: AbstractSoilModel{NF, GridType}
     "Spatial discretization"
     grid::GridType
 
@@ -33,20 +32,9 @@ function CoupledSoilAtmosphereModel(
     constants::PhysicalConstants = PhysicalConstants(NF),
     initializer::AbstractInitializer = DefaultInitializer()
 ) where {NF}
-    surface_fluxes = SoilBC(energy=GroundHeatFlux(), hydrology=DefaultBoundaryConditions())
-    coupled_soil_bcs = SoilBoundaryConditions(
-        NF,
-        top = surface_fluxes,
-        bottom = DefaultBoundaryConditions() # temporary hack
-    )
-    soil = setproperties(soil, boundary_conditions=coupled_soil_bcs)
     return CoupledSoilAtmosphereModel(soil.grid, atmosphere, surface_energy_balance, soil, constants, initializer)
 end
 
-get_boundary_conditions(model::CoupledSoilAtmosphereModel) = get_boundary_conditions(model.soil)
-
-get_time_stepping(model::CoupledSoilAtmosphereModel) = model.soil.time_stepping
-
 get_soil_energy_balance(model) = model.soil.energy
 
 get_soil_hydrology(model) = model.soil.hydrology
@@ -70,6 +58,19 @@ function variables(model::CoupledSoilAtmosphereModel)
     )
 end
 
+function initialize(
+    model::CoupledSoilAtmosphereModel{NF};
+    clock = Clock(time=zero(NF)),
+    boundary_conditions = (;),
+    fields = (;),
+    external_variables = ()
+) where {NF}
+    vars = Variables(variables(model)..., external_variables...)
+    surface_fluxes = GroundHeatFlux()
+    bcs = merge_recursive(boundary_conditions, surface_fluxes)
+    return StateVariables(vars, model.grid, clock, boundary_conditions=bcs)
+end
+
 function initialize!(state, model::CoupledSoilAtmosphereModel)
     initialize!(state, model, model.initializer)
     initialize!(state, model.soil)