From 9ea8bc8cc175ae85c083f2d0c1ea3c6f0b942e57 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Fri, 1 May 2026 00:26:05 +0200 Subject: [PATCH 1/7] Update SpeedyWeather version --- Project.toml | 2 +- examples/Project.toml | 3 --- 2 files changed, 1 insertion(+), 4 deletions(-) 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/examples/Project.toml b/examples/Project.toml index d5cb56b51..0506e53c2 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -19,8 +19,5 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Terrarium = "80418d68-07fa-499d-ae2b-c12e531f5cd8" -[compat] -SpeedyWeather = "0.18.1" - [sources.Terrarium] path = ".." From 24da6fd9a244c8bc37b3f8ffe7bc034755207d3c Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Fri, 1 May 2026 00:26:54 +0200 Subject: [PATCH 2/7] Always ensure that y-axis is removed in RingGrids Field converter --- src/grids/column_ring_grid.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 From ee623b0d2f3dd04a24fc9b4565cb80442e5cd6b2 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Fri, 1 May 2026 00:27:16 +0200 Subject: [PATCH 3/7] Initial kind of working downscaling --- examples/simulations/speedy_downscaling.jl | 242 +++++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 examples/simulations/speedy_downscaling.jl diff --git a/examples/simulations/speedy_downscaling.jl b/examples/simulations/speedy_downscaling.jl new file mode 100644 index 000000000..a9f5b9cf6 --- /dev/null +++ b/examples/simulations/speedy_downscaling.jl @@ -0,0 +1,242 @@ +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)" + buf_terrarium::FT + + "Reusable 2D buffer Field on Speedy ring grid (for downscaling Terrarium → Speedy)" + buf_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) + buf_terrarium = zeros(NF, terrarium_ring_grid) + buf_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(buf_terrarium), typeof(buf_speedy)}( + spectral_grid, geometry, integrator, interp_up, interp_down, buf_terrarium, buf_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 + terrarium_ring_grid = land.buf_terrarium.grid + # Downscale Terrarium initial state to Speedy resolution + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.temperature)[:, 1, end], terrarium_ring_grid), land.interp_down) + vars.prognostic.land.soil_temperature .= land.buf_speedy .+ NF(273.15) + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.saturation_water_ice)[:, 1, end], terrarium_ring_grid), land.interp_down) + vars.prognostic.land.soil_moisture .= land.buf_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 + +function speedy_timestep!( + vars::Speedy.Variables, + land::TerrariumWetLand{NF}, + ) where {NF} + # land constants + consts = land.integrator.model.constants + state = land.integrator.state + terrarium_ring_grid = land.buf_terrarium.grid + + terrarium_grid = land.integrator.model.grid + # Upscale Speedy fields to Terrarium resolution and update inputs + RingGrids.interpolate!(land.buf_terrarium, vars.grid.temperature[:, end], land.interp_up) + land.buf_terrarium.data .-= NF(273.15) ## convert K → °C before set! + set!(state.inputs.air_temperature, land.buf_terrarium) + RingGrids.interpolate!(land.buf_terrarium, vars.grid.pressure, land.interp_up) + land.buf_terrarium.data .= exp.(land.buf_terrarium.data) ## convert log(Pa) → Pa before set! + set!(state.inputs.air_pressure, land.buf_terrarium) + RingGrids.interpolate!(land.buf_terrarium, vars.grid.humidity[:, end], land.interp_up) + set!(state.inputs.specific_humidity, land.buf_terrarium) + RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.rain_rate, land.interp_up) + set!(state.inputs.rainfall, land.buf_terrarium) + RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.snow_rate, land.interp_up) + set!(state.inputs.snowfall, land.buf_terrarium) + RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.surface_wind_speed, land.interp_up) + set!(state.inputs.windspeed, land.buf_terrarium) + RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.surface_shortwave_down, land.interp_up) + set!(state.inputs.surface_shortwave_down, land.buf_terrarium) + RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.surface_longwave_down, land.interp_up) + set!(state.inputs.surface_longwave_down, land.buf_terrarium) + # 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 + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.skin_temperature)[:, 1, 1], terrarium_ring_grid), land.interp_down) + vars.prognostic.land.soil_temperature .= land.buf_speedy .+ NF(273.15) + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.saturation_water_ice)[:, 1, end], terrarium_ring_grid), land.interp_down) + vars.prognostic.land.soil_moisture .= land.buf_speedy + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.sensible_heat_flux)[:, 1, 1], terrarium_ring_grid), land.interp_down) + vars.prognostic.land.sensible_heat_flux .= land.buf_speedy + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.latent_heat_flux)[:, 1, 1], terrarium_ring_grid), land.interp_down) + vars.prognostic.land.surface_humidity_flux .= land.buf_speedy ./ consts.Llg + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.surface_longwave_up)[:, 1, 1], terrarium_ring_grid), land.interp_down) + vars.parameterizations.land.surface_longwave_up .= land.buf_speedy + RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.surface_shortwave_up)[:, 1, 1], terrarium_ring_grid), land.interp_down) + vars.parameterizations.land.surface_shortwave_up .= land.buf_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(48, speedy_arch) +terrarium_ring_grid_cpu = on_architecture(CPU(), terrarium_ring_grid) + +Nz = 30 +Δz_min = 0.05 +grid = ColumnRingGrid(terrarium_arch, Float32, ExponentialSpacing(; N = Nz, Δz_min), terrarium_ring_grid_cpu) +# 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.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.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 +period = Hour(1) +@info "Running simulation for $period" +@time Speedy.run!(sim_coupled; period) +Terrarium.checkfinite!(integrator.state.prognostic) + +# Plotting: interpolate HEALPix fields to FullGaussianGrid on CPU (HEALPix not supported for heatmap) +land_plot_grid = RingGrids.FullGaussianGrid(48) ## matches terrarium ring grid resolution +atm_plot_grid = RingGrids.FullGaussianGrid(24) ## matches speedy ring grid resolution +land_interp_plot = RingGrids.interpolator(land_plot_grid, terrarium_ring_grid_cpu) +atm_interp_plot = RingGrids.interpolator(atm_plot_grid, speedy_ring_grid_cpu) + +## Wrap a (GPU or CPU) array slice as a RingGrids.Field and interpolate to a FullGaussianField on CPU +function to_plotting_grid(data::AbstractArray, src_grid::RingGrids.AbstractGrid, plot_grid::RingGrids.AbstractGrid, interp) + src = RingGrids.Field(Array(data), src_grid) + out = zeros(eltype(src), plot_grid) + RingGrids.interpolate!(out, src, interp) + return out +end + +# Land variables +Tsoil_fig = heatmap(to_plotting_grid(interior(integrator.state.temperature)[:, 1, end - 2], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Soil temperature (°C)", size = (800, 400)) +Tsurf_fig = heatmap(to_plotting_grid(interior(integrator.state.skin_temperature)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Skin temperature (°C)", size = (800, 400)) +Hs_fig = heatmap(to_plotting_grid(interior(integrator.state.sensible_heat_flux)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Sensible heat flux (W/m²)", size = (800, 400)) +Hl_fig = heatmap(to_plotting_grid(interior(integrator.state.latent_heat_flux)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Latent heat flux (W/m²)", size = (800, 400)) +E_fig = heatmap(to_plotting_grid(interior(integrator.state.evaporation_ground)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Evaporation (m/s)", size = (800, 400)) +sat_fig = heatmap(to_plotting_grid(interior(integrator.state.saturation_water_ice)[:, 1, end], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Surface saturation", size = (800, 400)) + +# Atmosphere variables (Speedy v0.19 API) +Tair_fig = heatmap(to_plotting_grid(Array(sim_coupled.variables.grid.temperature[:, 8].data) .- 273.15f0, RingGrids.HEALPixGrid(24), atm_plot_grid, atm_interp_plot), title = "Air temperature (°C)", size = (800, 400)) +pres_fig = heatmap(to_plotting_grid(exp.(Array(sim_coupled.variables.grid.pressure.data)), RingGrids.HEALPixGrid(24), atm_plot_grid, atm_interp_plot), title = "Surface pressure (Pa)", size = (800, 400)) +srad_fig = heatmap(to_plotting_grid(Array(sim.variables.parameterizations.surface_shortwave_down.data), RingGrids.HEALPixGrid(24), atm_plot_grid, atm_interp_plot), title = "Surface shortwave down (W/m²)", size = (800, 400)) + +# Pick a point somewhere in the mid-latitudes and plot vertical profiles +T = Array(interior(integrator.state.temperature)[2000, 1, :]) +sat = Array(interior(integrator.state.saturation_water_ice)[2000, 1, :]) +f = Array(interior(integrator.state.liquid_water_fraction)[2000, 1, :]) +zs = znodes(integrator.state.temperature) + +# Plot temperature, saturation, and liquid fraction vertical profiles +Makie.scatterlines(T, zs) +Makie.scatterlines(sat, zs) +Makie.scatterlines(f, zs) From 334b358e4db8a1f11875afc04728aac52cad0000 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Sun, 3 May 2026 15:52:19 +0200 Subject: [PATCH 4/7] Use land-sea mask in downscaling Also adds plotting code for animation --- examples/simulations/speedy_downscaling.jl | 277 +++++++++++++++------ 1 file changed, 202 insertions(+), 75 deletions(-) diff --git a/examples/simulations/speedy_downscaling.jl b/examples/simulations/speedy_downscaling.jl index a9f5b9cf6..682615ca4 100644 --- a/examples/simulations/speedy_downscaling.jl +++ b/examples/simulations/speedy_downscaling.jl @@ -44,10 +44,10 @@ struct TerrariumWetLand{ interp_down::ID "Reusable 2D buffer Field on Terrarium ring grid (for upscaling Speedy → Terrarium)" - buf_terrarium::FT + buffer_terrarium::FT "Reusable 2D buffer Field on Speedy ring grid (for downscaling Terrarium → Speedy)" - buf_speedy::FS + buffer_speedy::FS function TerrariumWetLand( integrator::Terrarium.ModelIntegrator{NF, Arch, Grid}, @@ -58,13 +58,13 @@ struct TerrariumWetLand{ 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) - buf_terrarium = zeros(NF, terrarium_ring_grid) - buf_speedy = zeros(NF, speedy_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(buf_terrarium), typeof(buf_speedy)}( - spectral_grid, geometry, integrator, interp_up, interp_down, buf_terrarium, buf_speedy + 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 @@ -81,12 +81,14 @@ function Speedy.initialize!( ) where {NF} Terrarium.initialize!(land.integrator) state = land.integrator.state - terrarium_ring_grid = land.buf_terrarium.grid + grid = land.integrator.model.grid # Downscale Terrarium initial state to Speedy resolution - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.temperature)[:, 1, end], terrarium_ring_grid), land.interp_down) - vars.prognostic.land.soil_temperature .= land.buf_speedy .+ NF(273.15) - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.saturation_water_ice)[:, 1, end], terrarium_ring_grid), land.interp_down) - vars.prognostic.land.soil_moisture .= land.buf_speedy + 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 @@ -99,6 +101,14 @@ function Speedy.timestep!( 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}, @@ -106,44 +116,50 @@ function speedy_timestep!( # land constants consts = land.integrator.model.constants state = land.integrator.state - terrarium_ring_grid = land.buf_terrarium.grid - terrarium_grid = land.integrator.model.grid + mask = terrarium_grid.mask.data # Upscale Speedy fields to Terrarium resolution and update inputs - RingGrids.interpolate!(land.buf_terrarium, vars.grid.temperature[:, end], land.interp_up) - land.buf_terrarium.data .-= NF(273.15) ## convert K → °C before set! - set!(state.inputs.air_temperature, land.buf_terrarium) - RingGrids.interpolate!(land.buf_terrarium, vars.grid.pressure, land.interp_up) - land.buf_terrarium.data .= exp.(land.buf_terrarium.data) ## convert log(Pa) → Pa before set! - set!(state.inputs.air_pressure, land.buf_terrarium) - RingGrids.interpolate!(land.buf_terrarium, vars.grid.humidity[:, end], land.interp_up) - set!(state.inputs.specific_humidity, land.buf_terrarium) - RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.rain_rate, land.interp_up) - set!(state.inputs.rainfall, land.buf_terrarium) - RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.snow_rate, land.interp_up) - set!(state.inputs.snowfall, land.buf_terrarium) - RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.surface_wind_speed, land.interp_up) - set!(state.inputs.windspeed, land.buf_terrarium) - RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.surface_shortwave_down, land.interp_up) - set!(state.inputs.surface_shortwave_down, land.buf_terrarium) - RingGrids.interpolate!(land.buf_terrarium, vars.parameterizations.surface_longwave_down, land.interp_up) - set!(state.inputs.surface_longwave_down, land.buf_terrarium) + ## 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 - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.skin_temperature)[:, 1, 1], terrarium_ring_grid), land.interp_down) - vars.prognostic.land.soil_temperature .= land.buf_speedy .+ NF(273.15) - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.saturation_water_ice)[:, 1, end], terrarium_ring_grid), land.interp_down) - vars.prognostic.land.soil_moisture .= land.buf_speedy - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.sensible_heat_flux)[:, 1, 1], terrarium_ring_grid), land.interp_down) - vars.prognostic.land.sensible_heat_flux .= land.buf_speedy - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.latent_heat_flux)[:, 1, 1], terrarium_ring_grid), land.interp_down) - vars.prognostic.land.surface_humidity_flux .= land.buf_speedy ./ consts.Llg - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.surface_longwave_up)[:, 1, 1], terrarium_ring_grid), land.interp_down) - vars.parameterizations.land.surface_longwave_up .= land.buf_speedy - RingGrids.interpolate!(land.buf_speedy, RingGrids.Field(interior(state.surface_shortwave_up)[:, 1, 1], terrarium_ring_grid), land.interp_down) - vars.parameterizations.land.surface_shortwave_up .= land.buf_speedy + 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 @@ -157,12 +173,22 @@ 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(48, speedy_arch) +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) +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))) @@ -179,7 +205,7 @@ 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.RockyPlanetMask(land.spectral_grid) +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/") @@ -197,44 +223,145 @@ primitive_wet_coupled = Speedy.PrimitiveWetModel( Speedy.add!(primitive_wet_coupled.output, Speedy.SoilTemperatureOutput()) # initialize coupled simulation sim_coupled = Speedy.initialize!(primitive_wet_coupled) -# run it -period = Hour(1) +# run it (kept for reference) — we'll also run a stepwise loop below to build an animation +period = Month(1) @info "Running simulation for $period" @time Speedy.run!(sim_coupled; period) Terrarium.checkfinite!(integrator.state.prognostic) -# Plotting: interpolate HEALPix fields to FullGaussianGrid on CPU (HEALPix not supported for heatmap) -land_plot_grid = RingGrids.FullGaussianGrid(48) ## matches terrarium ring grid resolution -atm_plot_grid = RingGrids.FullGaussianGrid(24) ## matches speedy ring grid resolution -land_interp_plot = RingGrids.interpolator(land_plot_grid, terrarium_ring_grid_cpu) -atm_interp_plot = RingGrids.interpolator(atm_plot_grid, speedy_ring_grid_cpu) - -## Wrap a (GPU or CPU) array slice as a RingGrids.Field and interpolate to a FullGaussianField on CPU -function to_plotting_grid(data::AbstractArray, src_grid::RingGrids.AbstractGrid, plot_grid::RingGrids.AbstractGrid, interp) - src = RingGrids.Field(Array(data), src_grid) - out = zeros(eltype(src), plot_grid) - RingGrids.interpolate!(out, src, interp) - return out +# --- 2×2 Animation: land and atmosphere state evolution during the coupled run --- +# Re-initialize the simulation to start from the beginning for the animation. +Speedy.initialize!(sim_coupled; steps = 240, output = false) + +# 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_prec = speedy_frame(sim_coupled.variables.parameterizations.rain_rate) + +# 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 = (-30.0f0, 30.0f0) # °C — typical land surface temperature range +crange_evap = (0.0f0, 1.0f-5) # m s⁻¹ — bare-soil evaporation +crange_tair = (-30.0f0, 30.0f0) # °C — tropospheric temperature at lowest model level +crange_prec = (0.0f0, 2.0f-7) # m s⁻¹ — rain rate + +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_prec = heatmap_axis(fig_anim, 2, 3, "Precipitation (m s⁻¹)") + + obs_tsoil = Observable(mat_tsoil) + obs_evap = Observable(mat_evap) + obs_tair = Observable(mat_tair) + obs_prec = Observable(mat_prec) + + 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_prec = heatmap!(ax_prec, lond_lo, latd_lo, obs_prec; colormap = :Blues, colorrange = crange_prec) + + 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_prec; label = "m s⁻¹", 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:240; 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_prec[] = speedy_frame(sim_coupled.variables.parameterizations.rain_rate) + 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 -Tsoil_fig = heatmap(to_plotting_grid(interior(integrator.state.temperature)[:, 1, end - 2], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Soil temperature (°C)", size = (800, 400)) -Tsurf_fig = heatmap(to_plotting_grid(interior(integrator.state.skin_temperature)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Skin temperature (°C)", size = (800, 400)) -Hs_fig = heatmap(to_plotting_grid(interior(integrator.state.sensible_heat_flux)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Sensible heat flux (W/m²)", size = (800, 400)) -Hl_fig = heatmap(to_plotting_grid(interior(integrator.state.latent_heat_flux)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Latent heat flux (W/m²)", size = (800, 400)) -E_fig = heatmap(to_plotting_grid(interior(integrator.state.evaporation_ground)[:, 1, 1], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Evaporation (m/s)", size = (800, 400)) -sat_fig = heatmap(to_plotting_grid(interior(integrator.state.saturation_water_ice)[:, 1, end], terrarium_ring_grid_cpu, land_plot_grid, land_interp_plot), title = "Surface saturation", size = (800, 400)) +# 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) -Tair_fig = heatmap(to_plotting_grid(Array(sim_coupled.variables.grid.temperature[:, 8].data) .- 273.15f0, RingGrids.HEALPixGrid(24), atm_plot_grid, atm_interp_plot), title = "Air temperature (°C)", size = (800, 400)) -pres_fig = heatmap(to_plotting_grid(exp.(Array(sim_coupled.variables.grid.pressure.data)), RingGrids.HEALPixGrid(24), atm_plot_grid, atm_interp_plot), title = "Surface pressure (Pa)", size = (800, 400)) -srad_fig = heatmap(to_plotting_grid(Array(sim.variables.parameterizations.surface_shortwave_down.data), RingGrids.HEALPixGrid(24), atm_plot_grid, atm_interp_plot), title = "Surface shortwave down (W/m²)", 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 = Array(interior(integrator.state.temperature)[2000, 1, :]) -sat = Array(interior(integrator.state.saturation_water_ice)[2000, 1, :]) -f = Array(interior(integrator.state.liquid_water_fraction)[2000, 1, :]) -zs = znodes(integrator.state.temperature) +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) From 3dfadd598121c90a01ee7fc556c8574421329f21 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Sun, 3 May 2026 15:54:30 +0200 Subject: [PATCH 5/7] Update cpu vs. gpu benchmark plot --- .../gpu/soil_heat_hydrology_global.jl | 55 +++++++++++++++---- 1 file changed, 43 insertions(+), 12 deletions(-) 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 From 24ca6da4203167a112dcaabb3960b956d526d0aa Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Sun, 3 May 2026 23:21:25 +0200 Subject: [PATCH 6/7] Fix formatting issue --- src/state_variables.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 """ From 0edbcd8853a404e1238fafec3684b8c91e134c1a Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Sun, 3 May 2026 23:21:47 +0200 Subject: [PATCH 7/7] Adjust colorbar limits in downscaling plots --- examples/simulations/speedy_downscaling.jl | 29 +++++++++++++--------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/examples/simulations/speedy_downscaling.jl b/examples/simulations/speedy_downscaling.jl index 682615ca4..18de35fb3 100644 --- a/examples/simulations/speedy_downscaling.jl +++ b/examples/simulations/speedy_downscaling.jl @@ -224,14 +224,19 @@ 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 = Month(1) +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. -Speedy.initialize!(sim_coupled; steps = 240, output = false) +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) @@ -265,7 +270,7 @@ end 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_prec = speedy_frame(sim_coupled.variables.parameterizations.rain_rate) +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. @@ -283,10 +288,10 @@ end # Fixed colorbar ranges — standardised so colors are consistent across all frames. # Adjust these if your simulation uses very different conditions. -crange_tsoil = (-30.0f0, 30.0f0) # °C — typical land surface temperature range +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 = (-30.0f0, 30.0f0) # °C — tropospheric temperature at lowest model level -crange_prec = (0.0f0, 2.0f-7) # m s⁻¹ — rain rate +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)) @@ -294,34 +299,34 @@ Makie.with_theme(fontsize = 16) do 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_prec = heatmap_axis(fig_anim, 2, 3, "Precipitation (m s⁻¹)") + 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_prec = Observable(mat_prec) + 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_prec = heatmap!(ax_prec, lond_lo, latd_lo, obs_prec; colormap = :Blues, colorrange = crange_prec) + 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_prec; label = "m s⁻¹", 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:240; framerate = 12) do _ + 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_prec[] = speedy_frame(sim_coupled.variables.parameterizations.rain_rate) + obs_humid[] = speedy_frame(sim_coupled.variables.grid.humidity[:, end]) time_label[] = string(clock_anim.time) end end