Initial coupling to SpeedyWeather via example script

examples/Project.toml

@@ -1,7 +1,6 @@
 [deps]
 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"
@@ -12,11 +11,11 @@ Rasters = "a3a2b9e3-a471-40c9-b274-f788e487c689"
 RingGrids = "d1845624-ad4f-453b-8ff4-a8db365bf3a7"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
 SpeedyWeatherInternals = "34489162-d270-4603-8b96-37b04f830d73"
-TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Terrarium = "80418d68-07fa-499d-ae2b-c12e531f5cd8"
 
-[sources]
-Terrarium = {path = ".."}
-
 [compat]
 SpeedyWeatherInternals = "0.1.2"
+
+[sources.Terrarium]
+path = ".."

examples/example_model_notebook.jl

@@ -33,7 +33,7 @@ We begin by defining our model `struct` that subtypes `Terrarium.AbstractModel`:
 
 # ╔═╡ 4922e264-c80d-4a5b-8891-a7c8a3fdbfe7
 md""" 
-A `Terrarium.AbstractModel` typically constists of the fields 
+A "model" in Terrarium is a subtype of `Terrarium.AbstractModel` and is a `struct` type constisting of
  * `grid` which defines the discretization of the spatial domain
  * `initializer` which is responsible for initializing state variables
  * further fields that define processes, dynamics and submodels 
@@ -98,7 +98,7 @@ end
 md"""
 ## Defining the model behaviour 
 
-Now, we want to define our intended model behaviour. For this, we need to define the following routines: 
+Now, we want to define our intended model behaviour. For this, we need to define the following methods: 
 
 * `variables(::Model)` returns a tuple of variable metadata declaring the state variables. Variables must be one of three types: `prognostic`, `auxiliary` (sometimes referred to as “diagnostic”), or `input`. Prognostic variables fully characterize the state of the system at any given timestep and are updated according to their tendencies (i.e. ``u`` in our example). Tendencies are automatically allocated for each prognostic variable declared by the model. In this example we will treat the offset ``c`` as an auxiliary variable, though we could also just include it as a constant in the tendency computations.
 * `compute_auxiliary!(state, ::Model)` computes the values of all auxiliary variables (if necessary) assuming that the prognostic variables of the system in state are available for the current timestep.
