diff --git a/Project.toml b/Project.toml index ab7f75323..465741d10 100644 --- a/Project.toml +++ b/Project.toml @@ -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" @@ -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" @@ -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" diff --git a/docs/Project.toml b/docs/Project.toml index 82347bf37..17af42830 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -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 = ".."} diff --git a/docs/make.jl b/docs/make.jl index 2e5c289a4..4ac8fe534 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -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), diff --git a/examples/global_climate_simulation.jl b/examples/global_climate_simulation.jl index 6fbac8d07..9b3ceeef8 100644 --- a/examples/global_climate_simulation.jl +++ b/examples/global_climate_simulation.jl @@ -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 diff --git a/ext/NumericalEarthSpeedyWeatherExt/NumericalEarthSpeedyWeatherExt.jl b/ext/NumericalEarthSpeedyWeatherExt/NumericalEarthSpeedyWeatherExt.jl index a583861fb..a980103a3 100644 --- a/ext/NumericalEarthSpeedyWeatherExt/NumericalEarthSpeedyWeatherExt.jl +++ b/ext/NumericalEarthSpeedyWeatherExt/NumericalEarthSpeedyWeatherExt.jl @@ -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 diff --git a/ext/NumericalEarthSpeedyWeatherExt/speedy_regridder.jl b/ext/NumericalEarthSpeedyWeatherExt/speedy_regridder.jl deleted file mode 100644 index 85b2db32c..000000000 --- a/ext/NumericalEarthSpeedyWeatherExt/speedy_regridder.jl +++ /dev/null @@ -1,49 +0,0 @@ -using Oceananigans.Grids: AbstractGrid -using Oceananigans - -import XESMF: Regridder, xesmf_coordinates - -const Grids = Union{SpeedyWeather.SpectralGrid, AbstractGrid} - -function Regridder(src::Grids, dst::Grids; method::String="bilinear", periodic=true) - src_coords = xesmf_coordinates(src, Center(), Center(), Center()) - dst_coords = xesmf_coordinates(dst, Center(), Center(), Center()) - - return XESMF.Regridder(src_coords, dst_coords; method, periodic) -end - -two_dimensionalize(lat::Matrix, lon::Matrix) = lat, lon - -function two_dimensionalize(lat::AbstractVector, lon::AbstractVector) - Nx = length(lon) - Ny = length(lat) - lat = repeat(lat', Nx) - lon = repeat(lon, 1, Ny) - return lat, lon -end - -xesmf_coordinates(grid::SpeedyWeather.SpectralGrid, args...) = xesmf_coordinates(grid.grid, args...) -xesmf_coordinates(grid::SpeedyWeather.RingGrids.AbstractGrid, args...) = - throw(ArgumentError("xesmf_coordinates not implemented for grid type $(typeof(grid)), maybe you meant to pass a FullGrid?")) - -function xesmf_coordinates(grid::SpeedyWeather.RingGrids.AbstractFullGrid, args...) - lon = RingGrids.get_lond(grid) - lat = RingGrids.get_latd(grid) - dlon = lon[2] - lon[1] - - lat_b = [90, 0.5 .* (lat[1:end-1] .+ lat[2:end])..., -90] - lon_b = [lon[1] - dlon / 2, lon .+ dlon / 2...] - - lat, lon = two_dimensionalize(lat, lon) - lat_b, lon_b = two_dimensionalize(lat_b, lon_b) - - # Python's xESMF expects 2D arrays with (x, y) coordinates - # in which y varies in dim=1 and x varies in dim=2 - # therefore we transpose the coordinate matrices - coords_dictionary = Dict("lat" => permutedims(lat, (2, 1)), # φ is latitude - "lon" => permutedims(lon, (2, 1)), # λ is longitude - "lat_b" => permutedims(lat_b, (2, 1)), - "lon_b" => permutedims(lon_b, (2, 1))) - - return coords_dictionary -end diff --git a/ext/NumericalEarthSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/NumericalEarthSpeedyWeatherExt/speedy_weather_exchanger.jl index ec0dfa3e7..9a7b698e6 100644 --- a/ext/NumericalEarthSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/NumericalEarthSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -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 @@ -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), @@ -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 @@ -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) @@ -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 @@ -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 @@ -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 diff --git a/test/Project.toml b/test/Project.toml index 54e377ccc..1e0918dc2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -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 = ".."} @@ -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] diff --git a/test/runtests_setup.jl b/test/runtests_setup.jl index a871a354a..d0f04b42c 100644 --- a/test/runtests_setup.jl +++ b/test/runtests_setup.jl @@ -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 diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index 59cf8c58b..dc4f8e652 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -1,4 +1,4 @@ -using SpeedyWeather, XESMF +using SpeedyWeather using NumericalEarth using Oceananigans using Dates