Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
cd250d9
Replace XESMF with ConservativeRegridding as hard dependency
simone-silvestri Apr 14, 2026
5b55493
run example in PR
simone-silvestri Apr 14, 2026
03b5ce4
Fix vertex ordering in treeify to counter-clockwise
simone-silvestri Apr 14, 2026
2225b9a
Merge branch 'main' into ss/conservative-regridding
simone-silvestri Apr 14, 2026
d98db40
Add quadtree spatial acceleration for RingGrids treeify
simone-silvestri Apr 14, 2026
a67f3da
Use MultiTreeWrapper for reduced grids, revert regrid! to simple form
simone-silvestri Apr 14, 2026
de5224e
Clean up speedy_regridder: shared const, simpler helpers
simone-silvestri Apr 14, 2026
a2359a1
Add GeoInterface, GeometryOps, GeometryOpsCore to deps
simone-silvestri Apr 14, 2026
4d7e2b3
Add LinearAlgebra to deps for transpose in extension
simone-silvestri Apr 14, 2026
7cf0aae
Fix radius mismatch between Oceananigans and SpeedyWeather manifolds
simone-silvestri Apr 14, 2026
d25c51e
Update spectral grid configuration in simulation
simone-silvestri Apr 27, 2026
4f25e4c
Merge branch 'main' into ss/conservative-regridding
simone-silvestri Apr 27, 2026
1485cf9
Merge branch 'main' into ss/conservative-regridding
simone-silvestri Apr 27, 2026
64c9c67
Merge branch 'main' into ss/conservative-regridding
simone-silvestri Apr 29, 2026
264a922
Merge branch 'main' into ss/conservative-regridding
simone-silvestri May 1, 2026
cba76b8
update the regridding
simone-silvestri May 1, 2026
cf0c7b6
fix this
simone-silvestri May 1, 2026
b61b53d
Apply suggestion from @simone-silvestri
simone-silvestri May 1, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ authors = ["NumericalEarth contributors"]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
CFTime = "179af706-886a-5703-950a-314cd64e0468"
ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4"
ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343"
CubedSphere = "7445602f-e544-4518-8976-18f8e8ae6cdb"
DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
Expand Down Expand Up @@ -41,17 +42,19 @@ PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
WorldOceanAtlasTools = "04f20302-f1b9-11e8-29d9-7d841cb0a64a"
XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"

[extensions]
NumericalEarthBreezeExt = "Breeze"
NumericalEarthCDSAPIExt = "CDSAPI"
NumericalEarthCopernicusMarineExt = "CopernicusMarine"
NumericalEarthReactantExt = "Reactant"
NumericalEarthSpeedyWeatherExt = ["SpeedyWeather", "XESMF"]
NumericalEarthSpeedyWeatherExt = "SpeedyWeather"
NumericalEarthVerosExt = ["PythonCall", "CondaPkg"]
NumericalEarthWOAExt = "WorldOceanAtlasTools"

[sources]
ConservativeRegridding = {url = "https://github.com/JuliaGeo/ConservativeRegridding.jl", rev = "main"}

[compat]
Adapt = "4"
Breeze = "0.4"
Expand Down Expand Up @@ -84,7 +87,7 @@ SpeedyWeather = "0.19"
StaticArrays = "1"
Statistics = "<0.0.1, 1"
Thermodynamics = "0.15.3"
ConservativeRegridding = "0.2.1"
WorldOceanAtlasTools = "0.6"
XESMF = "0.1.6"
ZipFile = "0.10"
julia = "1.10"
1 change: 0 additions & 1 deletion docs/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"