@@ -115,9 +115,9 @@ Terrarium.variables(::ExpModel) = (
 
 # ╔═╡ d4d19de7-6f77-4873-9182-9832d1ca4381
 md"""
-Here, we defined our two variables with their name as a `Symbol` and whether they are 2D variables (`XY`) on the spatial grid or 3D variables (`XYZ`) that also vary along the vertical z-axis. Here we are considering only a simple scalar model so we choose 2D (`XY`).
+Here, we defined our two variables with their name as a `Symbol` and whether they are 2D variables (`XY`) on the spatial grid or 3D variables (`XYZ`) that also vary along the vertical z-axis. Here we are considering only a simple scalar model so we choose 2D (`XY`), bearing in mind that all points in the X and Y dimensions of `ColumnGrid` are independent of each other.
 
-We also need to define `compute_auxiliary!` and `compute_tendencies!` as discussed above. As is common in many Terrarium model definitions, we will simply forward these method calls to the process defined by the `dynamics` property.
+We also need to define `compute_auxiliary!` and `compute_tendencies!` as discussed above. We will use here a pattern which is commonly employed within Terrarium: we unpack the process from the model and forward the method calls to more specialzied ones defined for the `LinearDynamics` process.
 """
 
 # ╔═╡ 5ea313fc-3fbb-4092-a2cc-e0cd1f2fe641
@@ -149,7 +149,7 @@ function Terrarium.compute_auxiliary!(
 )
 	# set auxiliary variable for offset c
     state.auxiliary.c .= dynamics.c
-end 
+end
 
 # ╔═╡ 5c8be7e4-f150-492b-a75d-96887a11f6da
 # du/dt = u + c
@@ -159,11 +159,11 @@ function Terrarium.compute_tendencies!(
 	dynamics::LinearDynamics
 )
 	# define the dynamics; we'll use some special characters to make the equation nicer to look at :)
-	let u = state.prognostic.u,
+	let u = state.u,
 		∂u∂t = state.tendencies.u,
 		α = dynamics.alpha,
 		# note again that here we could also just use dynamics.c instead of defining an auxiliary variable! 
-		c = state.auxiliary.c;
+		c = state.c;
 		# Write into tendency variable ∂u∂t
 		∂u∂t .=  α * u + c
 	end
@@ -292,9 +292,9 @@ Well that's it. We defined and ran a simple exponential model following the Terr
 # ╠═94d82d31-42ec-41de-91e9-b5585b3a72d4
 # ╟─4922e264-c80d-4a5b-8891-a7c8a3fdbfe7
 # ╠═78f268ef-5385-4c63-bc35-2c973de69da5
-# ╠═054a8b11-250f-429f-966f-ca3c9a5dc2ef
+# ╟─054a8b11-250f-429f-966f-ca3c9a5dc2ef
 # ╠═407786db-607f-4508-b697-fe75b3ce0b25
-# ╠═575d920c-b12e-493f-95a7-5c962c3591fd
+# ╟─575d920c-b12e-493f-95a7-5c962c3591fd
 # ╠═82e45724-ba16-4806-9470-5cb4c43ea734
 # ╟─d4d19de7-6f77-4873-9182-9832d1ca4381
 # ╠═5ea313fc-3fbb-4092-a2cc-e0cd1f2fe641

examples/soil_heat_global_era5.jl

@@ -39,8 +39,7 @@ initializer = FieldInitializers(
 )
 model = SoilModel(grid; initializer, boundary_conditions)
 boundary_conditions = PrescribedSurfaceTemperature(:Tair)
-inputs = InputSources(Tair_forcing)
-integrator = initialize(model, ForwardEuler(); boundary_conditions, inputs)
+integrator = initialize(model, ForwardEuler(), Tair_forcing; boundary_conditions)
 @time timestep!(integrator)
 @time run!(integrator, period=Day(10), dt=120.0)
 

examples/speedy_dry_land.jl

@@ -0,0 +1,124 @@
+using Terrarium
+
+using CUDA
+using Dates
+using Rasters, NCDatasets
+using Statistics
+
+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, TMI<:Terrarium.ModelIntegrator{NF}} <: Speedy.AbstractDryLand
+    "Speedy spectral grid"
+    spectral_grid::Speedy.SpectralGrid
+
+    "Initialized Terrarium model integrator"
+    integrator::TMI
+
+    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
+
+function Speedy.initialize!(
+    progn_land::Speedy.PrognosticVariablesLand,  # for dispatch
+    progn::Speedy.PrognosticVariables{NF},
+    diagn::Speedy.DiagnosticVariables{NF},
+    land::TerrariumDryLand,
+    model::Speedy.PrimitiveEquation,
+) where {NF}
+    Tsoil = interior(land.integrator.state.temperature)[:, 1, end]
+    progn_land.soil_temperature .= Tsoil .+ NF(273.15)
+    Terrarium.initialize!(land.integrator)
+    return nothing
+end
+
+function Speedy.timestep!(
+    progn::Speedy.PrognosticVariables{NF},
+    diagn::Speedy.DiagnosticVariables{NF},
+    land::TerrariumDryLand,
+    model::Speedy.PrimitiveEquation,
+) where {NF}
+    # get speedy state variables
+    Tair = @view diagn.grid.temp_grid[:, end]
+    # 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.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) -> 0 - 0.02f0*z,
+    # fully saturated soil
+    saturation_water_ice = 1.0f0,
+)
+
+# Soil model with prescribed surface temperautre BC
+model = SoilModel(grid, initializer=soil_initializer)
+Tair_input = InputSource(eltype(grid), :air_temperature)
+bcs = PrescribedSurfaceTemperature(:air_temperature)
+integrator = initialize(model, ForwardEuler(eltype(grid)), Tair_input, boundary_conditions=bcs)
+
+# Initialize Terrarium-Speedy land model
+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/")
+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)
+
+# 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")
+
+# 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/Terrarium.jl

@@ -88,7 +88,7 @@ include("grids/vertical_discretization.jl")
 export ColumnGrid, ColumnRingGrid, get_field_grid
 include("grids/grids.jl")
 
