From 7b8944b5d14b4b455230107cd710ae8c66353c64 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Mon, 30 Mar 2026 23:44:50 +0200 Subject: [PATCH 1/2] WIP: Initial attempt at basic Terrarium extension --- .../NumericalEarthTerrariumExt.jl | 15 +++++ .../terrarium_exchanger.jl | 60 +++++++++++++++++++ .../terrarium_land_simulations.jl | 38 ++++++++++++ 3 files changed, 113 insertions(+) create mode 100644 ext/NumericalEarthTerrariumExt/NumericalEarthTerrariumExt.jl create mode 100644 ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl create mode 100644 ext/NumericalEarthTerrariumExt/terrarium_land_simulations.jl diff --git a/ext/NumericalEarthTerrariumExt/NumericalEarthTerrariumExt.jl b/ext/NumericalEarthTerrariumExt/NumericalEarthTerrariumExt.jl new file mode 100644 index 000000000..89666cb86 --- /dev/null +++ b/ext/NumericalEarthTerrariumExt/NumericalEarthTerrariumExt.jl @@ -0,0 +1,15 @@ +module NumericalEarthTerrariumExt + +using OffsetArrays +using KernelAbstractions +using Statistics + +import Terrarium +import Terrarium.RingGrids +import NumericalEarth +import Oceananigans + +include("terrarium_land_simulations.jl") +include("terrarium_exchanger.jl") + +end # module NumericalEarthTerrariumExt diff --git a/ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl b/ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl new file mode 100644 index 000000000..89e48bc36 --- /dev/null +++ b/ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl @@ -0,0 +1,60 @@ +using Oceananigans +using Oceananigans.BoundaryConditions: fill_halo_regions! +using Oceananigans.Grids: architecture +using Oceananigans.Fields: set!, interior + +import NumericalEarth.EarthSystemModels: update_net_fluxes!, interpolate_state! +import NumericalEarth.EarthSystemModels.InterfaceComputations: net_fluxes, ComponentExchanger + +net_fluxes(::LandSimulation) = nothing + +# Land exchanger constructor. +# For now, no regridder is needed since the exchange grid is assumed to match the land grid. +# The state holds a single surface temperature field that is communicated back to the atmosphere. +# TODO: add regridder when exchange grid differs from land grid. +function ComponentExchanger(land::LandSimulation, exchange_grid) + regridder = nothing + state = (; Ts = Field{Center, Center, Nothing}(exchange_grid)) + return ComponentExchanger(state, regridder) +end + +# Read the land surface state onto the exchange grid. +# Currently: exchange grid == land grid → direct copy, no regridding. +# TODO: regrid when exchange_grid differs from land grid. +function interpolate_state!(exchanger, exchange_grid, land::LandSimulation, coupled_model) + Ts_land = land.state.skin_temperature # °C + Ts_exchange = exchanger.state.Ts + + # Convert skin temperature from °C to K and copy to exchange grid field + set!(Ts_exchange, Ts_land + 273.15) + fill_halo_regions!(Ts_exchange) # TODO: is this necessary? + + return nothing +end + +# Update Terrarium land model inputs from the atmospheric state on the exchange grid. +# This is the primary atmosphere → land coupling step. +# TODO: regrid when land grid differs from the exchange grid. +function update_net_fluxes!(coupled_model, land::LandSimulation) + atmos_state = coupled_model.interfaces.exchanger.atmosphere.state + + # Air temperature: atmosphere provides K, Terrarium expects °C + set!(land.state.air_temperature, atmos_state.T - 273.15) + + # Wind speed: compute magnitude from (u, v) components + # TODO: implement as a proper GPU-compatible kernel + set!(land.state.windspeed, sqrt(atmos_state.u^2 + atmos_state.v^2)) + + # Remaining atmospheric scalars: direct copy (same units) + set!(land.state.specific_humidity, atmos_state.q) + set!(land.state.air_pressure, atmos_state.p) + set!(land.state.surface_shortwave_down, atmos_state.ℐꜜˢʷ) + set!(land.state.surface_longwave_down, atmos_state.ℐꜜˡʷ) + + # Total precipitation → rainfall; snowfall set to zero. + # TODO: partition rain/snow based on air temperature. + set!(land.state.rainfall, atmos_state.Jᶜ) + set!(land.state.snowfall, 0) + + return nothing +end diff --git a/ext/NumericalEarthTerrariumExt/terrarium_land_simulations.jl b/ext/NumericalEarthTerrariumExt/terrarium_land_simulations.jl new file mode 100644 index 000000000..7b1e71278 --- /dev/null +++ b/ext/NumericalEarthTerrariumExt/terrarium_land_simulations.jl @@ -0,0 +1,38 @@ +import NumericalEarth.Land: land_simulation + +const LandSimulation = Terrarium.ModelIntegrator +const LandEarthSystemModel = NumericalEarth.EarthSystemModel{<:Any, <:Any, <:LandSimulation} #TODO: fix when updated + +Base.summary(::LandSimulation) = "Terrarium.ModelIntegrator" + +# Take one time-step or more depending on the global timestep +function Oceananigans.TimeSteppers.time_step!(integrator::Terrarium.ModelIntegrator, Δt) + Δt_land = Terrarium.default_dt(integrator.timestepper) + nsteps = ceil(Int, Δt / Δt_land) + + if (Δt / Δt_land) % 1 != 0 + @warn "NumericalEarth only supports land timesteps that are integer divisors of the ESM timesteps" + end + + for _ in 1:nsteps + Terrarium.timestep!(integrator, Δt_land) + end + return +end + +""" + land_simulation(grid::Terrarium.AbstractLandGrid) + +Return a land simulation based on the given `AbstractLandGrid`. + +TODO: add more kwarg config options +""" +function land_simulation(grid::Terrarium.AbstractLandGrid) + # The land model + land_model = Terrarium.LandModel(grid) + + # Construct the Terrarium integrator + integrator = Terrarium.initialize(land_model) + + return integrator +end From 9d0ea78b84a2c0f65e22acfbcdf9a94faedfc4c3 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Wed, 1 Apr 2026 17:12:18 +0200 Subject: [PATCH 2/2] Update ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl Co-authored-by: Simone Silvestri --- ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl b/ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl index 89e48bc36..392509992 100644 --- a/ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl +++ b/ext/NumericalEarthTerrariumExt/terrarium_exchanger.jl @@ -26,8 +26,7 @@ function interpolate_state!(exchanger, exchange_grid, land::LandSimulation, coup Ts_exchange = exchanger.state.Ts # Convert skin temperature from °C to K and copy to exchange grid field - set!(Ts_exchange, Ts_land + 273.15) - fill_halo_regions!(Ts_exchange) # TODO: is this necessary? + parent(Ts_exchange) .= parent(Ts_land) .+ 273.15 return nothing end