[sources]
NumericalEarth = {path = ".."}
Expand Down
2 changes: 1 addition & 1 deletion docs/make.jl
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ examples = [
Example("Single-column ocean simulation", "single_column_os_papa_simulation", true),
Example("One-degree ocean--sea ice simulation", "one_degree_simulation", false),
Example("Near-global ocean simulation", "near_global_ocean_simulation", false),
Example("Global climate simulation", "global_climate_simulation", false),
Example("Global climate simulation", "global_climate_simulation", true),
Example("Veros ocean simulation", "veros_ocean_forced_simulation", false),
Example("Breeze over two oceans", "breeze_over_two_oceans", false),
Example("ERA5 winds and Stokes drift", "ERA5_winds_and_stokes_drift", true),
Expand Down
4 changes: 2 additions & 2 deletions examples/global_climate_simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@
# `SpeedyWeather.PrimitiveWetModel` (the atmosphere). All these are coupled and orchestrated by the `NumericalEarth.EarthSystemModel`
# (the coupled system).
#
# The XESMF.jl package is used to regrid fields between the atmosphere and ocean--sea ice components.
# The ConservativeRegridding.jl package is used to regrid fields between the atmosphere and ocean--sea ice components.

using Oceananigans, SpeedyWeather, XESMF, NumericalEarth
using Oceananigans, SpeedyWeather, NumericalEarth
using NCDatasets, CairoMakie
using Oceananigans.Units
using Printf, Statistics, Dates
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,12 @@ using OffsetArrays
using KernelAbstractions
using Statistics

import SpeedyWeather
import NumericalEarth
import Oceananigans
import SpeedyWeather
import NumericalEarth
import Oceananigans
import SpeedyWeather.RingGrids

include("speedy_atmosphere_simulations.jl")
include("speedy_regridder.jl")
include("speedy_weather_exchanger.jl")

end # module NumericalEarthSpeedyWeatherExt
49 changes: 0 additions & 49 deletions ext/NumericalEarthSpeedyWeatherExt/speedy_regridder.jl

This file was deleted.

72 changes: 39 additions & 33 deletions ext/NumericalEarthSpeedyWeatherExt/speedy_weather_exchanger.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,12 @@ using Oceananigans.BoundaryConditions
using Oceananigans.Grids: architecture
using Oceananigans.Utils: launch!
using Oceananigans.Operators: intrinsic_vector
using XESMF

using ConservativeRegridding: Regridder
import ConservativeRegridding: regrid!

using NumericalEarth.EarthSystemModels: sea_ice_concentration

# TODO: Implement conservative regridding when ready
# using ConservativeRegridding
# using GeoInterface: Polygon, LinearRing
import NumericalEarth.EarthSystemModels: update_net_fluxes!, interpolate_state!
import NumericalEarth.Atmospheres: atmosphere_regridder
import NumericalEarth.EarthSystemModels.InterfaceComputations: net_fluxes, ComponentExchanger
Expand All @@ -23,14 +22,16 @@ net_fluxes(::SpeedySimulation) = nothing
# If this work we can
# 1. Copy speedyweather gridarrays to the GPU
# 2. Perform the regridding on the GPU
function ComponentExchanger(atmosphere::SpeedySimulation, exchange_grid)
function ComponentExchanger(atmosphere::SpeedySimulation, exchange_grid)

spectral_grid = atmosphere.model.spectral_grid
# TODO: Implement a conservative regridder when ready
from_atmosphere = XESMF.Regridder(spectral_grid, exchange_grid)
to_atmosphere = XESMF.Regridder(exchange_grid, spectral_grid)
# Use the exchange_grid's manifold for both grids to avoid
# radius mismatch between Oceananigans and SpeedyWeather.
manifold = GOCore.best_manifold(exchange_grid)
from_atmosphere = Regridder(manifold, exchange_grid, spectral_grid)
to_atmosphere = Regridder(manifold, spectral_grid, exchange_grid)
regridder = (; to_atmosphere, from_atmosphere)

state = (; u = Field{Center, Center, Nothing}(exchange_grid),
v = Field{Center, Center, Nothing}(exchange_grid),
T = Field{Center, Center, Nothing}(exchange_grid),
Expand All @@ -39,18 +40,23 @@ function ComponentExchanger(atmosphere::SpeedySimulation, exchange_grid)
ℐꜜˢʷ = Field{Center, Center, Nothing}(exchange_grid),
ℐꜜˡʷ = Field{Center, Center, Nothing}(exchange_grid),
Jᶜ = Field{Center, Center, Nothing}(exchange_grid))

return ComponentExchanger(state, regridder)
end

@inline (regrid!::XESMF.Regridder)(field::Oceananigans.Field, data::AbstractArray) = regrid!(vec(interior(field)), data)
@inline (regrid!::XESMF.Regridder)(data::AbstractArray, field::Oceananigans.Field) = regrid!(data, vec(interior(field)))
function regrid!(field::Oceananigans.Field, regridder::Regridder, data::AbstractArray)
regrid!(vec(interior(field)), regridder, vec(data))
end

function regrid!(data::AbstractArray, regridder::Regridder, field::Oceananigans.Field)
regrid!(vec(data), regridder, vec(interior(field)))
end

# Regrid the atmospheric state on the exchange grid
function interpolate_state!(exchanger, exchange_grid, atmos::SpeedySimulation, coupled_model)
regrid! = exchanger.regridder.from_atmosphere
exchange_state = exchanger.state
surface_layer = atmos.model.spectral_grid.nlayers
from_atmosphere = exchanger.regridder.from_atmosphere
exchange_state = exchanger.state
surface_layer = atmos.model.spectral_grid.nlayers

ua = RingGrids.field_view(atmos.variables.grid.u, :, surface_layer).data
va = RingGrids.field_view(atmos.variables.grid.v, :, surface_layer).data
Expand All @@ -61,14 +67,14 @@ function interpolate_state!(exchanger, exchange_grid, atmos::SpeedySimulation, c
ℐꜜˡʷ = atmos.variables.parameterizations.surface_longwave_down.data
Jᶜ = atmos.variables.parameterizations.rain_rate.data .+ atmos.variables.parameterizations.snow_rate.data

regrid!(exchange_state.u, ua)
regrid!(exchange_state.v, va)
regrid!(exchange_state.T, Ta)
regrid!(exchange_state.q, qa)
regrid!(exchange_state.p, pa)
regrid!(exchange_state.ℐꜜˢʷ, ℐꜜˢʷ)
regrid!(exchange_state.ℐꜜˡʷ, ℐꜜˡʷ)
regrid!(exchange_state.Jᶜ, Jᶜ)
regrid!(exchange_state.u, from_atmosphere, ua)
regrid!(exchange_state.v, from_atmosphere, va)
regrid!(exchange_state.T, from_atmosphere, Ta)
regrid!(exchange_state.q, from_atmosphere, qa)
regrid!(exchange_state.p, from_atmosphere, pa)
regrid!(exchange_state.ℐꜜˢʷ, from_atmosphere, ℐꜜˢʷ)
regrid!(exchange_state.ℐꜜˡʷ, from_atmosphere, ℐꜜˡʷ)
regrid!(exchange_state.Jᶜ, from_atmosphere, Jᶜ)

arch = architecture(exchange_grid)

Expand All @@ -90,17 +96,17 @@ end

@kernel function _rotate_winds!(u, v, grid)
i, j = @index(Global, NTuple)
kᴺ = size(grid, 3)
kᴺ = size(grid, 3)
uₑ, vₑ = intrinsic_vector(i, j, kᴺ, grid, u, v)
@inbounds u[i, j, kᴺ] = uₑ
@inbounds v[i, j, kᴺ] = vₑ
end

# TODO: Fix the coupling with the sea ice model and make sure that
# TODO: Fix the coupling with the sea ice model and make sure that
# the this function works also for sea_ice=nothing and on GPUs without
# needing to allocate memory.
function update_net_fluxes!(coupled_model, atmos::SpeedySimulation)
regrid! = coupled_model.interfaces.exchanger.atmosphere.regridder.to_atmosphere
to_atmosphere = coupled_model.interfaces.exchanger.atmosphere.regridder.to_atmosphere
ao_fluxes = coupled_model.interfaces.atmosphere_ocean_interface.fluxes
ai_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes

Expand All @@ -120,16 +126,16 @@ function update_net_fluxes!(coupled_model, atmos::SpeedySimulation)
# TODO: Figure out how we are going to deal with upwelling radiation
# TODO: regrid longwave rather than a mixed surface temperature
# TODO: This does not work on GPUs!!
regrid!(𝒬ᵀ_speedy, vec(interior(𝒬ᵀᵃᵒ) .* (1 .- ℵ) .+ ℵ .* interior(𝒬ᵀᵃⁱ)))
regrid!(Jᵛ_speedy, vec(interior(Jᵛᵃᵒ) .* (1 .- ℵ) .+ ℵ .* interior(Jᵛᵃⁱ)))
regrid!(sst, vec(interior(To) .* (1 .- ℵ) .+ ℵ .* interior(Ti) .+ 273.15))
regrid!(𝒬ᵀ_speedy, to_atmosphere, vec(interior(𝒬ᵀᵃᵒ) .* (1 .- ℵ) .+ ℵ .* interior(𝒬ᵀᵃⁱ)))
regrid!(Jᵛ_speedy, to_atmosphere, vec(interior(Jᵛᵃᵒ) .* (1 .- ℵ) .+ ℵ .* interior(Jᵛᵃⁱ)))
regrid!(sst, to_atmosphere, vec(interior(To) .* (1 .- ℵ) .+ ℵ .* interior(Ti) .+ 273.15))

return nothing
end

# Simple case -> there is no sea ice!
function update_net_fluxes!(coupled_model::SpeedyNoSeaIceEarthSystemModel, atmos::SpeedySimulation)
regrid! = coupled_model.interfaces.exchanger.atmosphere.regridder.to_atmosphere
to_atmosphere = coupled_model.interfaces.exchanger.atmosphere.regridder.to_atmosphere
ao_fluxes = coupled_model.interfaces.atmosphere_ocean_interface.fluxes
𝒬ᵀᵃᵒ = ao_fluxes.sensible_heat
Jᵛᵃᵒ = ao_fluxes.water_vapor
Expand All @@ -142,9 +148,9 @@ function update_net_fluxes!(coupled_model::SpeedyNoSeaIceEarthSystemModel, atmos

# TODO: Figure out how we are going to deal with upwelling radiation
# TODO: This does not work on GPUs!!
regrid!(𝒬ᵀ_speedy, vec(interior(𝒬ᵀᵃᵒ)))
regrid!(Jᵛ_speedy, vec(interior(Jᵛᵃᵒ)))
regrid!(sst, vec(interior(To) .+ 273.15))
regrid!(𝒬ᵀ_speedy, to_atmosphere, vec(interior(𝒬ᵀᵃᵒ)))
regrid!(Jᵛ_speedy, to_atmosphere, vec(interior(Jᵛᵃᵒ)))
regrid!(sst, to_atmosphere, vec(interior(To) .+ 273.15))

return nothing
end
2 changes: 0 additions & 2 deletions test/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
WorldOceanAtlasTools = "04f20302-f1b9-11e8-29d9-7d841cb0a64a"
XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"

[sources]
NumericalEarth = {path = ".."}
Expand Down Expand Up @@ -56,7 +55,6 @@ Statistics = "<0.0.1, 1"
Test = "<0.0.1, 1"
Thermodynamics = "0.15.3"
WorldOceanAtlasTools = "0.6"
XESMF = "0.1.6"
julia = "1.10"

[extras]
Expand Down
8 changes: 0 additions & 8 deletions test/runtests_setup.jl
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,6 @@ test_fields = Dict(
EN4Monthly() => (:T, :S),
)

# Install XESMF for windows architectures
if Sys.iswindows() && Sys.ARCH == :x86_64
@info "Installing XESMF for windows architectures"
using Pkg
Pkg.add("CondaPkg")
using CondaPkg
CondaPkg.add(["esmf", "esmpy"])
end

#####
##### Test utilities
Expand Down
2 changes: 1 addition & 1 deletion test/test_speedy_coupling.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using SpeedyWeather, XESMF
using SpeedyWeather
using NumericalEarth
using Oceananigans
using Dates
Expand Down
Loading