-export InputSource
+export InputSource, InputSources, FieldInputSource, FieldTimeSeriesInputSource
 export update_inputs!
 include("input_output/input_sources.jl")
 

src/abstract_model.jl

@@ -136,7 +136,7 @@ abstract type AbstractLandModel{NF, GR} <: AbstractModel{NF, GR} end
 Convenience constructor for all `AbstractLandModel` types that allows the `grid` to be passed
 as the first positional argument.
 """
-(::Type{Model})(grid::AbstractLandGrid; kwargs...) where {Model<:AbstractModel} = Model(; grid, kwargs...)
+(::Type{ModelType})(grid::AbstractLandGrid, args...; kwargs...) where {ModelType<:AbstractModel} = ModelType(args...; grid, kwargs...)
 
 function Adapt.adapt_structure(to, model::AbstractModel)
     return setproperties(model, map(prop -> Adapt.adapt_structure(to, prop), getproperties(model)))

src/boundary_conditions.jl

@@ -10,11 +10,11 @@ Alias for a `NamedTuple` of `FieldBC` types where the keys correspond to field/v
 const FieldBCs{names, BCs} = NamedTuple{names, BCs} where {names, BCs<:Tuple{Vararg{FieldBC}}}
 
 """
-    boundary_conditions(bcs::FieldBCs...)
+    merge_boundary_conditions(bcs::FieldBCs...)
 
 Recursively merge an arbitrary number of field/variable boundary conditions.
 """
-boundary_conditions(bcs::FieldBCs...) = merge_recursive(bcs...)
+merge_boundary_conditions(bcs::FieldBCs...) = merge_recursive(bcs...)
 
 """
 Implementation of `Oceananigans.BoundaryConditions.getbc` for variable placeholders that retrieves the input `Field` from

src/input_output/input_sources.jl

@@ -41,7 +41,7 @@ InputSources(sources::InputSource...) = InputSources(Tuple(sources))
 
 variables(sources::InputSources) = tuplejoin(map(variables, sources.sources)...)
 
-function update_inputs!(fields, sources::InputSources, ::Clock)
+function update_inputs!(fields, sources::InputSources, clock::Clock)
     for source in sources.sources
         update_inputs!(fields, source, clock)
     end
@@ -50,31 +50,29 @@ end
 """
     $TYPEDEF
 
-Input source that reads directly from pre-specified Oceananigans `Field`s.
+Input source that defines `input` state variables with the given names which
+can then be directly modified by the user.
 """
-struct FieldInputSource{NF, VD<:VarDims, names, Fields<:Tuple{Vararg{AnyField{NF}}}} <: InputSource{NF}
+struct FieldInputSource{NF, VD<:VarDims, names} <: InputSource{NF}
     "Variable dimensions"
     dims::VD    
-
-    "Named tuple of input `Field`s"
-    fields::NamedTuple{names, Fields}
+
+    FieldInputSource(::Type{NF}, dims::VarDims, names::Symbol...) where {NF} = new{NF, typeof(dims), names}(dims)
 end
 
-function InputSource(named_fields::Pair{Symbol, <:Field}...)
-    dims = checkdims(; named_fields...)
-    return FieldInputSource(dims, (; named_fields...))
+"""
+    InputSource(::Type{NF}, names::Symbol...; dims = XY())
+
+Create a `FieldInputSource` with the given numeric type and input variable `names`.
+"""
+function InputSource(::Type{NF}, names::Symbol...; dims = XY()) where {NF}
+    return FieldInputSource(NF, dims, names...)
 end
 
 variables(source::FieldInputSource{NF, VD, names}) where {NF, VD, names} = map(name -> input(name, source.dims), names)
 
-function update_inputs!(fields, source::FieldInputSource, ::Clock)
-    for name in keys(source.fields)
-        if hasproperty(fields, name)
-            field = getproperty(fields, name)
-            set!(field, source.fields[name])
-        end
-    end
-end
+# For single fields; users can write directly into the allocated input variable
+update_inputs!(fields, ::FieldInputSource, ::Clock) = nothing
 
 """
 Type alias for a `FieldTimeSeries` with any X, Y, Z location or grid.

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