From 54548abac17a13f4edba2149714dd2279b7fb5bf Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 9 Apr 2025 18:21:46 -0400 Subject: [PATCH 01/56] Add SpeedyWeather extension --- ext/ClimaOceanSpeedyWeatherExt.jl | 104 ++++++++++++++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 ext/ClimaOceanSpeedyWeatherExt.jl diff --git a/ext/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt.jl new file mode 100644 index 000000000..9d0e8aa68 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt.jl @@ -0,0 +1,104 @@ +module ClimaOceanSpeedyWeatherExt + +using OffsetArrays +using KernelAbstractions +using Statistics + +import SpeedyWeather as SW +import ClimaOcean as CO +import Oceananigans as OC +import CUDA + +# This document lays out the functions that must be extended to +# use an atmospheric simulation in ClimaOcean. + +import Oceananigans.TimeSteppers: time_step! +import Oceananigans.Models: update_model_field_time_series! + +# Make sure the atmospheric parameters from SpeedyWeather can be used in the compute fluxes function +import ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: + thermodynamics_parameters, + boundary_layer_height, + surface_layer_height + +import ClimaOcean.OceanSeaIceModels: + compute_net_atmosphere_fluxes! + +import ClimaOcean.OceanSeaIceModels.InterfaceComputations: + atmosphere_exchanger, + initialize!, + StateExchanger, + interpolate_atmosphere_state! + +using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel + +# const ClimaCoupledModel = OceanSeaIceModel{<:Any, <:SpeedySimulation} + +# This can be left blank: +update_model_field_time_series!(::SpeedySimulation, time) = nothing + +# Take one time-step +function time_step!(atmos::SpeedySimulation, Δt) + # TODO: check if the time-step can be changed. + @assert Δt == atmos.integrator.dt + CA.SciMLBase.step!(atmos.integrator) + return nothing +end + +# The height of near-surface variables used in the turbulent flux solver +surface_layer_height(s::SpeedySimulation) = 10 # meters, for example + +# This is a parameter that is used in the computation of the fluxes, +# It probably should not be here but in the similarity theory type. +boundary_layer_height(atmos::SpeedySimulation) = 600 + +# Note: possibly, can use the atmos thermodynamic parameters directly here. +thermodynamics_parameters(atmos::SpeedySimulation) = + atmos.integrator.p.params.thermodynamics_params + +Base.summary(::SpeedySimulation) = "ClimaAtmos.AtmosSimulation" + +using Oceananigans.Grids: λnodes, φnodes, LatitudeLongitudeGrid +using Oceananigans.Fields: Center +using Thermodynamics + +""" + interpolate_atmospheric_state!(surface_atmosphere_state, + interpolated_prescribed_freshwater_flux, + atmos::ClimaAtmosSimulation, + grid, clock) + +Interpolate the atmospheric state in `atmos` to `surface_atmospheric_state`, a +the collection of `Field`s needed to compute turbulent fluxes. +""" +function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, coupled_model) + + interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp + exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state + + ue = parent(exchange_atmosphere_state.u) + ve = parent(exchange_atmosphere_state.v) + Te = parent(exchange_atmosphere_state.T) + qe = parent(exchange_atmosphere_state.q) + pe = parent(exchange_atmosphere_state.p) + + ue = dropdims(ue, dims=3) + ve = dropdims(ve, dims=3) + Te = dropdims(Te, dims=3) + qe = dropdims(qe, dims=3) + pe = dropdims(pe, dims=3) + + return nothing +end + +function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid) + return something +end + +initialize!(::StateExchanger, ::SpeedySimulation) = nothing + +function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) + return nothing +end + +end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file From 36a6bdf9217ece106242a15539d9754b0121faaf Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 9 Apr 2025 22:47:25 -0400 Subject: [PATCH 02/56] add a few things --- ext/ClimaOceanSpeedyWeatherExt.jl | 52 +++++++++++++++++++++++++++---- 1 file changed, 46 insertions(+), 6 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt.jl index 9d0e8aa68..2b34fd2cd 100644 --- a/ext/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt.jl @@ -32,7 +32,9 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel -# const ClimaCoupledModel = OceanSeaIceModel{<:Any, <:SpeedySimulation} +const SpeedySimulation = SpeedyWeather.Simulation +const ClimaCoupledModel = OceanSeaIceModel{<:Any, <:SpeedySimulation} +Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" # This can be left blank: update_model_field_time_series!(::SpeedySimulation, time) = nothing @@ -46,17 +48,22 @@ function time_step!(atmos::SpeedySimulation, Δt) end # The height of near-surface variables used in the turbulent flux solver -surface_layer_height(s::SpeedySimulation) = 10 # meters, for example +function surface_layer_height(s::SpeedySimulation) + T = s.model.atmosphere.temp_ref + g = s.model.planet.gravity + Φ = s.model.geopotential.Δp_geopot_full + return Φ[end] * T / g +end # This is a parameter that is used in the computation of the fluxes, # It probably should not be here but in the similarity theory type. boundary_layer_height(atmos::SpeedySimulation) = 600 -# Note: possibly, can use the atmos thermodynamic parameters directly here. -thermodynamics_parameters(atmos::SpeedySimulation) = - atmos.integrator.p.params.thermodynamics_params +Base.eltype(::EarthAtmosphere{FT}) where FT = FT -Base.summary(::SpeedySimulation) = "ClimaAtmos.AtmosSimulation" +# This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather +thermodynamics_parameters(atmos::SpeedyWeather.Simulation) = + PrescribedAtmosphereThermodynamicsParameters(eltype(atmos.model.atmosphere)) using Oceananigans.Grids: λnodes, φnodes, LatitudeLongitudeGrid using Oceananigans.Fields: Center @@ -73,6 +80,26 @@ the collection of `Field`s needed to compute turbulent fluxes. """ function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, coupled_model) + # Get the atmospheric state on the ocean grid + # ua = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.u) + # va = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.v) + # Ta = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.T) + # qa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.q) + # pa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.p) + # Qs = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qs) + # Qℓ = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qℓ) + # Mp = on_architecture(Oceananigans.CPU(), interpolated_prescribed_freshwater_flux) + # ρf = fluxes.freshwater_density + + # RingGrids.interpolate!(vec(view(ua, :, :, 1)), atmos.diagnostic_variables.grid.u_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(va, :, :, 1)), atmos.diagnostic_variables.grid.v_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(Ta, :, :, 1)), atmos.diagnostic_variables.grid.temp_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(qa, :, :, 1)), atmos.diagnostic_variables.grid.humid_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(pa, :, :, 1)), exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]), interpolator) + # RingGrids.interpolate!(vec(view(Qs, :, :, 1)), atmos.diagnostic_variables.physics.surface_shortwave_down, interpolator) + # RingGrids.interpolate!(vec(view(Qℓ, :, :, 1)), atmos.diagnostic_variables.physics.surface_longwave_down, interpolator) + # RingGrids.interpolate!(vec(view(Mp, :, :, 1)), atmosphere_precipitation, interpolator) + interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state @@ -92,6 +119,19 @@ function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, end function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid) + cpu_grid = on_architecture(CPU(), exchange_grid) + cpu_surface_state = ( + u = Field{Center, Center, Nothing}(cpu_grid), + v = Field{Center, Center, Nothing}(cpu_grid), + T = Field{Center, Center, Nothing}(cpu_grid), + q = Field{Center, Center, Nothing}(cpu_grid), + p = Field{Center, Center, Nothing}(cpu_grid), + Qs = Field{Center, Center, Nothing}(cpu_grid), + Qℓ = Field{Center, Center, Nothing}(cpu_grid), + ) + + exchanger = (; cpu_surface_state) + return something end From 6b19679cf8cff28b1dacdfaf8ef55b83ec1a8946 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 10 Apr 2025 11:29:21 +0200 Subject: [PATCH 03/56] start populating the model --- .../ClimaOceanSpeedyWeatherExt.jl | 14 +++++ .../conservative_regridding.jl | 59 +++++++++++++++++++ .../speedy_weather_fluxes.jl} | 44 +++++--------- 3 files changed, 88 insertions(+), 29 deletions(-) create mode 100644 ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl create mode 100644 ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl rename ext/{ClimaOceanSpeedyWeatherExt.jl => ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl} (84%) diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl new file mode 100644 index 000000000..4bb8e6a43 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -0,0 +1,14 @@ +module ClimaOceanSpeedyWeatherExt + +using OffsetArrays +using KernelAbstractions +using Statistics + +import SpeedyWeather +import ClimaOcean +import Oceananigans + +include("conservative_regridding.jl") +include("speedy_weather_fluxes.jl") + +end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl new file mode 100644 index 000000000..482903aa9 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl @@ -0,0 +1,59 @@ +using ConservativeRegridding + +using Oceananigans +using Oceananigans.Utils +using Oceananigans.Grids: λnode, φnode, on_architecture, AbstractGrid +using GeometryOps: Point2 + +using KernelAbstractions: @kernel, @index + +""" + list_cell_vertices(grid) + +Returns a list representing all horizontal grid cells in a curvilinear `grid`. +The output is an Array of 6 * M `Point2` elements where `M = Nx * Ny`. Each row lists the vertices associated with a +horizontal cell in clockwise order starting from the southwest (bottom left) corner. +""" +function list_cell_vertices(grid; add_nans=true) + Nx, Ny, _ = size(grid) + FT = eltype(grid) + + cpu_grid = on_architecture(Oceananigans.CPU(), grid) + + sw = fill(Point2{FT}(0, 0), 1, Nx*Ny) + nw = fill(Point2{FT}(0, 0), 1, Nx*Ny) + ne = fill(Point2{FT}(0, 0), 1, Nx*Ny) + se = fill(Point2{FT}(0, 0), 1, Nx*Ny) + nan = fill(Point2{FT}(NaN, NaN), 1, Nx*Ny) + + launch!(Oceananigans.CPU(), cpu_grid, :xy, _get_vertices!, sw, nw, ne, se, grid) + + vertices = vcat(sw, nw, ne, se, sw) + + if add_nans + vertices = vcat(vertices, nan) + end + + return vertices +end + +@kernel function _get_vertices!(sw, nw, ne, se, grid) + i, j = @index(Global, NTuple) + + FT = eltype(grid) + Nx = size(grid, 1) + λ⁻⁻ = λnode(i, j, 1, grid, Face(), Face(), nothing) + λ⁺⁻ = λnode(i, j+1, 1, grid, Face(), Face(), nothing) + λ⁻⁺ = λnode(i+1, j, 1, grid, Face(), Face(), nothing) + λ⁺⁺ = λnode(i+1, j+1, 1, grid, Face(), Face(), nothing) + + φ⁻⁻ = φnode(i, j, 1, grid, Face(), Face(), nothing) + φ⁺⁻ = φnode(i, j+1, 1, grid, Face(), Face(), nothing) + φ⁻⁺ = φnode(i+1, j, 1, grid, Face(), Face(), nothing) + φ⁺⁺ = φnode(i+1, j+1, 1, grid, Face(), Face(), nothing) + + sw[i+(j-1)*Nx] = Point2{FT}(λ⁻⁻, φ⁻⁻) + nw[i+(j-1)*Nx] = Point2{FT}(λ⁻⁺, φ⁻⁺) + ne[i+(j-1)*Nx] = Point2{FT}(λ⁺⁺, φ⁺⁺) + se[i+(j-1)*Nx] = Point2{FT}(λ⁺⁻, φ⁺⁻) +end \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl similarity index 84% rename from ext/ClimaOceanSpeedyWeatherExt.jl rename to ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl index 2b34fd2cd..b491c7bff 100644 --- a/ext/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl @@ -1,17 +1,5 @@ -module ClimaOceanSpeedyWeatherExt - -using OffsetArrays -using KernelAbstractions -using Statistics - -import SpeedyWeather as SW -import ClimaOcean as CO -import Oceananigans as OC -import CUDA - # This document lays out the functions that must be extended to # use an atmospheric simulation in ClimaOcean. - import Oceananigans.TimeSteppers: time_step! import Oceananigans.Models: update_model_field_time_series! @@ -33,7 +21,7 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel const SpeedySimulation = SpeedyWeather.Simulation -const ClimaCoupledModel = OceanSeaIceModel{<:Any, <:SpeedySimulation} +const SpeedyCoupledModel = OceanSeaIceModel{<:Any, <:SpeedySimulation} Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" # This can be left blank: @@ -59,7 +47,7 @@ end # It probably should not be here but in the similarity theory type. boundary_layer_height(atmos::SpeedySimulation) = 600 -Base.eltype(::EarthAtmosphere{FT}) where FT = FT +# Base.eltype(::EarthAtmosphere{FT}) where FT = FT # This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather thermodynamics_parameters(atmos::SpeedyWeather.Simulation) = @@ -100,20 +88,20 @@ function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, # RingGrids.interpolate!(vec(view(Qℓ, :, :, 1)), atmos.diagnostic_variables.physics.surface_longwave_down, interpolator) # RingGrids.interpolate!(vec(view(Mp, :, :, 1)), atmosphere_precipitation, interpolator) - interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp - exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state + # interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp + # exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state - ue = parent(exchange_atmosphere_state.u) - ve = parent(exchange_atmosphere_state.v) - Te = parent(exchange_atmosphere_state.T) - qe = parent(exchange_atmosphere_state.q) - pe = parent(exchange_atmosphere_state.p) + # ue = parent(exchange_atmosphere_state.u) + # ve = parent(exchange_atmosphere_state.v) + # Te = parent(exchange_atmosphere_state.T) + # qe = parent(exchange_atmosphere_state.q) + # pe = parent(exchange_atmosphere_state.p) - ue = dropdims(ue, dims=3) - ve = dropdims(ve, dims=3) - Te = dropdims(Te, dims=3) - qe = dropdims(qe, dims=3) - pe = dropdims(pe, dims=3) + # ue = dropdims(ue, dims=3) + # ve = dropdims(ve, dims=3) + # Te = dropdims(Te, dims=3) + # qe = dropdims(qe, dims=3) + # pe = dropdims(pe, dims=3) return nothing end @@ -139,6 +127,4 @@ initialize!(::StateExchanger, ::SpeedySimulation) = nothing function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) return nothing -end - -end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file +end \ No newline at end of file From e3fc3e8acac701188e61216fffe55739597e580c Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 10 Apr 2025 11:29:44 +0200 Subject: [PATCH 04/56] change project --- Project.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/Project.toml b/Project.toml index 9b43e5a5f..4b86a7729 100644 --- a/Project.toml +++ b/Project.toml @@ -9,10 +9,13 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CFTime = "179af706-886a-5703-950a-314cd64e0468" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4" +ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343" CubicSplines = "9c784101-8907-5a6d-9be6-98f00873c89b" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +GeometryOps = "3251bfac-6a57-4b6d-aa61-ac1fef2975ab" +GeometryOpsCore = "05efe853-fabf-41c8-927e-7063c8b9f013" ImageMorphology = "787d08f9-d448-5407-9aad-5290dd7ab264" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -24,6 +27,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Scratch = "6c6a2e73-6563-6170-7368-637461726353" SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" +SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f" @@ -34,6 +38,7 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" [extensions] ClimaOceanReactantExt = "Reactant" +ClimaOceanSpeedyWeatherExt = "SpeedyWeather" [compat] Adapt = "4" From 8fc92aae9c21e3eeb1911e80bb354f9413cde812 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 10 Apr 2025 09:17:34 -0400 Subject: [PATCH 05/56] add speedy --- Project.toml | 2 ++ test/test_speedy_coupling.jl | 12 ++++++++++++ 2 files changed, 14 insertions(+) create mode 100644 test/test_speedy_coupling.jl diff --git a/Project.toml b/Project.toml index 9b43e5a5f..7d6af08d9 100644 --- a/Project.toml +++ b/Project.toml @@ -24,6 +24,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Scratch = "6c6a2e73-6563-6170-7368-637461726353" SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" +SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f" @@ -54,6 +55,7 @@ PrecompileTools = "1" Reactant = "0.2.45" Scratch = "1" SeawaterPolynomials = "0.3.5" +SpeedyWeather = "0.13.0" StaticArrays = "1" Statistics = "1.9" SurfaceFluxes = "0.11, 0.12" diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl new file mode 100644 index 000000000..7983436ec --- /dev/null +++ b/test/test_speedy_coupling.jl @@ -0,0 +1,12 @@ +import SpeedyWeather +import ClimaOcean +using Oceananigans + +arch = CPU() +Nx = 180 +Ny = 85 +Nz = 10 +longitude = (0, 360) +latitude = (-85, 85) +grid = LatitudeLongitudeGrid(arch; size=(Nx, Ny, Nz), longitude, latitude, z=(-100, 0)) +ocean = ClimaOcean.ocean_simulation(grid) \ No newline at end of file From b702959b2397a1cdd8e683006bf0729f3a3d1a3e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 10 Apr 2025 16:50:49 -0400 Subject: [PATCH 06/56] building --- Project.toml | 5 +- .../ClimaOceanSpeedyWeatherExt.jl | 3 +- .../speedy_atmosphere_simulations.jl | 48 +++++++++++++ .../speedy_weather_fluxes.jl | 72 +++++++------------ src/ClimaOcean.jl | 2 + test/test_speedy_coupling.jl | 17 ++++- 6 files changed, 94 insertions(+), 53 deletions(-) create mode 100644 ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl diff --git a/Project.toml b/Project.toml index 3181a2a98..2e6c24360 100644 --- a/Project.toml +++ b/Project.toml @@ -9,7 +9,6 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CFTime = "179af706-886a-5703-950a-314cd64e0468" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4" -ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343" CubicSplines = "9c784101-8907-5a6d-9be6-98f00873c89b" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" @@ -59,7 +58,7 @@ PrecompileTools = "1" Reactant = "0.2.45" Scratch = "1" SeawaterPolynomials = "0.3.5" -SpeedyWeather = "0.13.0" +SpeedyWeather = "0.14.0" StaticArrays = "1" Statistics = "1.9" SurfaceFluxes = "0.11, 0.12" @@ -73,4 +72,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll"] diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index 4bb8e6a43..9666d53ca 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -8,7 +8,8 @@ import SpeedyWeather import ClimaOcean import Oceananigans -include("conservative_regridding.jl") +# include("conservative_regridding.jl") +include("speedy_atmosphere_simulations.jl") include("speedy_weather_fluxes.jl") end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl new file mode 100644 index 000000000..8aa9453f9 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -0,0 +1,48 @@ +# Make sure the atmospheric parameters from SpeedyWeather can be used in the compute fluxes function +import ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: + thermodynamics_parameters, + boundary_layer_height, + surface_layer_height + +import ClimaOcean: atmosphere_simulation + +# This can be left blank: +update_model_field_time_series!(::SpeedySimulation, time) = nothing + +# Take one time-step +time_step!(atmos::SpeedySimulation, Δt) = SpeedyWeather.timestep!(atmos) + +function atmosphere_simulation(grid::SpeedyWeather.SpectralGrid) + # orography = zeros(grid.Grid, grid.nlat_half)) + + surface_heat_flux = SpeedyWeather.PrescribedOceanHeatFlux() + surface_evaporation = SpeedyWeather.PrescribedOceanEvaporation() + atmos_model = SpeedyWeather.PrimitiveWetModel(grid; + # orography, + surface_heat_flux, + surface_evaporation) + + atmos_sim = SpeedyWeather.initialize!(atmos_model) + return atmos_sim +end + +# The height of near-surface variables used in the turbulent flux solver +function surface_layer_height(s::SpeedySimulation) + T = s.model.atmosphere.temp_ref + g = s.model.planet.gravity + Φ = s.model.geopotential.Δp_geopot_full + return Φ[end] * T / g +end + +# This is a parameter that is used in the computation of the fluxes, +# It probably should not be here but in the similarity theory type. +boundary_layer_height(atmos::SpeedySimulation) = 600 + +# Base.eltype(::EarthAtmosphere{FT}) where FT = FT + +# This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather +thermodynamics_parameters(atmos::SpeedyWeather.Simulation) = + ClimaOcean.OceanSeaIceModels.PrescribedAtmosphereThermodynamicsParameters(Float32) + + + diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl index b491c7bff..f9e3a95c7 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl @@ -3,12 +3,6 @@ import Oceananigans.TimeSteppers: time_step! import Oceananigans.Models: update_model_field_time_series! -# Make sure the atmospheric parameters from SpeedyWeather can be used in the compute fluxes function -import ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: - thermodynamics_parameters, - boundary_layer_height, - surface_layer_height - import ClimaOcean.OceanSeaIceModels: compute_net_atmosphere_fluxes! @@ -24,37 +18,7 @@ const SpeedySimulation = SpeedyWeather.Simulation const SpeedyCoupledModel = OceanSeaIceModel{<:Any, <:SpeedySimulation} Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" -# This can be left blank: -update_model_field_time_series!(::SpeedySimulation, time) = nothing - -# Take one time-step -function time_step!(atmos::SpeedySimulation, Δt) - # TODO: check if the time-step can be changed. - @assert Δt == atmos.integrator.dt - CA.SciMLBase.step!(atmos.integrator) - return nothing -end - -# The height of near-surface variables used in the turbulent flux solver -function surface_layer_height(s::SpeedySimulation) - T = s.model.atmosphere.temp_ref - g = s.model.planet.gravity - Φ = s.model.geopotential.Δp_geopot_full - return Φ[end] * T / g -end - -# This is a parameter that is used in the computation of the fluxes, -# It probably should not be here but in the similarity theory type. -boundary_layer_height(atmos::SpeedySimulation) = 600 - -# Base.eltype(::EarthAtmosphere{FT}) where FT = FT - -# This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather -thermodynamics_parameters(atmos::SpeedyWeather.Simulation) = - PrescribedAtmosphereThermodynamicsParameters(eltype(atmos.model.atmosphere)) - -using Oceananigans.Grids: λnodes, φnodes, LatitudeLongitudeGrid -using Oceananigans.Fields: Center +using Oceananigans using Thermodynamics """ @@ -68,6 +32,18 @@ the collection of `Field`s needed to compute turbulent fluxes. """ function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, coupled_model) + # Plan: + # 1. transfer atmos data to GPU (req ability to represent atmos on GPU) + # 2. interpolate from atmos grid to ocean grid + + # RingGrids.interpolate!(vec(view(ua, :, :, 1)), atmos.diagnostic_variables.grid.u_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(va, :, :, 1)), atmos.diagnostic_variables.grid.v_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(Ta, :, :, 1)), atmos.diagnostic_variables.grid.temp_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(qa, :, :, 1)), atmos.diagnostic_variables.grid.humid_grid[:, end], interpolator) + # RingGrids.interpolate!(vec(view(pa, :, :, 1)), exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]), interpolator) + # RingGrids.interpolate!(vec(view(Qs, :, :, 1)), atmos.diagnostic_variables.physics.surface_shortwave_down, interpolator) + # RingGrids.interpolate!(vec(view(Qℓ, :, :, 1)), atmos.diagnostic_variables.physics.surface_longwave_down, interpolator) + # Get the atmospheric state on the ocean grid # ua = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.u) # va = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.v) @@ -79,15 +55,6 @@ function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, # Mp = on_architecture(Oceananigans.CPU(), interpolated_prescribed_freshwater_flux) # ρf = fluxes.freshwater_density - # RingGrids.interpolate!(vec(view(ua, :, :, 1)), atmos.diagnostic_variables.grid.u_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(va, :, :, 1)), atmos.diagnostic_variables.grid.v_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(Ta, :, :, 1)), atmos.diagnostic_variables.grid.temp_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(qa, :, :, 1)), atmos.diagnostic_variables.grid.humid_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(pa, :, :, 1)), exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]), interpolator) - # RingGrids.interpolate!(vec(view(Qs, :, :, 1)), atmos.diagnostic_variables.physics.surface_shortwave_down, interpolator) - # RingGrids.interpolate!(vec(view(Qℓ, :, :, 1)), atmos.diagnostic_variables.physics.surface_longwave_down, interpolator) - # RingGrids.interpolate!(vec(view(Mp, :, :, 1)), atmosphere_precipitation, interpolator) - # interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp # exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state @@ -118,6 +85,17 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid) Qℓ = Field{Center, Center, Nothing}(cpu_grid), ) + # Figure this out: + spectral_grid = atmosphere.model.spectral_grid + # 1/4 degree? + interpolator = SpeedyWeather.RingGrids.AnvilInterpolator(Float32, + SpeedyWeather.FullClenshawGrid, 90, spectral_grid.npoints) + + arch = exchange_grid.architecture + tmp_grid = LatitudeLongitudeGrid(arch; size=(360, 179, 1), latitude=(-89.5, 89.5), longitude=(0, 360), z = (0, 1)) + londs, latds = SpeedyWeather.RingGrids.get_londlatds(spectral_grid.Grid, spectral_grid.nlat_half) + SpeedyWeather.RingGrids.update_locator!(interpolator, londs, latds) + exchanger = (; cpu_surface_state) return something @@ -127,4 +105,4 @@ initialize!(::StateExchanger, ::SpeedySimulation) = nothing function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) return nothing -end \ No newline at end of file +end diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 6d3eb0e37..260431fd3 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -73,6 +73,8 @@ end return NamedTuple{names}(vals) end +function atmosphere_simulation end + include("OceanSimulations/OceanSimulations.jl") include("SeaIceSimulations.jl") include("OceanSeaIceModels/OceanSeaIceModels.jl") diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index 7983436ec..1b12a2110 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -1,6 +1,7 @@ import SpeedyWeather import ClimaOcean using Oceananigans +using Dates arch = CPU() Nx = 180 @@ -8,5 +9,17 @@ Ny = 85 Nz = 10 longitude = (0, 360) latitude = (-85, 85) -grid = LatitudeLongitudeGrid(arch; size=(Nx, Ny, Nz), longitude, latitude, z=(-100, 0)) -ocean = ClimaOcean.ocean_simulation(grid) \ No newline at end of file +halo = (7, 7, 7) +ocean_grid = LatitudeLongitudeGrid(arch; size=(Nx, Ny, Nz), halo, longitude, latitude, z=(-100, 0)) +ocean = ClimaOcean.ocean_simulation(ocean_grid) + +atmos_grid = SpeedyWeather.SpectralGrid(trunc=31, nlayers=10) +atmosphere = ClimaOcean.atmosphere_simulation(atmos_grid) +SpeedyWeather.set!(atmosphere.model.time_stepping, Δt=Minute(10)) + +# EarthSystemModel +# const OceanSeaIceModel = +# const SeaIceOnlyModel = +# const LandOnlyModel = + +model = ClimaOcean.OceanSeaIceModel(ocean; atmosphere) \ No newline at end of file From 007aed0de35ebcdc0fcc5882264c04f25a41085b Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 10 Apr 2025 17:41:33 -0400 Subject: [PATCH 07/56] grid test --- .../ClimaOceanSpeedyWeatherExt.jl | 9 +++++++++ .../speedy_atmosphere_simulations.jl | 8 ++++++-- .../speedy_weather_fluxes.jl | 7 ------- test/test_speedy_coupling.jl | 20 ++++++++++++++++++- 4 files changed, 34 insertions(+), 10 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index 9666d53ca..d2f577d32 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -8,6 +8,15 @@ import SpeedyWeather import ClimaOcean import Oceananigans +function clenshaw_latitude_longitude_grid(arch, spectral_grid) + grid = LatitudeLongitudeGrid(arch; + size = (360, 179, 1), + latitude = (-89.5, 89.5), + longitude = (0, 360), + z = (0, 1)) + return grid +end + # include("conservative_regridding.jl") include("speedy_atmosphere_simulations.jl") include("speedy_weather_fluxes.jl") diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl index 8aa9453f9..59ea6173c 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -6,11 +6,15 @@ import ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: import ClimaOcean: atmosphere_simulation +const SpeedySimulation = SpeedyWeather.Simulation +const SpeedyCoupledModel = ClimaOcean.OceanSeaIceModel{<:Any, <:SpeedySimulation} +Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" + # This can be left blank: -update_model_field_time_series!(::SpeedySimulation, time) = nothing +Oceananigans.Models.update_model_field_time_series!(::SpeedySimulation, time) = nothing # Take one time-step -time_step!(atmos::SpeedySimulation, Δt) = SpeedyWeather.timestep!(atmos) +Oceananigans.TimeSteppers.time_step!(atmos::SpeedySimulation, Δt) = SpeedyWeather.timestep!(atmos) function atmosphere_simulation(grid::SpeedyWeather.SpectralGrid) # orography = zeros(grid.Grid, grid.nlat_half)) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl index f9e3a95c7..73b77ea4a 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl @@ -1,7 +1,5 @@ # This document lays out the functions that must be extended to # use an atmospheric simulation in ClimaOcean. -import Oceananigans.TimeSteppers: time_step! -import Oceananigans.Models: update_model_field_time_series! import ClimaOcean.OceanSeaIceModels: compute_net_atmosphere_fluxes! @@ -13,11 +11,6 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: interpolate_atmosphere_state! using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel - -const SpeedySimulation = SpeedyWeather.Simulation -const SpeedyCoupledModel = OceanSeaIceModel{<:Any, <:SpeedySimulation} -Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" - using Oceananigans using Thermodynamics diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index 1b12a2110..c1f3f08a9 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -2,6 +2,23 @@ import SpeedyWeather import ClimaOcean using Oceananigans using Dates +using Test + +ClimaOceanSpeedyWeatherExt = Base.get_extension(ClimaOcean, :ClimaOceanSpeedyWeatherExt) +@test !isnothing(ClimaOceanSpeedyWeatherExt) + +spectral_grid = SpeedyWeather.SpectralGrid(trunc=31, nlayers=10) +oceananigans_grid = ClimaOceanSpeedyWeatherExt.clenshaw_latitude_longitude_grid(CPU(), spectral_grid) +λd_sw, φd_sw = SpeedyWeather.RingGrids.get_londlatds(spectral_grid.Grid, spectral_grid.nlat_half) +λd_oc = λnodes(oceananigans_grid, Center(), Center(), Center()) +φd_oc = φnodes(oceananigans_grid, Center(), Center(), Center()) + +@test λd_sw ≈ λd_oc +@test φd_sw ≈ φd_oc + +# SpeedyWeather.RingGrids.update_locator!(interpolator, londs, latds) + +#= arch = CPU() Nx = 180 @@ -22,4 +39,5 @@ SpeedyWeather.set!(atmosphere.model.time_stepping, Δt=Minute(10)) # const SeaIceOnlyModel = # const LandOnlyModel = -model = ClimaOcean.OceanSeaIceModel(ocean; atmosphere) \ No newline at end of file +model = ClimaOcean.OceanSeaIceModel(ocean; atmosphere) +=# \ No newline at end of file From a3aa03f004d21272231c45a38672d159af59a39f Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 14 Apr 2025 09:47:09 +0200 Subject: [PATCH 08/56] better comment --- .../conservative_regridding.jl | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl index 482903aa9..d0520c512 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl @@ -11,10 +11,10 @@ using KernelAbstractions: @kernel, @index list_cell_vertices(grid) Returns a list representing all horizontal grid cells in a curvilinear `grid`. -The output is an Array of 6 * M `Point2` elements where `M = Nx * Ny`. Each row lists the vertices associated with a -horizontal cell in clockwise order starting from the southwest (bottom left) corner. +The output is an Array of 5 * M `Point2` elements where `M = Nx * Ny`. Each column lists the vertices +associated with a horizontal cell in clockwise order starting from the southwest (bottom left) corner. """ -function list_cell_vertices(grid; add_nans=true) +function list_cell_vertices(grid) Nx, Ny, _ = size(grid) FT = eltype(grid) @@ -24,16 +24,11 @@ function list_cell_vertices(grid; add_nans=true) nw = fill(Point2{FT}(0, 0), 1, Nx*Ny) ne = fill(Point2{FT}(0, 0), 1, Nx*Ny) se = fill(Point2{FT}(0, 0), 1, Nx*Ny) - nan = fill(Point2{FT}(NaN, NaN), 1, Nx*Ny) launch!(Oceananigans.CPU(), cpu_grid, :xy, _get_vertices!, sw, nw, ne, se, grid) vertices = vcat(sw, nw, ne, se, sw) - if add_nans - vertices = vcat(vertices, nan) - end - return vertices end @@ -56,4 +51,5 @@ end nw[i+(j-1)*Nx] = Point2{FT}(λ⁻⁺, φ⁻⁺) ne[i+(j-1)*Nx] = Point2{FT}(λ⁺⁺, φ⁺⁺) se[i+(j-1)*Nx] = Point2{FT}(λ⁺⁻, φ⁺⁻) -end \ No newline at end of file +end + From 9e5e4af6949744f034f6508f3a31d9b65cee417d Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 14 Apr 2025 10:23:42 +0200 Subject: [PATCH 09/56] this should start working? --- Project.toml | 2 + .../ClimaOceanSpeedyWeatherExt.jl | 5 +- .../conservative_regridding.jl | 8 +- .../speedy_weather_exchanger.jl | 95 +++++++++++++++++++ .../component_interfaces.jl | 4 +- 5 files changed, 106 insertions(+), 8 deletions(-) create mode 100644 ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl diff --git a/Project.toml b/Project.toml index 2e6c24360..67b9e5585 100644 --- a/Project.toml +++ b/Project.toml @@ -9,10 +9,12 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CFTime = "179af706-886a-5703-950a-314cd64e0468" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4" +ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343" CubicSplines = "9c784101-8907-5a6d-9be6-98f00873c89b" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f" GeometryOps = "3251bfac-6a57-4b6d-aa61-ac1fef2975ab" GeometryOpsCore = "05efe853-fabf-41c8-927e-7063c8b9f013" ImageMorphology = "787d08f9-d448-5407-9aad-5290dd7ab264" diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index d2f577d32..e9f8cb730 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -17,8 +17,9 @@ function clenshaw_latitude_longitude_grid(arch, spectral_grid) return grid end -# include("conservative_regridding.jl") +include("conservative_regridding.jl") include("speedy_atmosphere_simulations.jl") -include("speedy_weather_fluxes.jl") +include("speedy_weather_exchanger.jl") +# include("speedy_weather_fluxes.jl") end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl index d0520c512..e5ebede3f 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl @@ -20,10 +20,10 @@ function list_cell_vertices(grid) cpu_grid = on_architecture(Oceananigans.CPU(), grid) - sw = fill(Point2{FT}(0, 0), 1, Nx*Ny) - nw = fill(Point2{FT}(0, 0), 1, Nx*Ny) - ne = fill(Point2{FT}(0, 0), 1, Nx*Ny) - se = fill(Point2{FT}(0, 0), 1, Nx*Ny) + sw = fill(Point2{FT}(0, 0), 1, Nx*Ny) + nw = fill(Point2{FT}(0, 0), 1, Nx*Ny) + ne = fill(Point2{FT}(0, 0), 1, Nx*Ny) + se = fill(Point2{FT}(0, 0), 1, Nx*Ny) launch!(Oceananigans.CPU(), cpu_grid, :xy, _get_vertices!, sw, nw, ne, se, grid) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl new file mode 100644 index 000000000..ab128f4d5 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -0,0 +1,95 @@ +using ConservativeRegridding +using GeometryOps: CutAtAntimeridianAndPoles, ClosedRing +using GeoInterface: Polygon, LinearRing +using Oceananigans + +function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, exchange_atmosphere_state) + + # Figure this out: + spectral_grid = atmosphere.model.spectral_grid + grid_faces = ConservativeRegridding.get_faces(spectral_grid) + exchange_faces = list_cell_vertices(exchange_grid) + arch = architecture(exchange_grid) + + if arch isa CPU # In case of a CPU grid, we reuse the already allocated fields + cpu_surface_state = ( + u = vec(interior(exchange_atmosphere_state.u)), + v = vec(interior(exchange_atmosphere_state.v)), + T = vec(interior(exchange_atmosphere_state.T)), + q = vec(interior(exchange_atmosphere_state.q)), + p = vec(interior(exchange_atmosphere_state.p)), + Qs = vec(interior(exchange_atmosphere_state.Qs)), + Qℓ = vec(interior(exchange_atmosphere_state.Qℓ)) + ) + else # Otherwise we allocate new CPU fields + cpu_surface_state = ( + u = zeros(Float32, length(exchange_faces)), + v = zeros(Float32, length(exchange_faces)), + T = zeros(Float32, length(exchange_faces)), + q = zeros(Float32, length(exchange_faces)), + p = zeros(Float32, length(exchange_faces)), + Qs = zeros(Float32, length(exchange_faces)), + Qℓ = zeros(Float32, length(exchange_faces)), + ) + end + + regularizing_function = ClosedRing() ∘ CutAtAntimeridianAndPoles() + + # Magical incantation from ConservativeRegridding + polys1 = map(regularizing_function, (Polygon([LinearRing(f)]) for f in eachcol(grid_faces))) + polys2 = map(regularizing_function, (Polygon([LinearRing(f)]) for f in eachcol(exchange_faces))) + + regridder = ConservativeRegridding.intersection_areas(polys1, polys2) + + exchange_areas = vec(sum(A, dims=2)) + atmosphere_areas = vec(sum(A, dims=1)) + + exchanger = (; cpu_surface_state, regridder, exchange_areas, atmosphere_areas) + + return exchanger +end + +fill_exchange_fields!(::Oceananigans.CPU, args...) = nothing + +function fill_exchange_fields!(::Oceananigans.GPU, exchange_state, cpu_surface_state) + # Can I just copyto! here? +end + +function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coupled_model) + # Get the atmospheric state on the ocean grid + exchanger = interfaces.exchanger + regridder = exchanger.regridder + exchange_grid = interfaces.exchanger.exchange_grid + exchange_state = exchanger.exchange_atmosphere_state + arch = architecture(exchange_grid) + + ua = atmos.diagnostic_variables.grid.u_grid[:, end] + va = atmos.diagnostic_variables.grid.v_grid[:, end] + Ta = atmos.diagnostic_variables.grid.temp_grid[:, end] + qa = atmos.diagnostic_variables.grid.humid_grid[:, end] + pa = exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]) + Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down + Qla = atmos.diagnostic_variables.physics.surface_longwave_down + + exchange_areas = exchanger.exchange_areas + + ue = exchanger.cpu_surface_state.u + ve = exchanger.cpu_surface_state.v + Te = exchanger.cpu_surface_state.T + qe = exchanger.cpu_surface_state.q + pe = exchanger.cpu_surface_state.p + Qse = exchanger.cpu_surface_state.Qs + Qle = exchanger.cpu_surface_state.Qℓ + + regrid!(ua, ue, regridder, exchange_areas) + regrid!(va, ve, regridder, exchange_areas) + regrid!(Ta, Te, regridder, exchange_areas) + regrid!(qa, qe, regridder, exchange_areas) + regrid!(pa, pe, regridder, exchange_areas) + regrid!(Qsa, Qse, regridder, exchange_areas) + regrid!(Qla, Qle, regridder, exchange_areas) + + fill_exchange_fields!(arch, exchange_state, exchanger.cpu_surface_state) + + return nothing +end \ No newline at end of file diff --git a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl index ce4ae92f5..4aeed5522 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl @@ -99,12 +99,12 @@ function StateExchanger(ocean::Simulation, atmosphere) # TODO: generalize this exchange_grid = ocean.model.grid exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid) - exchanger = atmosphere_exchanger(atmosphere, exchange_grid) + exchanger = atmosphere_exchanger(atmosphere, exchange_grid, exchange_atmosphere_state) return StateExchanger(ocean.model.grid, exchange_atmosphere_state, exchanger) end -function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid) +function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid, exchange_atmosphere_state) atmos_grid = atmosphere.grid arch = architecture(exchange_grid) Nx, Ny, Nz = size(exchange_grid) From 93d054d117180766f7deae2da2cb6c924d9a8e6e Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 14 Apr 2025 10:29:29 +0200 Subject: [PATCH 10/56] comments --- .../speedy_weather_exchanger.jl | 37 ++++++++++--------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index ab128f4d5..5349fd509 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -41,7 +41,7 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha regridder = ConservativeRegridding.intersection_areas(polys1, polys2) - exchange_areas = vec(sum(A, dims=2)) + exchange_areas = vec(sum(A, dims=2)) atmosphere_areas = vec(sum(A, dims=1)) exchanger = (; cpu_surface_state, regridder, exchange_areas, atmosphere_areas) @@ -51,44 +51,45 @@ end fill_exchange_fields!(::Oceananigans.CPU, args...) = nothing +# TODO: add GPU support function fill_exchange_fields!(::Oceananigans.GPU, exchange_state, cpu_surface_state) # Can I just copyto! here? end +# Regrid the atmospheric state on the exchange grid function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coupled_model) - # Get the atmospheric state on the ocean grid exchanger = interfaces.exchanger regridder = exchanger.regridder exchange_grid = interfaces.exchanger.exchange_grid exchange_state = exchanger.exchange_atmosphere_state - arch = architecture(exchange_grid) - ua = atmos.diagnostic_variables.grid.u_grid[:, end] - va = atmos.diagnostic_variables.grid.v_grid[:, end] - Ta = atmos.diagnostic_variables.grid.temp_grid[:, end] - qa = atmos.diagnostic_variables.grid.humid_grid[:, end] - pa = exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]) + ua = atmos.diagnostic_variables.grid.u_grid[:, end] + va = atmos.diagnostic_variables.grid.v_grid[:, end] + Ta = atmos.diagnostic_variables.grid.temp_grid[:, end] + qa = atmos.diagnostic_variables.grid.humid_grid[:, end] + pa = exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]) Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down Qla = atmos.diagnostic_variables.physics.surface_longwave_down exchange_areas = exchanger.exchange_areas - ue = exchanger.cpu_surface_state.u - ve = exchanger.cpu_surface_state.v - Te = exchanger.cpu_surface_state.T - qe = exchanger.cpu_surface_state.q - pe = exchanger.cpu_surface_state.p + ue = exchanger.cpu_surface_state.u + ve = exchanger.cpu_surface_state.v + Te = exchanger.cpu_surface_state.T + qe = exchanger.cpu_surface_state.q + pe = exchanger.cpu_surface_state.p Qse = exchanger.cpu_surface_state.Qs Qle = exchanger.cpu_surface_state.Qℓ - regrid!(ua, ue, regridder, exchange_areas) - regrid!(va, ve, regridder, exchange_areas) - regrid!(Ta, Te, regridder, exchange_areas) - regrid!(qa, qe, regridder, exchange_areas) - regrid!(pa, pe, regridder, exchange_areas) + regrid!(ua, ue, regridder, exchange_areas) + regrid!(va, ve, regridder, exchange_areas) + regrid!(Ta, Te, regridder, exchange_areas) + regrid!(qa, qe, regridder, exchange_areas) + regrid!(pa, pe, regridder, exchange_areas) regrid!(Qsa, Qse, regridder, exchange_areas) regrid!(Qla, Qle, regridder, exchange_areas) + arch = architecture(exchange_grid) fill_exchange_fields!(arch, exchange_state, exchanger.cpu_surface_state) return nothing From 60c647691e8647f2318c738ae3d631b3aebedfba Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 14 Apr 2025 11:01:26 +0200 Subject: [PATCH 11/56] import stuff --- .../ClimaOceanSpeedyWeatherExt.jl | 2 +- .../speedy_weather_exchanger.jl | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index e9f8cb730..8bb759544 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -18,8 +18,8 @@ function clenshaw_latitude_longitude_grid(arch, spectral_grid) end include("conservative_regridding.jl") -include("speedy_atmosphere_simulations.jl") include("speedy_weather_exchanger.jl") +include("speedy_atmosphere_simulations.jl") # include("speedy_weather_fluxes.jl") end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 5349fd509..54bf0dab0 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -3,6 +3,15 @@ using GeometryOps: CutAtAntimeridianAndPoles, ClosedRing using GeoInterface: Polygon, LinearRing using Oceananigans +import ClimaOcean.OceanSeaIceModels: + compute_net_atmosphere_fluxes! + +import ClimaOcean.OceanSeaIceModels.InterfaceComputations: + atmosphere_exchanger, + initialize!, + StateExchanger, + interpolate_atmosphere_state! + function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, exchange_atmosphere_state) # Figure this out: From 3d313d7b207a6ea08866868d846287b8b26056b5 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 14 Apr 2025 11:28:20 +0200 Subject: [PATCH 12/56] some improvements --- .../ClimaOceanSpeedyWeatherExt.jl | 2 +- .../conservative_regridding.jl | 30 +++++++++++++++-- .../speedy_weather_exchanger.jl | 33 +++++++++++-------- 3 files changed, 48 insertions(+), 17 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index 8bb759544..e9f8cb730 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -18,8 +18,8 @@ function clenshaw_latitude_longitude_grid(arch, spectral_grid) end include("conservative_regridding.jl") -include("speedy_weather_exchanger.jl") include("speedy_atmosphere_simulations.jl") +include("speedy_weather_exchanger.jl") # include("speedy_weather_fluxes.jl") end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl index e5ebede3f..89b3b8b5e 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl @@ -7,14 +7,16 @@ using GeometryOps: Point2 using KernelAbstractions: @kernel, @index +using SpeedyWeather.RingGrids + """ - list_cell_vertices(grid) + get_faces(grid) Returns a list representing all horizontal grid cells in a curvilinear `grid`. The output is an Array of 5 * M `Point2` elements where `M = Nx * Ny`. Each column lists the vertices associated with a horizontal cell in clockwise order starting from the southwest (bottom left) corner. """ -function list_cell_vertices(grid) +function get_faces(grid::AbstractGrid) Nx, Ny, _ = size(grid) FT = eltype(grid) @@ -32,6 +34,30 @@ function list_cell_vertices(grid) return vertices end +# Move to SpeedyWeather? +function get_faces(grid::SpeedyWeather.SpectralGrid) + + Grid = grid.Grid + nlat_half = grid.nlat_half + npoints = RingGrids.get_npoints2D(Grid, nlat_half) + + # vertex east, south, west, north (i.e. clockwise for every grid point) + E, S, W, N = RingGrids.get_vertices(Grid, nlat_half) + + # allocate faces as Point2{Float64} so that no data copy has to be made in Makie + faces = Matrix{NTuple{2, Float64}}(undef, 5, npoints) + + @inbounds for ij in 1:npoints + faces[1, ij] = (E[1, ij], E[2, ij]) # clockwise + faces[2, ij] = (S[1, ij], S[2, ij]) + faces[3, ij] = (W[1, ij], W[2, ij]) + faces[4, ij] = (N[1, ij], N[2, ij]) + faces[5, ij] = (E[1, ij], E[2, ij]) # back to east to close the polygon + end + + return faces +end + @kernel function _get_vertices!(sw, nw, ne, se, grid) i, j = @index(Global, NTuple) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 54bf0dab0..71542c40c 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -2,6 +2,7 @@ using ConservativeRegridding using GeometryOps: CutAtAntimeridianAndPoles, ClosedRing using GeoInterface: Polygon, LinearRing using Oceananigans +using Oceananigans.Grids: architecture import ClimaOcean.OceanSeaIceModels: compute_net_atmosphere_fluxes! @@ -12,15 +13,21 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: StateExchanger, interpolate_atmosphere_state! +# For the moment the workflow is: +# 1. Perform the regridding on the CPU +# 2. Eventually copy the regridded fields to the GPU +# If this work we can +# 1. Copy speedyweather gridarrays to the GPU +# 2. Perform the regridding on the GPU function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, exchange_atmosphere_state) # Figure this out: spectral_grid = atmosphere.model.spectral_grid - grid_faces = ConservativeRegridding.get_faces(spectral_grid) - exchange_faces = list_cell_vertices(exchange_grid) + grid_faces = get_faces(spectral_grid) + exchange_faces = get_faces(exchange_grid) arch = architecture(exchange_grid) - if arch isa CPU # In case of a CPU grid, we reuse the already allocated fields + if arch isa Oceananigans.CPU # In case of a CPU grid, we reuse the already allocated fields cpu_surface_state = ( u = vec(interior(exchange_atmosphere_state.u)), v = vec(interior(exchange_atmosphere_state.v)), @@ -42,11 +49,9 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha ) end - regularizing_function = ClosedRing() ∘ CutAtAntimeridianAndPoles() - # Magical incantation from ConservativeRegridding - polys1 = map(regularizing_function, (Polygon([LinearRing(f)]) for f in eachcol(grid_faces))) - polys2 = map(regularizing_function, (Polygon([LinearRing(f)]) for f in eachcol(exchange_faces))) + polys1 = map(ClosedRing() ∘ CutAtAntimeridianAndPoles(), (Polygon([LinearRing(f)]) for f in eachcol(grid_faces))) + polys2 = map(ClosedRing() ∘ CutAtAntimeridianAndPoles(), (Polygon([LinearRing(f)]) for f in eachcol(exchange_faces))) regridder = ConservativeRegridding.intersection_areas(polys1, polys2) @@ -90,13 +95,13 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup Qse = exchanger.cpu_surface_state.Qs Qle = exchanger.cpu_surface_state.Qℓ - regrid!(ua, ue, regridder, exchange_areas) - regrid!(va, ve, regridder, exchange_areas) - regrid!(Ta, Te, regridder, exchange_areas) - regrid!(qa, qe, regridder, exchange_areas) - regrid!(pa, pe, regridder, exchange_areas) - regrid!(Qsa, Qse, regridder, exchange_areas) - regrid!(Qla, Qle, regridder, exchange_areas) + ConservativeRegridding.regrid!(ua, ue, regridder, exchange_areas) + ConservativeRegridding.regrid!(va, ve, regridder, exchange_areas) + ConservativeRegridding.regrid!(Ta, Te, regridder, exchange_areas) + ConservativeRegridding.regrid!(qa, qe, regridder, exchange_areas) + ConservativeRegridding.regrid!(pa, pe, regridder, exchange_areas) + ConservativeRegridding.regrid!(Qsa, Qse, regridder, exchange_areas) + ConservativeRegridding.regrid!(Qla, Qle, regridder, exchange_areas) arch = architecture(exchange_grid) fill_exchange_fields!(arch, exchange_state, exchanger.cpu_surface_state) From c663312e09de97ddd99c797bf9004f2939e363d0 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:03:29 +0200 Subject: [PATCH 13/56] add the regridding --- .../speedy_weather_exchanger.jl | 70 +++--- .../speedy_weather_fluxes.jl | 202 +++++++++--------- 2 files changed, 143 insertions(+), 129 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 71542c40c..91124b91b 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -13,6 +13,12 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: StateExchanger, interpolate_atmosphere_state! +const OCRExt = Base.get_extension(Oceananigans, :OceananigansConservativeRegriddingExt) +const SWGExt = Base.get_extension(SpeedyWeather, :SpeedyWeatherGeoMakieExt) + +get_cell_matrix(grid::SpeedyWeather.SpectralGrid) = SWGExt.get_faces(grid; add_nans=false) +get_cell_matrix(grid::Oceananigans.AbstractGrid) = OCRExt.compute_cell_matrix(grid, Center(), Center(), nothing) + # For the moment the workflow is: # 1. Perform the regridding on the CPU # 2. Eventually copy the regridded fields to the GPU @@ -23,8 +29,8 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha # Figure this out: spectral_grid = atmosphere.model.spectral_grid - grid_faces = get_faces(spectral_grid) - exchange_faces = get_faces(exchange_grid) + atmosphere_cell_matrix = get_cell_matrix(spectral_grid) + exchange_cell_matrix = get_cell_matrix(exchange_grid) arch = architecture(exchange_grid) if arch isa Oceananigans.CPU # In case of a CPU grid, we reuse the already allocated fields @@ -39,26 +45,19 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha ) else # Otherwise we allocate new CPU fields cpu_surface_state = ( - u = zeros(Float32, length(exchange_faces)), - v = zeros(Float32, length(exchange_faces)), - T = zeros(Float32, length(exchange_faces)), - q = zeros(Float32, length(exchange_faces)), - p = zeros(Float32, length(exchange_faces)), - Qs = zeros(Float32, length(exchange_faces)), - Qℓ = zeros(Float32, length(exchange_faces)), + u = zeros(Float32, length(exchange_cell_matrix)), + v = zeros(Float32, length(exchange_cell_matrix)), + T = zeros(Float32, length(exchange_cell_matrix)), + q = zeros(Float32, length(exchange_cell_matrix)), + p = zeros(Float32, length(exchange_cell_matrix)), + Qs = zeros(Float32, length(exchange_cell_matrix)), + Qℓ = zeros(Float32, length(exchange_cell_matrix)), ) end - # Magical incantation from ConservativeRegridding - polys1 = map(ClosedRing() ∘ CutAtAntimeridianAndPoles(), (Polygon([LinearRing(f)]) for f in eachcol(grid_faces))) - polys2 = map(ClosedRing() ∘ CutAtAntimeridianAndPoles(), (Polygon([LinearRing(f)]) for f in eachcol(exchange_faces))) - - regridder = ConservativeRegridding.intersection_areas(polys1, polys2) - - exchange_areas = vec(sum(A, dims=2)) - atmosphere_areas = vec(sum(A, dims=1)) + regridder = ConservativeRegridding.Regridder(exchange_cell_matrix, atmosphere_cell_matrix) - exchanger = (; cpu_surface_state, regridder, exchange_areas, atmosphere_areas) + exchanger = (; cpu_surface_state, regridder) return exchanger end @@ -85,8 +84,6 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down Qla = atmos.diagnostic_variables.physics.surface_longwave_down - exchange_areas = exchanger.exchange_areas - ue = exchanger.cpu_surface_state.u ve = exchanger.cpu_surface_state.v Te = exchanger.cpu_surface_state.T @@ -95,16 +92,33 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup Qse = exchanger.cpu_surface_state.Qs Qle = exchanger.cpu_surface_state.Qℓ - ConservativeRegridding.regrid!(ua, ue, regridder, exchange_areas) - ConservativeRegridding.regrid!(va, ve, regridder, exchange_areas) - ConservativeRegridding.regrid!(Ta, Te, regridder, exchange_areas) - ConservativeRegridding.regrid!(qa, qe, regridder, exchange_areas) - ConservativeRegridding.regrid!(pa, pe, regridder, exchange_areas) - ConservativeRegridding.regrid!(Qsa, Qse, regridder, exchange_areas) - ConservativeRegridding.regrid!(Qla, Qle, regridder, exchange_areas) + ConservativeRegridding.regrid!(ue, regridder, ua) + ConservativeRegridding.regrid!(ve, regridder, va) + ConservativeRegridding.regrid!(Te, regridder, Ta) + ConservativeRegridding.regrid!(qe, regridder, qa) + ConservativeRegridding.regrid!(pe, regridder, pa) + ConservativeRegridding.regrid!(Qse, regridder, Qsa) + ConservativeRegridding.regrid!(Qle, regridder, Qla) arch = architecture(exchange_grid) fill_exchange_fields!(arch, exchange_state, exchanger.cpu_surface_state) return nothing -end \ No newline at end of file +end + +function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) + atmos = coupled_model.atmosphere + + # All the location of these fluxes will change + Qs = similarity_theory_fields.sensible_heat + Mv = similarity_theory_fields.water_vapor + Qu = total_fluxes.ocean.heat.upwelling_radiation + + regridder = coupled_model.interfaces.exchanger.regridder + + ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, transpose(regridder), vec(interior(Qs))) + ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux, transpose(regridder), vec(interior(Mv))) + ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.surface_longwave_up, transpose(regridder), vec(interior(Qu))) + + return nothing +end diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl index 73b77ea4a..ab53591b5 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl @@ -1,101 +1,101 @@ -# This document lays out the functions that must be extended to -# use an atmospheric simulation in ClimaOcean. - -import ClimaOcean.OceanSeaIceModels: - compute_net_atmosphere_fluxes! - -import ClimaOcean.OceanSeaIceModels.InterfaceComputations: - atmosphere_exchanger, - initialize!, - StateExchanger, - interpolate_atmosphere_state! - -using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel -using Oceananigans -using Thermodynamics - -""" - interpolate_atmospheric_state!(surface_atmosphere_state, - interpolated_prescribed_freshwater_flux, - atmos::ClimaAtmosSimulation, - grid, clock) - -Interpolate the atmospheric state in `atmos` to `surface_atmospheric_state`, a -the collection of `Field`s needed to compute turbulent fluxes. -""" -function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, coupled_model) - - # Plan: - # 1. transfer atmos data to GPU (req ability to represent atmos on GPU) - # 2. interpolate from atmos grid to ocean grid - - # RingGrids.interpolate!(vec(view(ua, :, :, 1)), atmos.diagnostic_variables.grid.u_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(va, :, :, 1)), atmos.diagnostic_variables.grid.v_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(Ta, :, :, 1)), atmos.diagnostic_variables.grid.temp_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(qa, :, :, 1)), atmos.diagnostic_variables.grid.humid_grid[:, end], interpolator) - # RingGrids.interpolate!(vec(view(pa, :, :, 1)), exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]), interpolator) - # RingGrids.interpolate!(vec(view(Qs, :, :, 1)), atmos.diagnostic_variables.physics.surface_shortwave_down, interpolator) - # RingGrids.interpolate!(vec(view(Qℓ, :, :, 1)), atmos.diagnostic_variables.physics.surface_longwave_down, interpolator) - - # Get the atmospheric state on the ocean grid - # ua = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.u) - # va = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.v) - # Ta = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.T) - # qa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.q) - # pa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.p) - # Qs = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qs) - # Qℓ = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qℓ) - # Mp = on_architecture(Oceananigans.CPU(), interpolated_prescribed_freshwater_flux) - # ρf = fluxes.freshwater_density - - # interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp - # exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state - - # ue = parent(exchange_atmosphere_state.u) - # ve = parent(exchange_atmosphere_state.v) - # Te = parent(exchange_atmosphere_state.T) - # qe = parent(exchange_atmosphere_state.q) - # pe = parent(exchange_atmosphere_state.p) - - # ue = dropdims(ue, dims=3) - # ve = dropdims(ve, dims=3) - # Te = dropdims(Te, dims=3) - # qe = dropdims(qe, dims=3) - # pe = dropdims(pe, dims=3) - - return nothing -end - -function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid) - cpu_grid = on_architecture(CPU(), exchange_grid) - cpu_surface_state = ( - u = Field{Center, Center, Nothing}(cpu_grid), - v = Field{Center, Center, Nothing}(cpu_grid), - T = Field{Center, Center, Nothing}(cpu_grid), - q = Field{Center, Center, Nothing}(cpu_grid), - p = Field{Center, Center, Nothing}(cpu_grid), - Qs = Field{Center, Center, Nothing}(cpu_grid), - Qℓ = Field{Center, Center, Nothing}(cpu_grid), - ) - - # Figure this out: - spectral_grid = atmosphere.model.spectral_grid - # 1/4 degree? - interpolator = SpeedyWeather.RingGrids.AnvilInterpolator(Float32, - SpeedyWeather.FullClenshawGrid, 90, spectral_grid.npoints) - - arch = exchange_grid.architecture - tmp_grid = LatitudeLongitudeGrid(arch; size=(360, 179, 1), latitude=(-89.5, 89.5), longitude=(0, 360), z = (0, 1)) - londs, latds = SpeedyWeather.RingGrids.get_londlatds(spectral_grid.Grid, spectral_grid.nlat_half) - SpeedyWeather.RingGrids.update_locator!(interpolator, londs, latds) - - exchanger = (; cpu_surface_state) - - return something -end - -initialize!(::StateExchanger, ::SpeedySimulation) = nothing - -function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) - return nothing -end +# # This document lays out the functions that must be extended to +# # use an atmospheric simulation in ClimaOcean. + +# import ClimaOcean.OceanSeaIceModels: +# compute_net_atmosphere_fluxes! + +# import ClimaOcean.OceanSeaIceModels.InterfaceComputations: +# atmosphere_exchanger, +# initialize!, +# StateExchanger, +# interpolate_atmosphere_state! + +# using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel +# using Oceananigans +# using Thermodynamics + +# """ +# interpolate_atmospheric_state!(surface_atmosphere_state, +# interpolated_prescribed_freshwater_flux, +# atmos::ClimaAtmosSimulation, +# grid, clock) + +# Interpolate the atmospheric state in `atmos` to `surface_atmospheric_state`, a +# the collection of `Field`s needed to compute turbulent fluxes. +# """ +# function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, coupled_model) + +# # Plan: +# # 1. transfer atmos data to GPU (req ability to represent atmos on GPU) +# # 2. interpolate from atmos grid to ocean grid + +# # RingGrids.interpolate!(vec(view(ua, :, :, 1)), atmos.diagnostic_variables.grid.u_grid[:, end], interpolator) +# # RingGrids.interpolate!(vec(view(va, :, :, 1)), atmos.diagnostic_variables.grid.v_grid[:, end], interpolator) +# # RingGrids.interpolate!(vec(view(Ta, :, :, 1)), atmos.diagnostic_variables.grid.temp_grid[:, end], interpolator) +# # RingGrids.interpolate!(vec(view(qa, :, :, 1)), atmos.diagnostic_variables.grid.humid_grid[:, end], interpolator) +# # RingGrids.interpolate!(vec(view(pa, :, :, 1)), exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]), interpolator) +# # RingGrids.interpolate!(vec(view(Qs, :, :, 1)), atmos.diagnostic_variables.physics.surface_shortwave_down, interpolator) +# # RingGrids.interpolate!(vec(view(Qℓ, :, :, 1)), atmos.diagnostic_variables.physics.surface_longwave_down, interpolator) + +# # Get the atmospheric state on the ocean grid +# # ua = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.u) +# # va = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.v) +# # Ta = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.T) +# # qa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.q) +# # pa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.p) +# # Qs = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qs) +# # Qℓ = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qℓ) +# # Mp = on_architecture(Oceananigans.CPU(), interpolated_prescribed_freshwater_flux) +# # ρf = fluxes.freshwater_density + +# # interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp +# # exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state + +# # ue = parent(exchange_atmosphere_state.u) +# # ve = parent(exchange_atmosphere_state.v) +# # Te = parent(exchange_atmosphere_state.T) +# # qe = parent(exchange_atmosphere_state.q) +# # pe = parent(exchange_atmosphere_state.p) + +# # ue = dropdims(ue, dims=3) +# # ve = dropdims(ve, dims=3) +# # Te = dropdims(Te, dims=3) +# # qe = dropdims(qe, dims=3) +# # pe = dropdims(pe, dims=3) + +# return nothing +# end + +# function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid) +# cpu_grid = on_architecture(CPU(), exchange_grid) +# cpu_surface_state = ( +# u = Field{Center, Center, Nothing}(cpu_grid), +# v = Field{Center, Center, Nothing}(cpu_grid), +# T = Field{Center, Center, Nothing}(cpu_grid), +# q = Field{Center, Center, Nothing}(cpu_grid), +# p = Field{Center, Center, Nothing}(cpu_grid), +# Qs = Field{Center, Center, Nothing}(cpu_grid), +# Qℓ = Field{Center, Center, Nothing}(cpu_grid), +# ) + +# # Figure this out: +# spectral_grid = atmosphere.model.spectral_grid +# # 1/4 degree? +# interpolator = SpeedyWeather.RingGrids.AnvilInterpolator(Float32, +# SpeedyWeather.FullClenshawGrid, 90, spectral_grid.npoints) + +# arch = exchange_grid.architecture +# tmp_grid = LatitudeLongitudeGrid(arch; size=(360, 179, 1), latitude=(-89.5, 89.5), longitude=(0, 360), z = (0, 1)) +# londs, latds = SpeedyWeather.RingGrids.get_londlatds(spectral_grid.Grid, spectral_grid.nlat_half) +# SpeedyWeather.RingGrids.update_locator!(interpolator, londs, latds) + +# exchanger = (; cpu_surface_state) + +# return something +# end + +# initialize!(::StateExchanger, ::SpeedySimulation) = nothing + +# function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) +# return nothing +# end From e33b6afd0bc489f0205ff15422299b154309b3d2 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:11:46 +0200 Subject: [PATCH 14/56] this should wortk --- ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 91124b91b..28e6633ea 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -13,7 +13,7 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: StateExchanger, interpolate_atmosphere_state! -const OCRExt = Base.get_extension(Oceananigans, :OceananigansConservativeRegriddingExt) +const OCRExt = Base.get_extension(Oceananigans, :OceananigansConservativeRegriddingExt) const SWGExt = Base.get_extension(SpeedyWeather, :SpeedyWeatherGeoMakieExt) get_cell_matrix(grid::SpeedyWeather.SpectralGrid) = SWGExt.get_faces(grid; add_nans=false) From c83fa03060e04f66f2164a316396255b5faf41b5 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:16:12 +0200 Subject: [PATCH 15/56] do not use regrid! --- .../ClimaOceanSpeedyWeatherExt.jl | 1 - .../conservative_regridding.jl | 81 ------------------- src/InitialConditions/InitialConditions.jl | 45 ----------- test/test_speedy_coupling.jl | 36 +-------- 4 files changed, 4 insertions(+), 159 deletions(-) delete mode 100644 ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index e9f8cb730..d372bda3c 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -17,7 +17,6 @@ function clenshaw_latitude_longitude_grid(arch, spectral_grid) return grid end -include("conservative_regridding.jl") include("speedy_atmosphere_simulations.jl") include("speedy_weather_exchanger.jl") # include("speedy_weather_fluxes.jl") diff --git a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl b/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl deleted file mode 100644 index 89b3b8b5e..000000000 --- a/ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl +++ /dev/null @@ -1,81 +0,0 @@ -using ConservativeRegridding - -using Oceananigans -using Oceananigans.Utils -using Oceananigans.Grids: λnode, φnode, on_architecture, AbstractGrid -using GeometryOps: Point2 - -using KernelAbstractions: @kernel, @index - -using SpeedyWeather.RingGrids - -""" - get_faces(grid) - -Returns a list representing all horizontal grid cells in a curvilinear `grid`. -The output is an Array of 5 * M `Point2` elements where `M = Nx * Ny`. Each column lists the vertices -associated with a horizontal cell in clockwise order starting from the southwest (bottom left) corner. -""" -function get_faces(grid::AbstractGrid) - Nx, Ny, _ = size(grid) - FT = eltype(grid) - - cpu_grid = on_architecture(Oceananigans.CPU(), grid) - - sw = fill(Point2{FT}(0, 0), 1, Nx*Ny) - nw = fill(Point2{FT}(0, 0), 1, Nx*Ny) - ne = fill(Point2{FT}(0, 0), 1, Nx*Ny) - se = fill(Point2{FT}(0, 0), 1, Nx*Ny) - - launch!(Oceananigans.CPU(), cpu_grid, :xy, _get_vertices!, sw, nw, ne, se, grid) - - vertices = vcat(sw, nw, ne, se, sw) - - return vertices -end - -# Move to SpeedyWeather? -function get_faces(grid::SpeedyWeather.SpectralGrid) - - Grid = grid.Grid - nlat_half = grid.nlat_half - npoints = RingGrids.get_npoints2D(Grid, nlat_half) - - # vertex east, south, west, north (i.e. clockwise for every grid point) - E, S, W, N = RingGrids.get_vertices(Grid, nlat_half) - - # allocate faces as Point2{Float64} so that no data copy has to be made in Makie - faces = Matrix{NTuple{2, Float64}}(undef, 5, npoints) - - @inbounds for ij in 1:npoints - faces[1, ij] = (E[1, ij], E[2, ij]) # clockwise - faces[2, ij] = (S[1, ij], S[2, ij]) - faces[3, ij] = (W[1, ij], W[2, ij]) - faces[4, ij] = (N[1, ij], N[2, ij]) - faces[5, ij] = (E[1, ij], E[2, ij]) # back to east to close the polygon - end - - return faces -end - -@kernel function _get_vertices!(sw, nw, ne, se, grid) - i, j = @index(Global, NTuple) - - FT = eltype(grid) - Nx = size(grid, 1) - λ⁻⁻ = λnode(i, j, 1, grid, Face(), Face(), nothing) - λ⁺⁻ = λnode(i, j+1, 1, grid, Face(), Face(), nothing) - λ⁻⁺ = λnode(i+1, j, 1, grid, Face(), Face(), nothing) - λ⁺⁺ = λnode(i+1, j+1, 1, grid, Face(), Face(), nothing) - - φ⁻⁻ = φnode(i, j, 1, grid, Face(), Face(), nothing) - φ⁺⁻ = φnode(i, j+1, 1, grid, Face(), Face(), nothing) - φ⁻⁺ = φnode(i+1, j, 1, grid, Face(), Face(), nothing) - φ⁺⁺ = φnode(i+1, j+1, 1, grid, Face(), Face(), nothing) - - sw[i+(j-1)*Nx] = Point2{FT}(λ⁻⁻, φ⁻⁻) - nw[i+(j-1)*Nx] = Point2{FT}(λ⁻⁺, φ⁻⁺) - ne[i+(j-1)*Nx] = Point2{FT}(λ⁺⁺, φ⁺⁺) - se[i+(j-1)*Nx] = Point2{FT}(λ⁺⁻, φ⁺⁻) -end - diff --git a/src/InitialConditions/InitialConditions.jl b/src/InitialConditions/InitialConditions.jl index 097da7119..cf5abbcaa 100644 --- a/src/InitialConditions/InitialConditions.jl +++ b/src/InitialConditions/InitialConditions.jl @@ -24,51 +24,6 @@ using Oceananigans.Grids: cpu_face_constructor_x, cpu_face_constructor_z, topology -# Should we move this to grids?? -construct_grid(::Type{<:RectilinearGrid}, arch, size, extent, topology) = - RectilinearGrid(arch; size, x = extent[1], y = extent[2], z = extent[2], topology) - -construct_grid(::Type{<:LatitudeLongitudeGrid}, arch, size, extent, topology) = - LatitudeLongitudeGrid(arch; size, longitude = extent[1], latitude = extent[2], z = extent[3], topology) - -# Regrid a field in three dimensions -function three_dimensional_regrid!(a, b) - target_grid = a.grid isa ImmersedBoundaryGrid ? a.grid.underlying_grid : a.grid - source_grid = b.grid isa ImmersedBoundaryGrid ? b.grid.underlying_grid : b.grid - - topo = topology(target_grid) - arch = architecture(target_grid) - arch = child_architecture(arch) - - target_y = yt = cpu_face_constructor_y(target_grid) - target_z = zt = cpu_face_constructor_z(target_grid) - - target_size = Nt = size(target_grid) - - source_x = xs = cpu_face_constructor_x(source_grid) - source_y = ys = cpu_face_constructor_y(source_grid) - - source_size = Ns = size(source_grid) - - # Start by regridding in z - @debug "Regridding in z" - zgrid = construct_grid(typeof(target_grid), arch, (Ns[1], Ns[2], Nt[3]), (xs, ys, zt), topo) - field_z = Field(location(b), zgrid) - regrid!(field_z, zgrid, source_grid, b) - - # regrid in y - @debug "Regridding in y" - ygrid = construct_grid(typeof(target_grid), arch, (Ns[1], Nt[2], Nt[3]), (xs, yt, zt), topo) - field_y = Field(location(b), ygrid); - regrid!(field_y, ygrid, zgrid, field_z); - - # Finally regrid in x - @debug "Regridding in x" - regrid!(a, target_grid, ygrid, field_y) - - return a -end - include("diffuse_tracers.jl") end # module diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index c1f3f08a9..807ec159b 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -8,36 +8,8 @@ ClimaOceanSpeedyWeatherExt = Base.get_extension(ClimaOcean, :ClimaOceanSpeedyWea @test !isnothing(ClimaOceanSpeedyWeatherExt) spectral_grid = SpeedyWeather.SpectralGrid(trunc=31, nlayers=10) -oceananigans_grid = ClimaOceanSpeedyWeatherExt.clenshaw_latitude_longitude_grid(CPU(), spectral_grid) -λd_sw, φd_sw = SpeedyWeather.RingGrids.get_londlatds(spectral_grid.Grid, spectral_grid.nlat_half) -λd_oc = λnodes(oceananigans_grid, Center(), Center(), Center()) -φd_oc = φnodes(oceananigans_grid, Center(), Center(), Center()) +oceananigans_grid = LatitudeLongitudeGrid(CPU(); size=(200, 100, 1), latitude=(-80, 80), longitude=(0, 360), z = (0, 1)) -@test λd_sw ≈ λd_oc -@test φd_sw ≈ φd_oc - -# SpeedyWeather.RingGrids.update_locator!(interpolator, londs, latds) - -#= - -arch = CPU() -Nx = 180 -Ny = 85 -Nz = 10 -longitude = (0, 360) -latitude = (-85, 85) -halo = (7, 7, 7) -ocean_grid = LatitudeLongitudeGrid(arch; size=(Nx, Ny, Nz), halo, longitude, latitude, z=(-100, 0)) -ocean = ClimaOcean.ocean_simulation(ocean_grid) - -atmos_grid = SpeedyWeather.SpectralGrid(trunc=31, nlayers=10) -atmosphere = ClimaOcean.atmosphere_simulation(atmos_grid) -SpeedyWeather.set!(atmosphere.model.time_stepping, Δt=Minute(10)) - -# EarthSystemModel -# const OceanSeaIceModel = -# const SeaIceOnlyModel = -# const LandOnlyModel = - -model = ClimaOcean.OceanSeaIceModel(ocean; atmosphere) -=# \ No newline at end of file +ocean = ocean_simulation(oceananigans_grid; momentum_advection=nothing, tracer_advection=nothing, closure=nothing) +atmos = atmosphere_simulation(spectral_grid) +earth = OceanSeaIceModel(ocean; atmosphere=atmos) \ No newline at end of file From 292c3c8e11f071aa45771c7ae16e578853be8cd1 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:16:58 +0200 Subject: [PATCH 16/56] remove the regrid! --- src/DataWrangling/ECCO/ECCO.jl | 1 - src/InitialConditions/InitialConditions.jl | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO.jl b/src/DataWrangling/ECCO/ECCO.jl index 2e15720aa..0be5a5e87 100644 --- a/src/DataWrangling/ECCO/ECCO.jl +++ b/src/DataWrangling/ECCO/ECCO.jl @@ -7,7 +7,6 @@ export ECCOFieldTimeSeries, ECCORestoring, LinearlyTaperedPolarMask using ClimaOcean using ClimaOcean.DataWrangling using ClimaOcean.DataWrangling: inpaint_mask!, NearestNeighborInpainting, download_progress, compute_native_date_range -using ClimaOcean.InitialConditions: three_dimensional_regrid!, interpolate! using Oceananigans using Oceananigans: location diff --git a/src/InitialConditions/InitialConditions.jl b/src/InitialConditions/InitialConditions.jl index cf5abbcaa..6ed65f4c2 100644 --- a/src/InitialConditions/InitialConditions.jl +++ b/src/InitialConditions/InitialConditions.jl @@ -18,7 +18,7 @@ using JLD2 # Implementation of 3-dimensional regridding # TODO: move all the following to Oceananigans! -using Oceananigans.Fields: regrid!, interpolate! +using Oceananigans.Fields: interpolate! using Oceananigans.Grids: cpu_face_constructor_x, cpu_face_constructor_y, cpu_face_constructor_z, From a81ab5f25208defccaac6f0a6286eb267b2ace51 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:35:12 +0200 Subject: [PATCH 17/56] this works? --- .../ClimaOceanSpeedyWeatherExt.jl | 41 +++++++++++++++++++ .../speedy_weather_exchanger.jl | 28 ++++++------- .../component_interfaces.jl | 2 +- test/test_speedy_coupling.jl | 10 ++--- 4 files changed, 61 insertions(+), 20 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index d372bda3c..c20d38e58 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -17,6 +17,47 @@ function clenshaw_latitude_longitude_grid(arch, spectral_grid) return grid end +# Put it here for the moment, +# but use the one from Speecy before merging +function get_faces( + Grid::Type{<:SpeedyWeather.AbstractGridArray}, + nlat_half::Integer; + add_nan::Bool = false, +) + npoints = SpeedyWeather.RingGrids.get_npoints2D(Grid, nlat_half) + + # vertex east, south, west, north (i.e. clockwise for every grid point) + E, S, W, N = SpeedyWeather.RingGrids.get_vertices(Grid, nlat_half) + + @boundscheck size(N) == size(W) == size(S) == size(E) || throw(BoundsError("Vertices must have the same size")) + @boundscheck size(N) == (2, npoints) || throw(BoundsError("Number of vertices and npoints do not agree")) + + # number of vertices = 4, 5 to close the polygon, 6 to add a nan + # to prevent grid lines to be drawn between cells + nvertices = add_nan ? 6 : 5 + + # allocate faces as Point2{Float64} so that no data copy has to be made in Makie + faces = Matrix{NTuple{2, Float64}}(undef, nvertices, npoints) + + @inbounds for ij in 1:npoints + faces[1, ij] = (E[1, ij], E[2, ij]) # clockwise + faces[2, ij] = (S[1, ij], S[2, ij]) + faces[3, ij] = (W[1, ij], W[2, ij]) + faces[4, ij] = (N[1, ij], N[2, ij]) + faces[5, ij] = (E[1, ij], E[2, ij]) # back to east to close the polygon + end + + if add_nan # add a NaN to separate grid cells + for ij in 1:npoints + faces[6, ij] = (NaN, NaN) + end + end + + return faces +end + +get_faces(grid::SpeedyWeather.AbstractGridArray; kwargs...) = get_faces(typeof(grid), grid.nlat_half; kwargs...) + include("speedy_atmosphere_simulations.jl") include("speedy_weather_exchanger.jl") # include("speedy_weather_fluxes.jl") diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 28e6633ea..fad031c9a 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -13,11 +13,11 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: StateExchanger, interpolate_atmosphere_state! -const OCRExt = Base.get_extension(Oceananigans, :OceananigansConservativeRegriddingExt) +const OCRExt = Base.get_extension(Oceananigans, :OceananigansConservativeRegriddingExt) const SWGExt = Base.get_extension(SpeedyWeather, :SpeedyWeatherGeoMakieExt) -get_cell_matrix(grid::SpeedyWeather.SpectralGrid) = SWGExt.get_faces(grid; add_nans=false) -get_cell_matrix(grid::Oceananigans.AbstractGrid) = OCRExt.compute_cell_matrix(grid, Center(), Center(), nothing) +get_cell_matrix(grid::SpeedyWeather.SpectralGrid) = get_faces(grid.Grid, grid.nlat_half; add_nan=false) +get_cell_matrix(grid::Oceananigans.AbstractGrid) = OCRExt.compute_cell_matrix(grid, (Center, Center, Nothing)) # For the moment the workflow is: # 1. Perform the regridding on the CPU @@ -71,10 +71,10 @@ end # Regrid the atmospheric state on the exchange grid function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coupled_model) - exchanger = interfaces.exchanger - regridder = exchanger.regridder + atmosphere_exchanger = interfaces.exchanger.atmosphere_exchanger + regridder = atmosphere_exchanger.regridder exchange_grid = interfaces.exchanger.exchange_grid - exchange_state = exchanger.exchange_atmosphere_state + exchange_state = interfaces.exchanger.exchange_atmosphere_state ua = atmos.diagnostic_variables.grid.u_grid[:, end] va = atmos.diagnostic_variables.grid.v_grid[:, end] @@ -84,13 +84,13 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down Qla = atmos.diagnostic_variables.physics.surface_longwave_down - ue = exchanger.cpu_surface_state.u - ve = exchanger.cpu_surface_state.v - Te = exchanger.cpu_surface_state.T - qe = exchanger.cpu_surface_state.q - pe = exchanger.cpu_surface_state.p - Qse = exchanger.cpu_surface_state.Qs - Qle = exchanger.cpu_surface_state.Qℓ + ue = atmosphere_exchanger.cpu_surface_state.u + ve = atmosphere_exchanger.cpu_surface_state.v + Te = atmosphere_exchanger.cpu_surface_state.T + qe = atmosphere_exchanger.cpu_surface_state.q + pe = atmosphere_exchanger.cpu_surface_state.p + Qse = atmosphere_exchanger.cpu_surface_state.Qs + Qle = atmosphere_exchanger.cpu_surface_state.Qℓ ConservativeRegridding.regrid!(ue, regridder, ua) ConservativeRegridding.regrid!(ve, regridder, va) @@ -114,7 +114,7 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) Mv = similarity_theory_fields.water_vapor Qu = total_fluxes.ocean.heat.upwelling_radiation - regridder = coupled_model.interfaces.exchanger.regridder + regridder = coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, transpose(regridder), vec(interior(Qs))) ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux, transpose(regridder), vec(interior(Mv))) diff --git a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl index 4aeed5522..2919bcfbc 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl @@ -123,7 +123,7 @@ end initialize!(exchanger::StateExchanger, ::Nothing) = nothing -function initialize!(exchanger::StateExchanger, atmosphere) +function initialize!(exchanger::StateExchanger, atmosphere::PrescribedAtmosphere) atmos_grid = atmosphere.grid exchange_grid = exchanger.exchange_grid arch = architecture(exchange_grid) diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index 807ec159b..2fbd81f0c 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -1,5 +1,5 @@ -import SpeedyWeather -import ClimaOcean +using SpeedyWeather +using ClimaOcean using Oceananigans using Dates using Test @@ -8,8 +8,8 @@ ClimaOceanSpeedyWeatherExt = Base.get_extension(ClimaOcean, :ClimaOceanSpeedyWea @test !isnothing(ClimaOceanSpeedyWeatherExt) spectral_grid = SpeedyWeather.SpectralGrid(trunc=31, nlayers=10) -oceananigans_grid = LatitudeLongitudeGrid(CPU(); size=(200, 100, 1), latitude=(-80, 80), longitude=(0, 360), z = (0, 1)) +oceananigans_grid = LatitudeLongitudeGrid(Oceananigans.CPU(); size=(200, 100, 1), latitude=(-80, 80), longitude=(0, 360), z = (0, 1)) -ocean = ocean_simulation(oceananigans_grid; momentum_advection=nothing, tracer_advection=nothing, closure=nothing) -atmos = atmosphere_simulation(spectral_grid) +ocean = ClimaOcean.OceanSimulations.ocean_simulation(oceananigans_grid; momentum_advection=nothing, tracer_advection=nothing, closure=nothing) +atmos = ClimaOcean.atmosphere_simulation(spectral_grid) earth = OceanSeaIceModel(ocean; atmosphere=atmos) \ No newline at end of file From 24c88facfec16f97ede787e00b620280756e7235 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:37:45 +0200 Subject: [PATCH 18/56] going --- .../speedy_weather_exchanger.jl | 18 ++++++++++-------- test/test_speedy_coupling.jl | 3 ++- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index fad031c9a..a3d2c6663 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -101,24 +101,26 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup ConservativeRegridding.regrid!(Qle, regridder, Qla) arch = architecture(exchange_grid) - fill_exchange_fields!(arch, exchange_state, exchanger.cpu_surface_state) + fill_exchange_fields!(arch, exchange_state, atmosphere_exchanger.cpu_surface_state) return nothing end +# TODO: For the moment this is just ciupling between ocean and atmosphere. +# we will also need to add the coupling with the sea-ice model function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) atmos = coupled_model.atmosphere # All the location of these fluxes will change - Qs = similarity_theory_fields.sensible_heat - Mv = similarity_theory_fields.water_vapor - Qu = total_fluxes.ocean.heat.upwelling_radiation + Qs = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat + Mv = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.water_vapor - regridder = coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder + regridder = transpose(coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder) - ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, transpose(regridder), vec(interior(Qs))) - ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux, transpose(regridder), vec(interior(Mv))) - ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.surface_longwave_up, transpose(regridder), vec(interior(Qu))) + ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, regridder, vec(interior(Qs))) + ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux, regridder, vec(interior(Mv))) + + # TODO: Figure out how we are going to deal with upwelling radiation return nothing end diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index 2fbd81f0c..312766df7 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -12,4 +12,5 @@ oceananigans_grid = LatitudeLongitudeGrid(Oceananigans.CPU(); size=(200, 100, 1) ocean = ClimaOcean.OceanSimulations.ocean_simulation(oceananigans_grid; momentum_advection=nothing, tracer_advection=nothing, closure=nothing) atmos = ClimaOcean.atmosphere_simulation(spectral_grid) -earth = OceanSeaIceModel(ocean; atmosphere=atmos) \ No newline at end of file +earth = OceanSeaIceModel(ocean; atmosphere=atmos) + From 2152712d2c15819dd9523c8d185a47c91fa3743f Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:41:20 +0200 Subject: [PATCH 19/56] this seems to work --- ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl | 3 +-- test/test_speedy_coupling.jl | 6 ++++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index a3d2c6663..aa6745e25 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -117,10 +117,9 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) regridder = transpose(coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder) + # TODO: Figure out how we are going to deal with upwelling radiation ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, regridder, vec(interior(Qs))) ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux, regridder, vec(interior(Mv))) - # TODO: Figure out how we are going to deal with upwelling radiation - return nothing end diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index 312766df7..f5021d979 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -14,3 +14,9 @@ ocean = ClimaOcean.OceanSimulations.ocean_simulation(oceananigans_grid; momentum atmos = ClimaOcean.atmosphere_simulation(spectral_grid) earth = OceanSeaIceModel(ocean; atmosphere=atmos) +fluxes = earth.interfaces.atmosphere_ocean_interface.fluxes + +Oceananigans.set!(fluxes.sensible_heat, (x, y) -> exp(- y^2 / 10^2)) + +# Regridding to speedy +ClimaOcean.OceanSeaIceModel.InterfaceComputations.compute_net_atmosphere_fluxes!(earth) \ No newline at end of file From 269ca7676b8af6196cb1aeb55af7c4af0265d2b5 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 12:49:22 +0200 Subject: [PATCH 20/56] still some NaNs to fix --- test/test_speedy_coupling.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl index f5021d979..5c2f15436 100644 --- a/test/test_speedy_coupling.jl +++ b/test/test_speedy_coupling.jl @@ -7,13 +7,16 @@ using Test ClimaOceanSpeedyWeatherExt = Base.get_extension(ClimaOcean, :ClimaOceanSpeedyWeatherExt) @test !isnothing(ClimaOceanSpeedyWeatherExt) -spectral_grid = SpeedyWeather.SpectralGrid(trunc=31, nlayers=10) +spectral_grid = SpeedyWeather.SpectralGrid(trunc=51, nlayers=3) oceananigans_grid = LatitudeLongitudeGrid(Oceananigans.CPU(); size=(200, 100, 1), latitude=(-80, 80), longitude=(0, 360), z = (0, 1)) ocean = ClimaOcean.OceanSimulations.ocean_simulation(oceananigans_grid; momentum_advection=nothing, tracer_advection=nothing, closure=nothing) atmos = ClimaOcean.atmosphere_simulation(spectral_grid) -earth = OceanSeaIceModel(ocean; atmosphere=atmos) +# We initialize the atmosphere simulation with some speed and some temperature +SpeedyWeather.initialize!(atmos) + +earth = OceanSeaIceModel(ocean; atmosphere=atmos) fluxes = earth.interfaces.atmosphere_ocean_interface.fluxes Oceananigans.set!(fluxes.sensible_heat, (x, y) -> exp(- y^2 / 10^2)) From 876b3a71ac8f4a15337b98ac5312ef53462a9a0a Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 15 Apr 2025 13:25:22 +0200 Subject: [PATCH 21/56] suggested changes --- .../speedy_weather_exchanger.jl | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index aa6745e25..659d14630 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -32,6 +32,7 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha atmosphere_cell_matrix = get_cell_matrix(spectral_grid) exchange_cell_matrix = get_cell_matrix(exchange_grid) arch = architecture(exchange_grid) + FT = eltype(exchange_atmosphere_state.u) if arch isa Oceananigans.CPU # In case of a CPU grid, we reuse the already allocated fields cpu_surface_state = ( @@ -45,13 +46,13 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha ) else # Otherwise we allocate new CPU fields cpu_surface_state = ( - u = zeros(Float32, length(exchange_cell_matrix)), - v = zeros(Float32, length(exchange_cell_matrix)), - T = zeros(Float32, length(exchange_cell_matrix)), - q = zeros(Float32, length(exchange_cell_matrix)), - p = zeros(Float32, length(exchange_cell_matrix)), - Qs = zeros(Float32, length(exchange_cell_matrix)), - Qℓ = zeros(Float32, length(exchange_cell_matrix)), + u = zeros(FT, length(exchange_cell_matrix)), + v = zeros(FT, length(exchange_cell_matrix)), + T = zeros(FT, length(exchange_cell_matrix)), + q = zeros(FT, length(exchange_cell_matrix)), + p = zeros(FT, length(exchange_cell_matrix)), + Qs = zeros(FT, length(exchange_cell_matrix)), + Qℓ = zeros(FT, length(exchange_cell_matrix)), ) end @@ -112,13 +113,13 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) atmos = coupled_model.atmosphere # All the location of these fluxes will change - Qs = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat + Qc = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat Mv = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.water_vapor regridder = transpose(coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder) # TODO: Figure out how we are going to deal with upwelling radiation - ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, regridder, vec(interior(Qs))) + ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, regridder, vec(interior(Qc))) ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux, regridder, vec(interior(Mv))) return nothing From 2dc2b86703e3a24eaa8827d8825368a20c5330c1 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 21 Aug 2025 17:32:43 +0200 Subject: [PATCH 22/56] starting stuff --- CondaPkg.toml | 7 +- Project.toml | 14 +- experiments/coupled_simulation.jl | 88 ++++++++++++ .../ClimaOceanSpeedyWeatherExt.jl | 56 +------- .../bilinear_interpolator.jl | 136 ++++++++++++++++++ .../speedy_atmosphere_simulations.jl | 23 +-- .../speedy_weather_exchanger.jl | 104 ++++++++------ .../speedy_weather_fluxes.jl | 101 ------------- src/ClimaOcean.jl | 2 - .../assemble_net_fluxes.jl | 6 +- .../atmosphere_sea_ice_fluxes.jl | 2 +- src/OceanSimulations/ocean_simulation.jl | 4 +- 12 files changed, 310 insertions(+), 233 deletions(-) create mode 100644 experiments/coupled_simulation.jl create mode 100644 ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl delete mode 100644 ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl diff --git a/CondaPkg.toml b/CondaPkg.toml index dde757656..4d2ec277f 100644 --- a/CondaPkg.toml +++ b/CondaPkg.toml @@ -1,7 +1,10 @@ [pip.deps] -copernicusmarine = ">=2.0.0" xarray = ">=2024.7.0" -numpy = ">=2.0.0" jax = ">=0.6" +copernicusmarine = ">=2.0.0" +numpy = ">=2.0.0" tensorflow = ">=2.17" +[deps] +xesmf = "" +scipy = "==1.16.0" diff --git a/Project.toml b/Project.toml index 03ef7df54..590490735 100644 --- a/Project.toml +++ b/Project.toml @@ -9,15 +9,11 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CFTime = "179af706-886a-5703-950a-314cd64e0468" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4" -ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343" CubicSplines = "9c784101-8907-5a6d-9be6-98f00873c89b" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" -GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f" -GeometryOps = "3251bfac-6a57-4b6d-aa61-ac1fef2975ab" -GeometryOpsCore = "05efe853-fabf-41c8-927e-7063c8b9f013" ImageMorphology = "787d08f9-d448-5407-9aad-5290dd7ab264" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -29,7 +25,6 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Scratch = "6c6a2e73-6563-6170-7368-637461726353" SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" -SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f" @@ -40,11 +35,15 @@ ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" +PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [extensions] ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"] -ClimaOceanReactantExt = "Reactant" ClimaOceanSpeedyWeatherExt = "SpeedyWeather" +ClimaOceanReactantExt = "Reactant" [compat] Adapt = "4" @@ -68,7 +67,6 @@ PythonCall = "0.9" Reactant = "0.2.45" Scratch = "1" SeawaterPolynomials = "0.3.5" -SpeedyWeather = "0.14.0" StaticArrays = "1" Statistics = "1.9" SurfaceFluxes = "0.11, 0.12" @@ -83,4 +81,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "PythonCall", "CondaPkg"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll"] diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl new file mode 100644 index 000000000..5476cbdf8 --- /dev/null +++ b/experiments/coupled_simulation.jl @@ -0,0 +1,88 @@ +##### +##### A one degree coupled Atmosphere - Ocean - Sea Ice model +##### + +using Oceananigans, ClimaSeaIce, SpeedyWeather, ClimaOcean +using Oceananigans, Oceananigans.Units +using Printf, Statistics, Dates + +##### +##### The Ocean!!! +##### + +Nx = 360 # Longitudinal direction +Ny = 180 # Meridional direction +Nz = 60 # Vertical levels + +r_faces = ExponentialCoordinate(Nz, -6000, 0) + +grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(5, 5, 4)) + +# Regridding the bathymetry... +bottom_height = regrid_bathymetry(grid; minimum_depth=15, major_basins=1, interpolation_passes=15) +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) + +# Advection +momentum_advection = WENOVectorInvariant(order=5) +tracer_advection = WENO(order=5) + +# Free Surface +free_surface = SplitExplicitFreeSurface(grid; substeps=70) + +# Parameterizations +catke_closure = RiBasedVerticalDiffusivity() +eddy_closure = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=1e3, κ_symmetric=1e3) +closures = (catke_closure, eddy_closure, VerticalScalarDiffusivity(κ=1e-5, ν=1e-4)) + +# The ocean simulation +ocean = ocean_simulation(grid; + momentum_advection, + tracer_advection, + free_surface, + timestepper = :SplitRungeKutta3, + closure = closures) + +##### +##### The Atmosphere!!! +##### + +spectral_grid = SpectralGrid(trunc=31, nlayers=8, Grid=FullClenshawGrid) +sea_ice = NoSeaIce() + +humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) +humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) +surface_humidity_flux = SurfaceHumidityFlux(ocean=humidity_flux_ocean, land=humidity_flux_land) + +ocean_heat_flux = PrescribedOceanHeatFlux(spectral_grid) +land_heat_flux = SurfaceLandHeatFlux(spectral_grid) +surface_heat_flux = SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) + +atmosphere_model = PrimitiveWetModel(spectral_grid; + surface_heat_flux, + surface_humidity_flux, + sea_ice) + +atmosphere = initialize!(atmosphere_model) + +##### +##### The Sea-Ice!!! +##### + +dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean; solver=ClimaSeaIce.SeaIceDynamics.SplitExplicitSolver(substeps=150)) +sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) + +##### +##### Coupled model +##### + +Δt = atmosphere_model.time_stepping.Δt_sec + +# Remember in the future that reflected radiatio +# is computed independently by speedy so we need to +# communicate albedo in some way if this reflected radiation +# is absorbed by clouds +radiation = Radiation() +earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) +earth = Simulation(earth_model; Δt, stop_time=365days) + +run!(earth) \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl index c20d38e58..ed214dbb2 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -7,59 +7,13 @@ using Statistics import SpeedyWeather import ClimaOcean import Oceananigans - -function clenshaw_latitude_longitude_grid(arch, spectral_grid) - grid = LatitudeLongitudeGrid(arch; - size = (360, 179, 1), - latitude = (-89.5, 89.5), - longitude = (0, 360), - z = (0, 1)) - return grid -end - -# Put it here for the moment, -# but use the one from Speecy before merging -function get_faces( - Grid::Type{<:SpeedyWeather.AbstractGridArray}, - nlat_half::Integer; - add_nan::Bool = false, -) - npoints = SpeedyWeather.RingGrids.get_npoints2D(Grid, nlat_half) - - # vertex east, south, west, north (i.e. clockwise for every grid point) - E, S, W, N = SpeedyWeather.RingGrids.get_vertices(Grid, nlat_half) - - @boundscheck size(N) == size(W) == size(S) == size(E) || throw(BoundsError("Vertices must have the same size")) - @boundscheck size(N) == (2, npoints) || throw(BoundsError("Number of vertices and npoints do not agree")) - - # number of vertices = 4, 5 to close the polygon, 6 to add a nan - # to prevent grid lines to be drawn between cells - nvertices = add_nan ? 6 : 5 - - # allocate faces as Point2{Float64} so that no data copy has to be made in Makie - faces = Matrix{NTuple{2, Float64}}(undef, nvertices, npoints) - - @inbounds for ij in 1:npoints - faces[1, ij] = (E[1, ij], E[2, ij]) # clockwise - faces[2, ij] = (S[1, ij], S[2, ij]) - faces[3, ij] = (W[1, ij], W[2, ij]) - faces[4, ij] = (N[1, ij], N[2, ij]) - faces[5, ij] = (E[1, ij], E[2, ij]) # back to east to close the polygon - end - - if add_nan # add a NaN to separate grid cells - for ij in 1:npoints - faces[6, ij] = (NaN, NaN) - end - end - - return faces -end - -get_faces(grid::SpeedyWeather.AbstractGridArray; kwargs...) = get_faces(typeof(grid), grid.nlat_half; kwargs...) +import SpeedyWeather.RingGrids include("speedy_atmosphere_simulations.jl") + +# TODO: Remove this when we have a proper interface +# for regridding between tripolar and latitude-longitude grids +include("bilinear_interpolator.jl") include("speedy_weather_exchanger.jl") -# include("speedy_weather_fluxes.jl") end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl new file mode 100644 index 000000000..5c096a59a --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -0,0 +1,136 @@ +using PyCall +using Oceananigans.Grids: AbstractGrid +using SparseArrays +using LinearAlgebra + +const Grids = Union{SpeedyWeather.SpectralGrid, AbstractGrid} + +struct BilinearInterpolator{W1, W2} + set1 :: W1 + set2 :: W2 +end + +function BilinearInterpolator(grid1::Grids, grid2::Grids) + W1 = regridder_weights(grid1, grid2; method="bilinear") + W2 = regridder_weights(grid2, grid1; method="bilinear") + return BilinearInterpolator(W1, W2) +end + +regrid!(dst, weights, scr) = LinearAlgebra.mul!(vec(dst), weights, vec(scr)) + +# Import and store as constants for submodules +const numpy = Ref(pyimport("numpy")) +const xesmf = Ref(pyimport("xesmf")) +const xarray = Ref(pyimport("xarray")) + +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 + +coordinate_dataset(grid::SpeedyWeather.SpectralGrid) = coordinate_dataset(grid.grid) + +function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractGrid) + lon, lat = RingGrids.get_londlatds(grid) + return numpy[].array(PyObject(Dict("lat" => lat, "lon" => lon))) +end + +function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractFullGrid) + lon = RingGrids.get_lond(grid) + lat = RingGrids.get_latd(grid) + return numpy[].array(PyObject(Dict("lat" => lat, "lon" => lon))) +end + +function coordinate_dataset(grid::AbstractGrid) + lat = Array(φnodes(grid, Center(), Center(), Center())) + lon = Array(λnodes(grid, Center(), Center(), Center())) + + lat_b = Array(φnodes(grid, Face(), Face(), Center())) + lon_b = Array(λnodes(grid, Face(), Face(), Center())) + + lat, lon = two_dimensionalize(lat, lon) + lat_b, lon_b = two_dimensionalize(lat_b, lon_b) + + lat_np = numpy[].array(lat) + lon_np = numpy[].array(lon) + + lat_b_np = numpy[].array(lat_b) + lon_b_np = numpy[].array(lon_b) + + return structure_coordinate_dataset(lat_np, lon_np, lat_b_np, lon_b_np) +end + +function structure_coordinate_dataset(lat, lon, lat_b, lon_b) + + ds_lat = xarray[].DataArray( + lat', + dims=["y", "x"], + coords= PyObject(Dict( + "lat" => (["y", "x"], lat'), + "lon" => (["y", "x"], lon') + )), + name="latitude" + ) + + ds_lon = xarray[].DataArray( + lon', + dims=["y", "x"], + coords= PyObject(Dict( + "lat" => (["y", "x"], lat'), + "lon" => (["y", "x"], lon') + )), + name="longitude" + ) + + ds_lat_b = xarray[].DataArray( + lat_b', + dims=["y_b", "x_b"], + coords= PyObject(Dict( + "lat_b" => (["y_b", "x_b"], lat_b'), + "lon_b" => (["y_b", "x_b"], lon_b') + )), + ) + + ds_lon_b = xarray[].DataArray( + lon_b', + dims=["y_b", "x_b"], + coords= PyObject(Dict( + "lat_b" => (["y_b", "x_b"], lat_b'), + "lon_b" => (["y_b", "x_b"], lon_b') + )), + ) + + return xarray[].Dataset( + PyObject( + Dict("lat" => ds_lat, + "lon" => ds_lon, + "lat_b" => ds_lat_b, + "lon_b" => ds_lon_b)) + ) +end + +function regridder_weights(dst::Grids, src::Grids; method::String="bilinear") + + src_ds = coordinate_dataset(src) + dst_ds = coordinate_dataset(dst) + + regridder = xesmf[].Regridder(src_ds, dst_ds, method) + + # Move back to Julia + # Convert the regridder weights to a Julia sparse matrix + coords = regridder.weights.data + shape = Tuple(Int.(coords[:shape])) + vals = Float64.(coords[:data]) + coords = coords[:coords] + rows = coords[1,:].+1 + cols = coords[2,:].+1 + + W = sparse(rows, cols, vals, shape[1], shape[2]) + + return W +end diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl index 59ea6173c..016162f2e 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -4,32 +4,13 @@ import ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: boundary_layer_height, surface_layer_height -import ClimaOcean: atmosphere_simulation - const SpeedySimulation = SpeedyWeather.Simulation const SpeedyCoupledModel = ClimaOcean.OceanSeaIceModel{<:Any, <:SpeedySimulation} Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" -# This can be left blank: -Oceananigans.Models.update_model_field_time_series!(::SpeedySimulation, time) = nothing - # Take one time-step Oceananigans.TimeSteppers.time_step!(atmos::SpeedySimulation, Δt) = SpeedyWeather.timestep!(atmos) -function atmosphere_simulation(grid::SpeedyWeather.SpectralGrid) - # orography = zeros(grid.Grid, grid.nlat_half)) - - surface_heat_flux = SpeedyWeather.PrescribedOceanHeatFlux() - surface_evaporation = SpeedyWeather.PrescribedOceanEvaporation() - atmos_model = SpeedyWeather.PrimitiveWetModel(grid; - # orography, - surface_heat_flux, - surface_evaporation) - - atmos_sim = SpeedyWeather.initialize!(atmos_model) - return atmos_sim -end - # The height of near-surface variables used in the turbulent flux solver function surface_layer_height(s::SpeedySimulation) T = s.model.atmosphere.temp_ref @@ -42,11 +23,9 @@ end # It probably should not be here but in the similarity theory type. boundary_layer_height(atmos::SpeedySimulation) = 600 -# Base.eltype(::EarthAtmosphere{FT}) where FT = FT - # This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather thermodynamics_parameters(atmos::SpeedyWeather.Simulation) = - ClimaOcean.OceanSeaIceModels.PrescribedAtmosphereThermodynamicsParameters(Float32) + ClimaOcean.OceanSeaIceModels.AtmosphereThermodynamicsParameters(Float32) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 659d14630..6a81abbeb 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -1,9 +1,11 @@ -using ConservativeRegridding -using GeometryOps: CutAtAntimeridianAndPoles, ClosedRing -using GeoInterface: Polygon, LinearRing using Oceananigans using Oceananigans.Grids: architecture +using Oceananigans.Utils: launch! +using Oceananigans.Operators: intrinsic_vector +# TODO: Implement conservative regridding when ready +# using ConservativeRegridding +# using GeoInterface: Polygon, LinearRing import ClimaOcean.OceanSeaIceModels: compute_net_atmosphere_fluxes! @@ -13,12 +15,6 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations: StateExchanger, interpolate_atmosphere_state! -const OCRExt = Base.get_extension(Oceananigans, :OceananigansConservativeRegriddingExt) -const SWGExt = Base.get_extension(SpeedyWeather, :SpeedyWeatherGeoMakieExt) - -get_cell_matrix(grid::SpeedyWeather.SpectralGrid) = get_faces(grid.Grid, grid.nlat_half; add_nan=false) -get_cell_matrix(grid::Oceananigans.AbstractGrid) = OCRExt.compute_cell_matrix(grid, (Center, Center, Nothing)) - # For the moment the workflow is: # 1. Perform the regridding on the CPU # 2. Eventually copy the regridded fields to the GPU @@ -29,8 +25,6 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha # Figure this out: spectral_grid = atmosphere.model.spectral_grid - atmosphere_cell_matrix = get_cell_matrix(spectral_grid) - exchange_cell_matrix = get_cell_matrix(exchange_grid) arch = architecture(exchange_grid) FT = eltype(exchange_atmosphere_state.u) @@ -42,22 +36,24 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha q = vec(interior(exchange_atmosphere_state.q)), p = vec(interior(exchange_atmosphere_state.p)), Qs = vec(interior(exchange_atmosphere_state.Qs)), - Qℓ = vec(interior(exchange_atmosphere_state.Qℓ)) + Qℓ = vec(interior(exchange_atmosphere_state.Qℓ)), + Mp = vec(interior(exchange_atmosphere_state.Mp)) ) else # Otherwise we allocate new CPU fields cpu_surface_state = ( - u = zeros(FT, length(exchange_cell_matrix)), - v = zeros(FT, length(exchange_cell_matrix)), - T = zeros(FT, length(exchange_cell_matrix)), - q = zeros(FT, length(exchange_cell_matrix)), - p = zeros(FT, length(exchange_cell_matrix)), - Qs = zeros(FT, length(exchange_cell_matrix)), - Qℓ = zeros(FT, length(exchange_cell_matrix)), + u = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), + v = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), + T = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), + q = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), + p = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), + Qs = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), + Qℓ = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), + Mp = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)) ) end - regridder = ConservativeRegridding.Regridder(exchange_cell_matrix, atmosphere_cell_matrix) - + # TODO: Implement a conservative regridder when ready + regridder = BilinearInterpolator(exchange_grid, spectral_grid) exchanger = (; cpu_surface_state, regridder) return exchanger @@ -65,9 +61,17 @@ end fill_exchange_fields!(::Oceananigans.CPU, args...) = nothing -# TODO: add GPU support -function fill_exchange_fields!(::Oceananigans.GPU, exchange_state, cpu_surface_state) - # Can I just copyto! here? +# TODO: improve GPU support +function fill_exchange_fields!(::Oceananigans.GPU, state, cpustate) + set!(state.u, reshape(cpustate.u, size(state.u))) + set!(state.v, reshape(cpustate.v, size(state.v))) + set!(state.T, reshape(cpustate.T, size(state.T))) + set!(state.q, reshape(cpustate.q, size(state.q))) + set!(state.p, reshape(cpustate.p, size(state.p))) + set!(state.Qs, reshape(cpustate.Qs, size(state.Qs))) + set!(state.Qℓ, reshape(cpustate.Qℓ, size(state.Qℓ))) + set!(state.Mp, reshape(cpustate.Mp, size(state.Mp))) + return nothing end # Regrid the atmospheric state on the exchange grid @@ -76,14 +80,16 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup regridder = atmosphere_exchanger.regridder exchange_grid = interfaces.exchanger.exchange_grid exchange_state = interfaces.exchanger.exchange_atmosphere_state + surface_layer = atmos.model.spectral_grid.nlayers - ua = atmos.diagnostic_variables.grid.u_grid[:, end] - va = atmos.diagnostic_variables.grid.v_grid[:, end] - Ta = atmos.diagnostic_variables.grid.temp_grid[:, end] - qa = atmos.diagnostic_variables.grid.humid_grid[:, end] - pa = exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]) + ua = RingGrids.field_view(atmos.diagnostic_variables.grid.u_grid, :, surface_layer) + va = RingGrids.field_view(atmos.diagnostic_variables.grid.v_grid, :, surface_layer) + Ta = RingGrids.field_view(atmos.diagnostic_variables.grid.temp_grid, :, surface_layer) + qa = RingGrids.field_view(atmos.diagnostic_variables.grid.humid_grid, :, surface_layer) + pa = exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]) Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down - Qla = atmos.diagnostic_variables.physics.surface_longwave_down + Qla = atmos.diagnostic_variables.physics.surface_longwave_down + Mpa = atmos.diagnostic_variables.physics.total_precipitation_rate ue = atmosphere_exchanger.cpu_surface_state.u ve = atmosphere_exchanger.cpu_surface_state.v @@ -92,21 +98,36 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup pe = atmosphere_exchanger.cpu_surface_state.p Qse = atmosphere_exchanger.cpu_surface_state.Qs Qle = atmosphere_exchanger.cpu_surface_state.Qℓ - - ConservativeRegridding.regrid!(ue, regridder, ua) - ConservativeRegridding.regrid!(ve, regridder, va) - ConservativeRegridding.regrid!(Te, regridder, Ta) - ConservativeRegridding.regrid!(qe, regridder, qa) - ConservativeRegridding.regrid!(pe, regridder, pa) - ConservativeRegridding.regrid!(Qse, regridder, Qsa) - ConservativeRegridding.regrid!(Qle, regridder, Qla) + Mpe = atmosphere_exchanger.cpu_surface_state.Mp + + regrid!(ue, regridder.set1, ua) + regrid!(ve, regridder.set1, va) + regrid!(Te, regridder.set1, Ta) + regrid!(qe, regridder.set1, qa) + regrid!(pe, regridder.set1, pa) + regrid!(Qse, regridder.set1, Qsa) + regrid!(Qle, regridder.set1, Qla) + regrid!(Mpe, regridder.set1, Mpa) arch = architecture(exchange_grid) fill_exchange_fields!(arch, exchange_state, atmosphere_exchanger.cpu_surface_state) + u = exchange_state.u + v = exchange_state.v + + launch!(arch, exchange_grid, :xy, _rotate_winds!, u, exchange_grid, v) + return nothing end +@kernel function _rotate_winds!(u, grid, v) + i, j = @index(Global, NTuple) + 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: For the moment this is just ciupling between ocean and atmosphere. # we will also need to add the coupling with the sea-ice model function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) @@ -116,11 +137,12 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) Qc = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat Mv = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.water_vapor - regridder = transpose(coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder) + Qca = atmos.prognostic_variables.ocean.sensible_heat_flux + Mva = atmos.prognostic_variables.ocean.surface_humidity_flux # TODO: Figure out how we are going to deal with upwelling radiation - ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.sensible_heat_flux, regridder, vec(interior(Qc))) - ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux, regridder, vec(interior(Mv))) + regrid!(Qca, regridder.set2, interior(Qc)) + regrid!(Mva, regridder.set2, interior(Mv)) return nothing end diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl deleted file mode 100644 index ab53591b5..000000000 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl +++ /dev/null @@ -1,101 +0,0 @@ -# # This document lays out the functions that must be extended to -# # use an atmospheric simulation in ClimaOcean. - -# import ClimaOcean.OceanSeaIceModels: -# compute_net_atmosphere_fluxes! - -# import ClimaOcean.OceanSeaIceModels.InterfaceComputations: -# atmosphere_exchanger, -# initialize!, -# StateExchanger, -# interpolate_atmosphere_state! - -# using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel -# using Oceananigans -# using Thermodynamics - -# """ -# interpolate_atmospheric_state!(surface_atmosphere_state, -# interpolated_prescribed_freshwater_flux, -# atmos::ClimaAtmosSimulation, -# grid, clock) - -# Interpolate the atmospheric state in `atmos` to `surface_atmospheric_state`, a -# the collection of `Field`s needed to compute turbulent fluxes. -# """ -# function interpolate_atmosphere_state!(interfaces, atmosphere::SpeedySimulation, coupled_model) - -# # Plan: -# # 1. transfer atmos data to GPU (req ability to represent atmos on GPU) -# # 2. interpolate from atmos grid to ocean grid - -# # RingGrids.interpolate!(vec(view(ua, :, :, 1)), atmos.diagnostic_variables.grid.u_grid[:, end], interpolator) -# # RingGrids.interpolate!(vec(view(va, :, :, 1)), atmos.diagnostic_variables.grid.v_grid[:, end], interpolator) -# # RingGrids.interpolate!(vec(view(Ta, :, :, 1)), atmos.diagnostic_variables.grid.temp_grid[:, end], interpolator) -# # RingGrids.interpolate!(vec(view(qa, :, :, 1)), atmos.diagnostic_variables.grid.humid_grid[:, end], interpolator) -# # RingGrids.interpolate!(vec(view(pa, :, :, 1)), exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]), interpolator) -# # RingGrids.interpolate!(vec(view(Qs, :, :, 1)), atmos.diagnostic_variables.physics.surface_shortwave_down, interpolator) -# # RingGrids.interpolate!(vec(view(Qℓ, :, :, 1)), atmos.diagnostic_variables.physics.surface_longwave_down, interpolator) - -# # Get the atmospheric state on the ocean grid -# # ua = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.u) -# # va = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.v) -# # Ta = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.T) -# # qa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.q) -# # pa = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.p) -# # Qs = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qs) -# # Qℓ = on_architecture(Oceananigans.CPU(), surface_atmosphere_state.Qℓ) -# # Mp = on_architecture(Oceananigans.CPU(), interpolated_prescribed_freshwater_flux) -# # ρf = fluxes.freshwater_density - -# # interpolator = interfaces.exchanger.atmosphere_exchanger.to_exchange_interp -# # exchange_atmosphere_state = interfaces.exchanger.exchange_atmosphere_state - -# # ue = parent(exchange_atmosphere_state.u) -# # ve = parent(exchange_atmosphere_state.v) -# # Te = parent(exchange_atmosphere_state.T) -# # qe = parent(exchange_atmosphere_state.q) -# # pe = parent(exchange_atmosphere_state.p) - -# # ue = dropdims(ue, dims=3) -# # ve = dropdims(ve, dims=3) -# # Te = dropdims(Te, dims=3) -# # qe = dropdims(qe, dims=3) -# # pe = dropdims(pe, dims=3) - -# return nothing -# end - -# function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid) -# cpu_grid = on_architecture(CPU(), exchange_grid) -# cpu_surface_state = ( -# u = Field{Center, Center, Nothing}(cpu_grid), -# v = Field{Center, Center, Nothing}(cpu_grid), -# T = Field{Center, Center, Nothing}(cpu_grid), -# q = Field{Center, Center, Nothing}(cpu_grid), -# p = Field{Center, Center, Nothing}(cpu_grid), -# Qs = Field{Center, Center, Nothing}(cpu_grid), -# Qℓ = Field{Center, Center, Nothing}(cpu_grid), -# ) - -# # Figure this out: -# spectral_grid = atmosphere.model.spectral_grid -# # 1/4 degree? -# interpolator = SpeedyWeather.RingGrids.AnvilInterpolator(Float32, -# SpeedyWeather.FullClenshawGrid, 90, spectral_grid.npoints) - -# arch = exchange_grid.architecture -# tmp_grid = LatitudeLongitudeGrid(arch; size=(360, 179, 1), latitude=(-89.5, 89.5), longitude=(0, 360), z = (0, 1)) -# londs, latds = SpeedyWeather.RingGrids.get_londlatds(spectral_grid.Grid, spectral_grid.nlat_half) -# SpeedyWeather.RingGrids.update_locator!(interpolator, londs, latds) - -# exchanger = (; cpu_surface_state) - -# return something -# end - -# initialize!(::StateExchanger, ::SpeedySimulation) = nothing - -# function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) -# return nothing -# end diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 2f2afd38b..7f4edb40e 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -77,8 +77,6 @@ end return NamedTuple{names}(vals) end -function atmosphere_simulation end - include("OceanSimulations/OceanSimulations.jl") include("SeaIceSimulations.jl") include("OceanSeaIceModels/OceanSeaIceModels.jl") diff --git a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl index 1532a3be7..e34aaddba 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl @@ -117,8 +117,8 @@ end Qs = downwelling_radiation.Qs[i, j, 1] # Downwelling shortwave radiation Qℓ = downwelling_radiation.Qℓ[i, j, 1] # Downwelling longwave radiation Qc = atmos_ocean_fluxes.sensible_heat[i, j, 1] # sensible or "conductive" heat flux - Qv = atmos_ocean_fluxes.latent_heat[i, j, 1] # latent heat flux - Mv = atmos_ocean_fluxes.water_vapor[i, j, 1] # mass flux of water vapor + Qv = atmos_ocean_fluxes.latent_heat[i, j, 1] # latent heat flux + Mv = atmos_ocean_fluxes.water_vapor[i, j, 1] # mass flux of water vapor end # Compute radiation fluxes (radiation is multiplied by the fraction of ocean, 1 - sea ice concentration) @@ -170,7 +170,7 @@ end Qio = sea_ice_ocean_fluxes.interface_heat[i, j, 1] Jᵀao = ΣQao * ρₒ⁻¹ / cₒ - Jˢao = - Sₒ * ΣFao + Jˢao = - Sₒ * ΣFao # salinity flux is the opposite of a water vapor flux Jᵀio = Qio * ρₒ⁻¹ / cₒ Jˢio = sea_ice_ocean_fluxes.salt[i, j, 1] * ℵᵢ diff --git a/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl index d12ae69b3..25a0fcfee 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl @@ -29,7 +29,7 @@ function compute_atmosphere_sea_ice_fluxes!(coupled_model) Qs = atmosphere_fields.Qs.data, Qℓ = atmosphere_fields.Qℓ.data, Mp = atmosphere_fields.Mp.data, - h_bℓ = atmosphere.boundary_layer_height) + h_bℓ = boundary_layer_height(atmosphere)) flux_formulation = coupled_model.interfaces.atmosphere_sea_ice_interface.flux_formulation interface_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes diff --git a/src/OceanSimulations/ocean_simulation.jl b/src/OceanSimulations/ocean_simulation.jl index f823440fc..4944f0bd5 100644 --- a/src/OceanSimulations/ocean_simulation.jl +++ b/src/OceanSimulations/ocean_simulation.jl @@ -26,8 +26,8 @@ using Statistics: mean @inline v_quadratic_bottom_drag(i, j, grid, c, Φ, μ) = @inbounds - μ * Φ.v[i, j, 1] * spᶜᶠᶜ(i, j, 1, grid, Φ) # Keep a constant linear drag parameter independent on vertical level -@inline u_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, fields) -@inline v_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, fields) +@inline u_immersed_bottom_drag(i, j, k, grid, clock, Φ, μ) = @inbounds - μ * Φ.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, Φ) +@inline v_immersed_bottom_drag(i, j, k, grid, clock, Φ, μ) = @inbounds - μ * Φ.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, Φ) ##### ##### Defaults From 4148fd74efd15036befb1b0dc36e59abc7ce872c Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 21 Aug 2025 17:58:05 +0200 Subject: [PATCH 23/56] simulation is running! --- experiments/coupled_simulation.jl | 11 +++++++++-- .../bilinear_interpolator.jl | 2 +- .../speedy_weather_exchanger.jl | 1 + 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index 5476cbdf8..ef014ed66 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -42,6 +42,9 @@ ocean = ocean_simulation(grid; timestepper = :SplitRungeKutta3, closure = closures) +Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()), + S=Metadatum(:salinity, dataset=ECCO4Monthly())) + ##### ##### The Atmosphere!!! ##### @@ -63,6 +66,7 @@ atmosphere_model = PrimitiveWetModel(spectral_grid; sea_ice) atmosphere = initialize!(atmosphere_model) +initialize!(atmosphere) ##### ##### The Sea-Ice!!! @@ -71,6 +75,9 @@ atmosphere = initialize!(atmosphere_model) dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean; solver=ClimaSeaIce.SeaIceDynamics.SplitExplicitSolver(substeps=150)) sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) +Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Monthly()), + ℵ=Metadatum(:sea_ice_concentration, dataset=ECCO4Monthly())) + ##### ##### Coupled model ##### @@ -83,6 +90,6 @@ sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) # is absorbed by clouds radiation = Radiation() earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) -earth = Simulation(earth_model; Δt, stop_time=365days) +earth = Oceananigans.Simulation(earth_model; Δt, stop_time=20days) -run!(earth) \ No newline at end of file +Oceananigans.run!(earth) \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl index 5c096a59a..3aaafad1c 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -119,7 +119,7 @@ function regridder_weights(dst::Grids, src::Grids; method::String="bilinear") src_ds = coordinate_dataset(src) dst_ds = coordinate_dataset(dst) - regridder = xesmf[].Regridder(src_ds, dst_ds, method) + regridder = xesmf[].Regridder(src_ds, dst_ds, method, periodic=PyObject(true)) # Move back to Julia # Convert the regridder weights to a Julia sparse matrix diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 6a81abbeb..e2d64242a 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -132,6 +132,7 @@ end # we will also need to add the coupling with the sea-ice model function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) atmos = coupled_model.atmosphere + regridder = coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder # All the location of these fluxes will change Qc = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat From 122e27b0b063df608777b6611cd3ac873e79aa00 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 22 Aug 2025 15:28:42 +0200 Subject: [PATCH 24/56] now everything works --- .../bilinear_interpolator.jl | 28 ++++++++++++------- .../speedy_atmosphere_simulations.jl | 5 +--- .../speedy_weather_exchanger.jl | 13 +++++++-- 3 files changed, 30 insertions(+), 16 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl index 3aaafad1c..15a303a4c 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -41,9 +41,17 @@ function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractGrid) end function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractFullGrid) - lon = RingGrids.get_lond(grid) - lat = RingGrids.get_latd(grid) - return numpy[].array(PyObject(Dict("lat" => lat, "lon" => lon))) + 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) + + return structured_coordinate_dataset(lat, lon, lat_b, lon_b) end function coordinate_dataset(grid::AbstractGrid) @@ -56,16 +64,16 @@ function coordinate_dataset(grid::AbstractGrid) lat, lon = two_dimensionalize(lat, lon) lat_b, lon_b = two_dimensionalize(lat_b, lon_b) - lat_np = numpy[].array(lat) - lon_np = numpy[].array(lon) + return structured_coordinate_dataset(lat, lon, lat_b, lon_b) +end - lat_b_np = numpy[].array(lat_b) - lon_b_np = numpy[].array(lon_b) +function structured_coordinate_dataset(lat, lon, lat_b, lon_b) - return structure_coordinate_dataset(lat_np, lon_np, lat_b_np, lon_b_np) -end + lat = numpy[].array(lat) + lon = numpy[].array(lon) -function structure_coordinate_dataset(lat, lon, lat_b, lon_b) + lat_b = numpy[].array(lat_b) + lon_b = numpy[].array(lon_b) ds_lat = xarray[].DataArray( lat', diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl index 016162f2e..3c493b72a 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -24,8 +24,5 @@ end boundary_layer_height(atmos::SpeedySimulation) = 600 # This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather -thermodynamics_parameters(atmos::SpeedyWeather.Simulation) = +thermodynamics_parameters(atmos::SpeedySimulation) = ClimaOcean.OceanSeaIceModels.AtmosphereThermodynamicsParameters(Float32) - - - diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index e2d64242a..8062470b1 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -1,4 +1,5 @@ using Oceananigans +using Oceananigans.BoundaryConditions using Oceananigans.Grids: architecture using Oceananigans.Utils: launch! using Oceananigans.Operators: intrinsic_vector @@ -108,7 +109,7 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup regrid!(Qse, regridder.set1, Qsa) regrid!(Qle, regridder.set1, Qla) regrid!(Mpe, regridder.set1, Mpa) - + arch = architecture(exchange_grid) fill_exchange_fields!(arch, exchange_state, atmosphere_exchanger.cpu_surface_state) @@ -117,6 +118,14 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup launch!(arch, exchange_grid, :xy, _rotate_winds!, u, exchange_grid, v) + fill_halo_regions!((u, v)) + fill_halo_regions!(exchange_state.T) + fill_halo_regions!(exchange_state.q) + fill_halo_regions!(exchange_state.p) + fill_halo_regions!(exchange_state.Qs) + fill_halo_regions!(exchange_state.Qℓ) + fill_halo_regions!(exchange_state.Mp) + return nothing end @@ -128,7 +137,7 @@ end @inbounds v[i, j, kᴺ] = vₑ end -# TODO: For the moment this is just ciupling between ocean and atmosphere. +# TODO: For the moment this is just coupling between ocean and atmosphere. # we will also need to add the coupling with the sea-ice model function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) atmos = coupled_model.atmosphere From c7b2949593406321ade56ce8ed52f5b52fc6df66 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 22 Aug 2025 16:19:54 +0200 Subject: [PATCH 25/56] use functions --- experiments/coupled_simulation.jl | 74 +++++++++++++++---- .../bilinear_interpolator.jl | 31 ++++---- .../speedy_weather_exchanger.jl | 35 ++++++--- 3 files changed, 100 insertions(+), 40 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index ef014ed66..eca221acd 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -15,8 +15,7 @@ Ny = 180 # Meridional direction Nz = 60 # Vertical levels r_faces = ExponentialCoordinate(Nz, -6000, 0) - -grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(5, 5, 4)) +grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(5, 5, 4)) # Regridding the bathymetry... bottom_height = regrid_bathymetry(grid; minimum_depth=15, major_basins=1, interpolation_passes=15) @@ -27,7 +26,7 @@ momentum_advection = WENOVectorInvariant(order=5) tracer_advection = WENO(order=5) # Free Surface -free_surface = SplitExplicitFreeSurface(grid; substeps=70) +free_surface = SplitExplicitFreeSurface(grid; substeps=80) # Parameterizations catke_closure = RiBasedVerticalDiffusivity() @@ -40,7 +39,8 @@ ocean = ocean_simulation(grid; tracer_advection, free_surface, timestepper = :SplitRungeKutta3, - closure = closures) + closure = closures, + radiative_forcing = nothing) Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()), S=Metadatum(:salinity, dataset=ECCO4Monthly())) @@ -50,7 +50,6 @@ Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()) ##### spectral_grid = SpectralGrid(trunc=31, nlayers=8, Grid=FullClenshawGrid) -sea_ice = NoSeaIce() humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) @@ -63,16 +62,35 @@ surface_heat_flux = SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) atmosphere_model = PrimitiveWetModel(spectral_grid; surface_heat_flux, surface_humidity_flux, - sea_ice) + sea_ice=NoSeaIce()) # This is provided by ClimaSeaIce atmosphere = initialize!(atmosphere_model) initialize!(atmosphere) +function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) + progn, diagn, model = SpeedyWeather.unpack(simulation) + + (; time) = progn.clock # current time + + # set the tendencies back to zero for accumulation + fill!(diagn.tendencies, 0, typeof(model)) + + if model.physics + # calculate all parameterizations + SpeedyWeather.parameterization_tendencies!(diagn, progn, time, model) + end + + return nothing +end + +initialize_atmospheric_state!(atmosphere) +atmosphere.feedback.verbose = false + ##### ##### The Sea-Ice!!! ##### -dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean; solver=ClimaSeaIce.SeaIceDynamics.SplitExplicitSolver(substeps=150)) +dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean) sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Monthly()), @@ -82,14 +100,42 @@ Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Mo ##### Coupled model ##### -Δt = atmosphere_model.time_stepping.Δt_sec +Δt = convert(eltype(grid), atmosphere_model.time_stepping.Δt_sec) -# Remember in the future that reflected radiatio -# is computed independently by speedy so we need to -# communicate albedo in some way if this reflected radiation -# is absorbed by clouds -radiation = Radiation() +# Remember in the future that reflected radiation is computed independently by speedy +# so we need to communicate albedo in some way if this reflected radiation is to be +# absorbed by clouds +radiation = Radiation(ocean_emissivity=0, sea_ice_emissivity=0) earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) earth = Oceananigans.Simulation(earth_model; Δt, stop_time=20days) -Oceananigans.run!(earth) \ No newline at end of file +function progress(sim) + sea_ice = sim.model.sea_ice + ocean = sim.model.ocean + hmax = maximum(sea_ice.model.ice_thickness) + ℵmax = maximum(sea_ice.model.ice_concentration) + uimax = maximum(abs, sea_ice.model.velocities.u) + vimax = maximum(abs, sea_ice.model.velocities.v) + uomax = maximum(abs, ocean.model.velocities.u) + vomax = maximum(abs, ocean.model.velocities.v) + + step_time = 1e-9 * (time_ns() - wall_time[]) + + msg1 = @sprintf("time: %s, iteration: %d, Δt: %s, ", prettytime(sim), iteration(sim), prettytime(sim.Δt)) + msg2 = @sprintf("max(h): %.2e m, max(ℵ): %.2e ", hmax, ℵmax) + msg3 = @sprintf("max uᵢ: (%.2f, %.2f) m s⁻¹, ", uimax, vimax) + msg4 = @sprintf("max uₒ: (%.2f, %.2f) m s⁻¹, ", uomax, vomax) + msg5 = @sprintf("wall time: %s \n", prettytime(step_time)) + + @info msg1 * msg2 * msg3 * msg4 * msg5 + + wall_time[] = time_ns() + + return nothing +end + +# And add it as a callback to the simulation. +add_callback!(earth, progress, IterationInterval(10)) + +# Oceananigans.run!(earth) +time_step!(earth) \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl index 15a303a4c..85c61f212 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -19,9 +19,9 @@ end regrid!(dst, weights, scr) = LinearAlgebra.mul!(vec(dst), weights, vec(scr)) # Import and store as constants for submodules -const numpy = Ref(pyimport("numpy")) -const xesmf = Ref(pyimport("xesmf")) -const xarray = Ref(pyimport("xarray")) +get_numpy() = pyimport("numpy") +get_xesmf() = pyimport("xesmf") +get_xarray() = pyimport("xarray") two_dimensionalize(lat::Matrix, lon::Matrix) = lat, lon @@ -36,8 +36,9 @@ end coordinate_dataset(grid::SpeedyWeather.SpectralGrid) = coordinate_dataset(grid.grid) function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractGrid) + numpy = get_numpy() lon, lat = RingGrids.get_londlatds(grid) - return numpy[].array(PyObject(Dict("lat" => lat, "lon" => lon))) + return numpy.array(PyObject(Dict("lat" => lat, "lon" => lon))) end function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractFullGrid) @@ -68,14 +69,16 @@ function coordinate_dataset(grid::AbstractGrid) end function structured_coordinate_dataset(lat, lon, lat_b, lon_b) + numpy = get_numpy() + xarray = get_xarray() - lat = numpy[].array(lat) - lon = numpy[].array(lon) + lat = numpy.array(lat) + lon = numpy.array(lon) - lat_b = numpy[].array(lat_b) - lon_b = numpy[].array(lon_b) + lat_b = numpy.array(lat_b) + lon_b = numpy.array(lon_b) - ds_lat = xarray[].DataArray( + ds_lat = xarray.DataArray( lat', dims=["y", "x"], coords= PyObject(Dict( @@ -85,7 +88,7 @@ function structured_coordinate_dataset(lat, lon, lat_b, lon_b) name="latitude" ) - ds_lon = xarray[].DataArray( + ds_lon = xarray.DataArray( lon', dims=["y", "x"], coords= PyObject(Dict( @@ -95,7 +98,7 @@ function structured_coordinate_dataset(lat, lon, lat_b, lon_b) name="longitude" ) - ds_lat_b = xarray[].DataArray( + ds_lat_b = xarray.DataArray( lat_b', dims=["y_b", "x_b"], coords= PyObject(Dict( @@ -104,7 +107,7 @@ function structured_coordinate_dataset(lat, lon, lat_b, lon_b) )), ) - ds_lon_b = xarray[].DataArray( + ds_lon_b = xarray.DataArray( lon_b', dims=["y_b", "x_b"], coords= PyObject(Dict( @@ -113,7 +116,7 @@ function structured_coordinate_dataset(lat, lon, lat_b, lon_b) )), ) - return xarray[].Dataset( + return xarray.Dataset( PyObject( Dict("lat" => ds_lat, "lon" => ds_lon, @@ -127,7 +130,7 @@ function regridder_weights(dst::Grids, src::Grids; method::String="bilinear") src_ds = coordinate_dataset(src) dst_ds = coordinate_dataset(dst) - regridder = xesmf[].Regridder(src_ds, dst_ds, method, periodic=PyObject(true)) + regridder = get_xesmf().Regridder(src_ds, dst_ds, method, periodic=PyObject(true)) # Move back to Julia # Convert the regridder weights to a Julia sparse matrix diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 8062470b1..96dd16ef7 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -4,6 +4,8 @@ using Oceananigans.Grids: architecture using Oceananigans.Utils: launch! using Oceananigans.Operators: intrinsic_vector +using ClimaOcean.OceanSeaIceModels: sea_ice_concentration + # TODO: Implement conservative regridding when ready # using ConservativeRegridding # using GeoInterface: Polygon, LinearRing @@ -52,7 +54,7 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha Mp = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)) ) end - + # TODO: Implement a conservative regridder when ready regridder = BilinearInterpolator(exchange_grid, spectral_grid) exchanger = (; cpu_surface_state, regridder) @@ -83,11 +85,11 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup exchange_state = interfaces.exchanger.exchange_atmosphere_state surface_layer = atmos.model.spectral_grid.nlayers - ua = RingGrids.field_view(atmos.diagnostic_variables.grid.u_grid, :, surface_layer) - va = RingGrids.field_view(atmos.diagnostic_variables.grid.v_grid, :, surface_layer) - Ta = RingGrids.field_view(atmos.diagnostic_variables.grid.temp_grid, :, surface_layer) + ua = RingGrids.field_view(atmos.diagnostic_variables.grid.u_grid, :, surface_layer) + va = RingGrids.field_view(atmos.diagnostic_variables.grid.v_grid, :, surface_layer) + Ta = RingGrids.field_view(atmos.diagnostic_variables.grid.temp_grid, :, surface_layer) qa = RingGrids.field_view(atmos.diagnostic_variables.grid.humid_grid, :, surface_layer) - pa = exp.(atmos.diagnostic_variables.grid.pres_grid[:, end]) + pa = exp.(atmos.diagnostic_variables.grid.pres_grid) Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down Qla = atmos.diagnostic_variables.physics.surface_longwave_down Mpa = atmos.diagnostic_variables.physics.total_precipitation_rate @@ -137,22 +139,31 @@ end @inbounds v[i, j, kᴺ] = vₑ end -# TODO: For the moment this is just coupling between ocean and atmosphere. -# we will also need to add the coupling with the sea-ice model + +# 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 compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) atmos = coupled_model.atmosphere + grid = coupled_model.interfaces.exchanger.exchange_grid regridder = coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder - # All the location of these fluxes will change - Qc = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat - Mv = coupled_model.interfaces.atmosphere_ocean_interface.fluxes.water_vapor + ao_fluxes = coupled_model.interfaces.atmosphere_ocean_interface.fluxes + ai_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes + + Qco = ao_fluxes.sensible_heat + Qci = ai_fluxes.sensible_heat + Mvo = ao_fluxes.water_vapor + Mvi = ai_fluxes.water_vapor + ℵ = sea_ice_concentration(coupled_model.sea_ice) + # All the location of these fluxes will change Qca = atmos.prognostic_variables.ocean.sensible_heat_flux Mva = atmos.prognostic_variables.ocean.surface_humidity_flux # TODO: Figure out how we are going to deal with upwelling radiation - regrid!(Qca, regridder.set2, interior(Qc)) - regrid!(Mva, regridder.set2, interior(Mv)) + regrid!(Qca, regridder.set2, interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci)) + regrid!(Mva, regridder.set2, interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi)) return nothing end From 8c2328050fc0e75c62f8611c76d84a52a702acdc Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 22 Aug 2025 16:58:44 +0200 Subject: [PATCH 26/56] speed up stuff --- experiments/coupled_simulation.jl | 6 +- .../speedy_weather_exchanger.jl | 77 ++++++++----------- 2 files changed, 37 insertions(+), 46 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index eca221acd..c39ef19ed 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -84,7 +84,7 @@ function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) end initialize_atmospheric_state!(atmosphere) -atmosphere.feedback.verbose = false +atmosphere.model.feedback.verbose = false ##### ##### The Sea-Ice!!! @@ -109,6 +109,8 @@ radiation = Radiation(ocean_emissivity=0, sea_ice_emissivity=0) earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) earth = Oceananigans.Simulation(earth_model; Δt, stop_time=20days) +wall_time = Ref(time_ns()) + function progress(sim) sea_ice = sim.model.sea_ice ocean = sim.model.ocean @@ -129,7 +131,7 @@ function progress(sim) @info msg1 * msg2 * msg3 * msg4 * msg5 - wall_time[] = time_ns() + wall_time[] = time_ns() return nothing end diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 96dd16ef7..e038206e4 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -31,29 +31,8 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha arch = architecture(exchange_grid) FT = eltype(exchange_atmosphere_state.u) - if arch isa Oceananigans.CPU # In case of a CPU grid, we reuse the already allocated fields - cpu_surface_state = ( - u = vec(interior(exchange_atmosphere_state.u)), - v = vec(interior(exchange_atmosphere_state.v)), - T = vec(interior(exchange_atmosphere_state.T)), - q = vec(interior(exchange_atmosphere_state.q)), - p = vec(interior(exchange_atmosphere_state.p)), - Qs = vec(interior(exchange_atmosphere_state.Qs)), - Qℓ = vec(interior(exchange_atmosphere_state.Qℓ)), - Mp = vec(interior(exchange_atmosphere_state.Mp)) - ) - else # Otherwise we allocate new CPU fields - cpu_surface_state = ( - u = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), - v = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), - T = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), - q = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), - p = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), - Qs = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), - Qℓ = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)), - Mp = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)) - ) - end + # sparse matrix multiply requires contiguous memory to speed up the computation + wrk = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)) # TODO: Implement a conservative regridder when ready regridder = BilinearInterpolator(exchange_grid, spectral_grid) @@ -81,7 +60,7 @@ end function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coupled_model) atmosphere_exchanger = interfaces.exchanger.atmosphere_exchanger regridder = atmosphere_exchanger.regridder - exchange_grid = interfaces.exchanger.exchange_grid + exchange_grid = interfaces.exchanger.exchange_grid exchange_state = interfaces.exchanger.exchange_atmosphere_state surface_layer = atmos.model.spectral_grid.nlayers @@ -93,27 +72,33 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down Qla = atmos.diagnostic_variables.physics.surface_longwave_down Mpa = atmos.diagnostic_variables.physics.total_precipitation_rate + wrk = atmosphere_exchanger.cpu_surface_state.wrk - ue = atmosphere_exchanger.cpu_surface_state.u - ve = atmosphere_exchanger.cpu_surface_state.v - Te = atmosphere_exchanger.cpu_surface_state.T - qe = atmosphere_exchanger.cpu_surface_state.q - pe = atmosphere_exchanger.cpu_surface_state.p - Qse = atmosphere_exchanger.cpu_surface_state.Qs - Qle = atmosphere_exchanger.cpu_surface_state.Qℓ - Mpe = atmosphere_exchanger.cpu_surface_state.Mp - - regrid!(ue, regridder.set1, ua) - regrid!(ve, regridder.set1, va) - regrid!(Te, regridder.set1, Ta) - regrid!(qe, regridder.set1, qa) - regrid!(pe, regridder.set1, pa) - regrid!(Qse, regridder.set1, Qsa) - regrid!(Qle, regridder.set1, Qla) - regrid!(Mpe, regridder.set1, Mpa) + regrid!(wrk, regridder.set1, ua) + Ocenananigans.set!(exchange_state.u, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, va) + Ocenananigans.set!(exchange_state.v, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, Ta) + Ocenananigans.set!(exchange_state.T, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, qa) + Ocenananigans.set!(exchange_state.q, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, pa) + Ocenananigans.set!(exchange_state.p, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, Qsa) + Ocenananigans.set!(exchange_state.Qs, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, Qla) + Ocenananigans.set!(exchange_state.Qℓ, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, Mpa) + Ocenananigans.set!(exchange_state.Mp, reshape(wrk, size(exchange_state.u))) arch = architecture(exchange_grid) - fill_exchange_fields!(arch, exchange_state, atmosphere_exchanger.cpu_surface_state) u = exchange_state.u v = exchange_state.v @@ -147,6 +132,7 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) atmos = coupled_model.atmosphere grid = coupled_model.interfaces.exchanger.exchange_grid regridder = coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder + wrk = coupled_model.interfaces.exchanger.atmosphere_exchanger.wrk ao_fluxes = coupled_model.interfaces.atmosphere_ocean_interface.fluxes ai_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes @@ -161,9 +147,12 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) Qca = atmos.prognostic_variables.ocean.sensible_heat_flux Mva = atmos.prognostic_variables.ocean.surface_humidity_flux + # TODO: Figure out how we are going to deal with upwelling radiation - regrid!(Qca, regridder.set2, interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci)) - regrid!(Mva, regridder.set2, interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi)) + wrk .= interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci) + regrid!(Qca, regridder.set2, wrk) + wrk .= interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi) + regrid!(Mva, regridder.set2, wrk) return nothing end From 90ea8bcf4f52525ca453e8d3a9a5fce76d44b468 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 22 Aug 2025 17:47:40 +0200 Subject: [PATCH 27/56] coupled simulation that runs --- experiments/coupled_simulation.jl | 4 +- .../speedy_weather_exchanger.jl | 59 +++++++------------ 2 files changed, 24 insertions(+), 39 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index c39ef19ed..9fe89431e 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -84,7 +84,7 @@ function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) end initialize_atmospheric_state!(atmosphere) -atmosphere.model.feedback.verbose = false +# atmosphere.model.feedback.verbose = false ##### ##### The Sea-Ice!!! @@ -137,7 +137,7 @@ function progress(sim) end # And add it as a callback to the simulation. -add_callback!(earth, progress, IterationInterval(10)) +# add_callback!(earth, progress, IterationInterval(10)) # Oceananigans.run!(earth) time_step!(earth) \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index e038206e4..f50ad6534 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -36,26 +36,11 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha # TODO: Implement a conservative regridder when ready regridder = BilinearInterpolator(exchange_grid, spectral_grid) - exchanger = (; cpu_surface_state, regridder) + exchanger = (; wrk, regridder) return exchanger end -fill_exchange_fields!(::Oceananigans.CPU, args...) = nothing - -# TODO: improve GPU support -function fill_exchange_fields!(::Oceananigans.GPU, state, cpustate) - set!(state.u, reshape(cpustate.u, size(state.u))) - set!(state.v, reshape(cpustate.v, size(state.v))) - set!(state.T, reshape(cpustate.T, size(state.T))) - set!(state.q, reshape(cpustate.q, size(state.q))) - set!(state.p, reshape(cpustate.p, size(state.p))) - set!(state.Qs, reshape(cpustate.Qs, size(state.Qs))) - set!(state.Qℓ, reshape(cpustate.Qℓ, size(state.Qℓ))) - set!(state.Mp, reshape(cpustate.Mp, size(state.Mp))) - return nothing -end - # Regrid the atmospheric state on the exchange grid function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coupled_model) atmosphere_exchanger = interfaces.exchanger.atmosphere_exchanger @@ -64,39 +49,39 @@ function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coup exchange_state = interfaces.exchanger.exchange_atmosphere_state surface_layer = atmos.model.spectral_grid.nlayers - ua = RingGrids.field_view(atmos.diagnostic_variables.grid.u_grid, :, surface_layer) - va = RingGrids.field_view(atmos.diagnostic_variables.grid.v_grid, :, surface_layer) - Ta = RingGrids.field_view(atmos.diagnostic_variables.grid.temp_grid, :, surface_layer) - qa = RingGrids.field_view(atmos.diagnostic_variables.grid.humid_grid, :, surface_layer) - pa = exp.(atmos.diagnostic_variables.grid.pres_grid) - Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down - Qla = atmos.diagnostic_variables.physics.surface_longwave_down - Mpa = atmos.diagnostic_variables.physics.total_precipitation_rate - wrk = atmosphere_exchanger.cpu_surface_state.wrk + ua = RingGrids.field_view(atmos.diagnostic_variables.grid.u_grid, :, surface_layer).data + va = RingGrids.field_view(atmos.diagnostic_variables.grid.v_grid, :, surface_layer).data + Ta = RingGrids.field_view(atmos.diagnostic_variables.grid.temp_grid, :, surface_layer).data + qa = RingGrids.field_view(atmos.diagnostic_variables.grid.humid_grid, :, surface_layer).data + pa = exp.(atmos.diagnostic_variables.grid.pres_grid.data) + Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down.data + Qla = atmos.diagnostic_variables.physics.surface_longwave_down.data + Mpa = atmos.diagnostic_variables.physics.total_precipitation_rate.data + wrk = atmosphere_exchanger.wrk regrid!(wrk, regridder.set1, ua) - Ocenananigans.set!(exchange_state.u, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.u, reshape(wrk, size(exchange_state.u))) regrid!(wrk, regridder.set1, va) - Ocenananigans.set!(exchange_state.v, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.v, reshape(wrk, size(exchange_state.v))) regrid!(wrk, regridder.set1, Ta) - Ocenananigans.set!(exchange_state.T, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.T, reshape(wrk, size(exchange_state.T))) regrid!(wrk, regridder.set1, qa) - Ocenananigans.set!(exchange_state.q, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.q, reshape(wrk, size(exchange_state.q))) regrid!(wrk, regridder.set1, pa) - Ocenananigans.set!(exchange_state.p, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.p, reshape(wrk, size(exchange_state.p))) regrid!(wrk, regridder.set1, Qsa) - Ocenananigans.set!(exchange_state.Qs, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.Qs, reshape(wrk, size(exchange_state.Qs))) regrid!(wrk, regridder.set1, Qla) - Ocenananigans.set!(exchange_state.Qℓ, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.Qℓ, reshape(wrk, size(exchange_state.Qℓ))) regrid!(wrk, regridder.set1, Mpa) - Ocenananigans.set!(exchange_state.Mp, reshape(wrk, size(exchange_state.u))) + Oceananigans.set!(exchange_state.Mp, reshape(wrk, size(exchange_state.Mp))) arch = architecture(exchange_grid) @@ -144,14 +129,14 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) ℵ = sea_ice_concentration(coupled_model.sea_ice) # All the location of these fluxes will change - Qca = atmos.prognostic_variables.ocean.sensible_heat_flux - Mva = atmos.prognostic_variables.ocean.surface_humidity_flux + Qca = atmos.prognostic_variables.ocean.sensible_heat_flux.data + Mva = atmos.prognostic_variables.ocean.surface_humidity_flux.data # TODO: Figure out how we are going to deal with upwelling radiation - wrk .= interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci) + copyto!(wrk, interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci)) regrid!(Qca, regridder.set2, wrk) - wrk .= interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi) + copyto!(wrk, interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi)) regrid!(Mva, regridder.set2, wrk) return nothing From 58727374090e9691e7659cdb0c17aa70d7c2c32d Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 22 Aug 2025 18:01:03 +0200 Subject: [PATCH 28/56] using trunc=63 --- experiments/coupled_simulation.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index 9fe89431e..01b88e046 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -49,7 +49,7 @@ Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()) ##### The Atmosphere!!! ##### -spectral_grid = SpectralGrid(trunc=31, nlayers=8, Grid=FullClenshawGrid) +spectral_grid = SpectralGrid(trunc=63, nlayers=8, Grid=FullClenshawGrid) humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) @@ -62,7 +62,8 @@ surface_heat_flux = SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) atmosphere_model = PrimitiveWetModel(spectral_grid; surface_heat_flux, surface_humidity_flux, - sea_ice=NoSeaIce()) # This is provided by ClimaSeaIce + ocean = nothing, + sea_ice = nothing) # This is provided by ClimaSeaIce atmosphere = initialize!(atmosphere_model) initialize!(atmosphere) @@ -107,7 +108,7 @@ Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Mo # absorbed by clouds radiation = Radiation(ocean_emissivity=0, sea_ice_emissivity=0) earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) -earth = Oceananigans.Simulation(earth_model; Δt, stop_time=20days) +earth = Oceananigans.Simulation(earth_model; Δt, stop_time=60days) wall_time = Ref(time_ns()) @@ -139,5 +140,4 @@ end # And add it as a callback to the simulation. # add_callback!(earth, progress, IterationInterval(10)) -# Oceananigans.run!(earth) -time_step!(earth) \ No newline at end of file +Oceananigans.run!(earth) From e433d9e75acf6af979257f8d83b7ce273da4ad5a Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 25 Aug 2025 10:13:29 +0200 Subject: [PATCH 29/56] small fix --- ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index f50ad6534..3ee85cd8b 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -109,7 +109,6 @@ end @inbounds v[i, j, kᴺ] = vₑ end - # 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. From 3ad33d7734e98b2dc4a04adcdd9930992068b047 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 25 Aug 2025 10:32:14 +0200 Subject: [PATCH 30/56] add the sst --- ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 3ee85cd8b..09385d938 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -130,13 +130,15 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) # All the location of these fluxes will change Qca = atmos.prognostic_variables.ocean.sensible_heat_flux.data Mva = atmos.prognostic_variables.ocean.surface_humidity_flux.data - + sst = atmos.prognostic_variables.ocean.sea_surface_temperature # TODO: Figure out how we are going to deal with upwelling radiation copyto!(wrk, interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci)) regrid!(Qca, regridder.set2, wrk) copyto!(wrk, interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi)) regrid!(Mva, regridder.set2, wrk) + copyto!(wrk, interior(coupled_model.interfaces.atmosphere_ocean_interface.temperature)) + regrid!(sst, regridder.set2, wrk) return nothing end From 0fb3f267ee7071316cbc49a6dabaacd913c252f3 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 25 Aug 2025 10:33:06 +0200 Subject: [PATCH 31/56] no arch --- ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 09385d938..baa8b622b 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -28,7 +28,6 @@ function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, excha # Figure this out: spectral_grid = atmosphere.model.spectral_grid - arch = architecture(exchange_grid) FT = eltype(exchange_atmosphere_state.u) # sparse matrix multiply requires contiguous memory to speed up the computation From 58323201443217fe1e99a4d71cf7e623c1c10695 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 25 Aug 2025 11:25:09 +0200 Subject: [PATCH 32/56] more changes --- ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index baa8b622b..2931bcf39 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -131,12 +131,16 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) Mva = atmos.prognostic_variables.ocean.surface_humidity_flux.data sst = atmos.prognostic_variables.ocean.sea_surface_temperature + To = coupled_model.interfaces.atmosphere_ocean_interface.temperature + Ti = coupled_model.interfaces.atmosphere_sea_ice_interface.temperature + # TODO: Figure out how we are going to deal with upwelling radiation + # TODO: regrid longwave rather than a mixed surface temperature copyto!(wrk, interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci)) regrid!(Qca, regridder.set2, wrk) copyto!(wrk, interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi)) regrid!(Mva, regridder.set2, wrk) - copyto!(wrk, interior(coupled_model.interfaces.atmosphere_ocean_interface.temperature)) + copyto!(wrk, interior(To) .* (1 - ℵ) .+ ℵ .* interior(Ti)) regrid!(sst, regridder.set2, wrk) return nothing From 34d54f38add16d39b31c9fb462b69fe5a9ddf08c Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 25 Aug 2025 11:26:18 +0200 Subject: [PATCH 33/56] make sure it is kelvin --- ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 2931bcf39..5051d21e4 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -140,7 +140,7 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) regrid!(Qca, regridder.set2, wrk) copyto!(wrk, interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi)) regrid!(Mva, regridder.set2, wrk) - copyto!(wrk, interior(To) .* (1 - ℵ) .+ ℵ .* interior(Ti)) + copyto!(wrk, interior(To) .* (1 - ℵ) .+ ℵ .* interior(Ti) .+ 273.15) regrid!(sst, regridder.set2, wrk) return nothing From 5e054256e96ea023255d7f9f5e45704c78ddb218 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 26 Aug 2025 14:13:57 +0200 Subject: [PATCH 34/56] coupled --- experiments/coupled_simulation.jl | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index 01b88e046..d444de76e 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -39,8 +39,7 @@ ocean = ocean_simulation(grid; tracer_advection, free_surface, timestepper = :SplitRungeKutta3, - closure = closures, - radiative_forcing = nothing) + closure = closures) Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()), S=Metadatum(:salinity, dataset=ECCO4Monthly())) @@ -105,8 +104,10 @@ Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Mo # Remember in the future that reflected radiation is computed independently by speedy # so we need to communicate albedo in some way if this reflected radiation is to be -# absorbed by clouds -radiation = Radiation(ocean_emissivity=0, sea_ice_emissivity=0) +# absorbed by clouds. At the moment, speedy does not provide longwave down, until this +# is provided we just estimate a lower emissivity +# *** THIS IS A HACK TO BE REMOVED *** +radiation = Radiation(ocean_emissivity=0.1, sea_ice_emissivity=0.1) earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) earth = Oceananigans.Simulation(earth_model; Δt, stop_time=60days) From 7524c1d0178fb05d4879b31dfc7da61056645056 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Wed, 27 Aug 2025 10:10:42 +0200 Subject: [PATCH 35/56] add this to tests --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 590490735..38012466e 100644 --- a/Project.toml +++ b/Project.toml @@ -81,4 +81,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather"] From 3b1a9770eeceff81f8d2720ac8f63e616c7e57cd Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Wed, 27 Aug 2025 10:11:19 +0200 Subject: [PATCH 36/56] test this out --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 518b187da..65a2481fe 100644 --- a/Project.toml +++ b/Project.toml @@ -81,4 +81,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather", "Reactant", "CondaPkg", "PyCall"] From e68d16b8ca532704fdb5028d9a4e16edd2379301 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Wed, 27 Aug 2025 12:08:44 +0200 Subject: [PATCH 37/56] add this --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 65a2481fe..982319db4 100644 --- a/Project.toml +++ b/Project.toml @@ -42,7 +42,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [extensions] ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"] -ClimaOceanSpeedyWeatherExt = "SpeedyWeather" +ClimaOceanSpeedyWeatherExt = ["SpeedyWeather", "PyCall"] ClimaOceanReactantExt = "Reactant" [compat] From 99cefaa9076bf4ad112a6f0904abac736d308243 Mon Sep 17 00:00:00 2001 From: simone-silvestri Date: Mon, 1 Sep 2025 08:49:17 -0400 Subject: [PATCH 38/56] try high res --- experiments/coupled_simulation.jl | 66 +++++++++++++++++++++++++------ 1 file changed, 53 insertions(+), 13 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index 01b88e046..68049ba3a 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -10,20 +10,20 @@ using Printf, Statistics, Dates ##### The Ocean!!! ##### -Nx = 360 # Longitudinal direction -Ny = 180 # Meridional direction -Nz = 60 # Vertical levels +Nx = 4320 # Longitudinal direction +Ny = 2160 # Meridional direction +Nz = 10 # Vertical levels -r_faces = ExponentialCoordinate(Nz, -6000, 0) -grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(5, 5, 4)) +r_faces = ExponentialCoordinate(Nz, -4000, 0) +grid = TripolarGrid(GPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(7, 7, 6)) # Regridding the bathymetry... bottom_height = regrid_bathymetry(grid; minimum_depth=15, major_basins=1, interpolation_passes=15) grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) # Advection -momentum_advection = WENOVectorInvariant(order=5) -tracer_advection = WENO(order=5) +momentum_advection = WENOVectorInvariant() +tracer_advection = WENO(order=7) # Free Surface free_surface = SplitExplicitFreeSurface(grid; substeps=80) @@ -31,7 +31,7 @@ free_surface = SplitExplicitFreeSurface(grid; substeps=80) # Parameterizations catke_closure = RiBasedVerticalDiffusivity() eddy_closure = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=1e3, κ_symmetric=1e3) -closures = (catke_closure, eddy_closure, VerticalScalarDiffusivity(κ=1e-5, ν=1e-4)) +closures = (catke_closure, VerticalScalarDiffusivity(κ=1e-5, ν=1e-4)) # The ocean simulation ocean = ocean_simulation(grid; @@ -49,7 +49,7 @@ Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()) ##### The Atmosphere!!! ##### -spectral_grid = SpectralGrid(trunc=63, nlayers=8, Grid=FullClenshawGrid) +spectral_grid = SpectralGrid(trunc=255, nlayers=8, Grid=FullClenshawGrid) humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) @@ -87,11 +87,16 @@ end initialize_atmospheric_state!(atmosphere) # atmosphere.model.feedback.verbose = false +atmosphere.model.output.output_dt = Hour(3) +atmosphere.model.output.active = true +add!(atmosphere.model, SpeedyWeather.SurfaceFluxesOutput()...) +add!(atmosphere.model, SpeedyWeather.PrecipitationOutput()...) + ##### ##### The Sea-Ice!!! ##### -dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean) +dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean; solver=ClimaSeaIce.SeaIceDynamics.SplitExplicitSolver(100)) sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Monthly()), @@ -106,9 +111,9 @@ Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Mo # Remember in the future that reflected radiation is computed independently by speedy # so we need to communicate albedo in some way if this reflected radiation is to be # absorbed by clouds -radiation = Radiation(ocean_emissivity=0, sea_ice_emissivity=0) +radiation = Radiation(ocean_emissivity=0.1, sea_ice_emissivity=0.1) earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) -earth = Oceananigans.Simulation(earth_model; Δt, stop_time=60days) +earth = Oceananigans.Simulation(earth_model; Δt, stop_time=360days) wall_time = Ref(time_ns()) @@ -137,7 +142,42 @@ function progress(sim) return nothing end +outputs = merge(ocean.model.velocities, ocean.model.tracers) + +ocean.output_writers[:free_surf] = JLD2OutputWriter(ocean.model, (; η=ocean.model.free_surface.η); + overwrite_exisiting=true, + schedule=TimeInterval(3600 * 3), + filename="ocean_free_surface.jld2") + +ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; + overwrite_exisiting=true, + schedule=TimeInterval(3600 * 3), + filename="ocean_surface_fields.jld2", + indices=(:, :, grid.Nz)) + +sea_ice.output_writers[:fields] = JLD2OutputWriter(sea_ice.model, Oceananigans.fields(sea_ice.model); + overwrite_exisiting=true, + schedule=TimeInterval(3600 * 3), + filename="sea_ice_fields.jld2") + +Qcao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat +Qvao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.latent_heat +τxao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.x_momentum +τyao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.y_momentum +Qcai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.sensible_heat +Qvai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.latent_heat +τxai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.x_momentum +τyai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.y_momentum +Qoi = earth.model.interfaces.net_fluxes.sea_ice_bottom.heat +Soi = earth.model.interfaces.net_fluxes.sea_ice_bottom.salt +fluxes = (; Qcao, Qvao, τxao, τyao, Qcai, Qvai, τxai, τyai, Qoi, Soi) + +earth.output_writers[:fluxes] = JLD2OutputWriter(earth.model, fluxes; + overwrite_exisiting=true, + schedule=TimeInterval(3600 * 3), + filename="intercomponent_fluxes.jld2") + # And add it as a callback to the simulation. -# add_callback!(earth, progress, IterationInterval(10)) +add_callback!(earth, progress, IterationInterval(10)) Oceananigans.run!(earth) From 4281c447447d09c0ee3512bb1c6c3cb624fa3af1 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 1 Sep 2025 15:38:34 +0200 Subject: [PATCH 39/56] add more packages --- Project.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 982319db4..374c6a7ab 100644 --- a/Project.toml +++ b/Project.toml @@ -33,17 +33,17 @@ ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" -PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [extensions] ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"] -ClimaOceanSpeedyWeatherExt = ["SpeedyWeather", "PyCall"] ClimaOceanReactantExt = "Reactant" +ClimaOceanSpeedyWeatherExt = ["SpeedyWeather", "PyCall"] [compat] Adapt = "4" @@ -81,4 +81,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather", "Reactant", "CondaPkg", "PyCall"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather", "Reactant", "CondaPkg", "PyCall", "SparseArrays", "LinearAlgebra"] From 2c38620cfd3c1f338a1c3839f6b936eb0c2e2169 Mon Sep 17 00:00:00 2001 From: simone-silvestri Date: Mon, 1 Sep 2025 09:52:09 -0400 Subject: [PATCH 40/56] thos now should work --- Project.toml | 2 +- experiments/coupled_simulation.jl | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 982319db4..65a2481fe 100644 --- a/Project.toml +++ b/Project.toml @@ -42,7 +42,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [extensions] ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"] -ClimaOceanSpeedyWeatherExt = ["SpeedyWeather", "PyCall"] +ClimaOceanSpeedyWeatherExt = "SpeedyWeather" ClimaOceanReactantExt = "Reactant" [compat] diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index 436d4d6a8..a6c123234 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -10,15 +10,15 @@ using Printf, Statistics, Dates ##### The Ocean!!! ##### -Nx = 4320 # Longitudinal direction -Ny = 2160 # Meridional direction +Nx = 2160 # Longitudinal direction +Ny = 1080 # Meridional direction Nz = 10 # Vertical levels -r_faces = ExponentialCoordinate(Nz, -4000, 0) +r_faces = ExponentialCoordinate(Nz, -3000, 0) grid = TripolarGrid(GPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(7, 7, 6)) # Regridding the bathymetry... -bottom_height = regrid_bathymetry(grid; minimum_depth=15, major_basins=1, interpolation_passes=15) +bottom_height = regrid_bathymetry(grid; major_basins=1, interpolation_passes=15) grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) # Advection From 0eaac46ce153df14c56b2b29bc3b25f81a93cdb0 Mon Sep 17 00:00:00 2001 From: simone-silvestri Date: Mon, 1 Sep 2025 10:14:30 -0400 Subject: [PATCH 41/56] we do not need PyCall --- CondaPkg.toml | 2 ++ Project.toml | 7 +++---- ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl | 3 ++- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/CondaPkg.toml b/CondaPkg.toml index 4d2ec277f..67292270b 100644 --- a/CondaPkg.toml +++ b/CondaPkg.toml @@ -7,4 +7,6 @@ numpy = ">=2.0.0" tensorflow = ">=2.17" [deps] xesmf = "" +xarray = "" +numpy = "" scipy = "==1.16.0" diff --git a/Project.toml b/Project.toml index 8f7eb4d1c..3a1a34f26 100644 --- a/Project.toml +++ b/Project.toml @@ -34,16 +34,15 @@ ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" [extensions] ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"] -ClimaOceanSpeedyWeatherExt = "SpeedyWeather" ClimaOceanReactantExt = "Reactant" +ClimaOceanSpeedyWeatherExt = "SpeedyWeather" [compat] Adapt = "4" @@ -81,4 +80,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather", "Reactant", "CondaPkg", "PyCall", "SparseArrays", "LinearAlgebra"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather", "Reactant", "CondaPkg", "SparseArrays", "LinearAlgebra"] diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl index 85c61f212..899616b80 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -1,4 +1,5 @@ -using PyCall +using PythonCall +using CondaPkg using Oceananigans.Grids: AbstractGrid using SparseArrays using LinearAlgebra From e819c1b86d92fc20eca1003ca4585d1a87d12766 Mon Sep 17 00:00:00 2001 From: simone-silvestri Date: Mon, 1 Sep 2025 10:38:20 -0400 Subject: [PATCH 42/56] make it work with PythonCall --- .../bilinear_interpolator.jl | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl index 899616b80..2e2adbf2d 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -73,6 +73,11 @@ function structured_coordinate_dataset(lat, lon, lat_b, lon_b) numpy = get_numpy() xarray = get_xarray() + lat = Array(lat') + lon = Array(lon') + lat_b = Array(lat_b') + lon_b = Array(lon_b') + lat = numpy.array(lat) lon = numpy.array(lon) @@ -80,40 +85,40 @@ function structured_coordinate_dataset(lat, lon, lat_b, lon_b) lon_b = numpy.array(lon_b) ds_lat = xarray.DataArray( - lat', + lat, dims=["y", "x"], coords= PyObject(Dict( - "lat" => (["y", "x"], lat'), - "lon" => (["y", "x"], lon') + "lat" => (["y", "x"], lat), + "lon" => (["y", "x"], lon) )), name="latitude" ) ds_lon = xarray.DataArray( - lon', + lon, dims=["y", "x"], coords= PyObject(Dict( - "lat" => (["y", "x"], lat'), - "lon" => (["y", "x"], lon') + "lat" => (["y", "x"], lat), + "lon" => (["y", "x"], lon) )), name="longitude" ) ds_lat_b = xarray.DataArray( - lat_b', + lat_b, dims=["y_b", "x_b"], coords= PyObject(Dict( - "lat_b" => (["y_b", "x_b"], lat_b'), - "lon_b" => (["y_b", "x_b"], lon_b') + "lat_b" => (["y_b", "x_b"], lat_b), + "lon_b" => (["y_b", "x_b"], lon_b) )), ) ds_lon_b = xarray.DataArray( - lon_b', + lon_b, dims=["y_b", "x_b"], coords= PyObject(Dict( - "lat_b" => (["y_b", "x_b"], lat_b'), - "lon_b" => (["y_b", "x_b"], lon_b') + "lat_b" => (["y_b", "x_b"], lat_b), + "lon_b" => (["y_b", "x_b"], lon_b) )), ) From c04115fea11245ac584bc702452a9c69791d8695 Mon Sep 17 00:00:00 2001 From: simone-silvestri Date: Mon, 1 Sep 2025 11:21:48 -0400 Subject: [PATCH 43/56] probably better like this? --- CondaPkg.toml | 4 +- Project.toml | 2 +- .../bilinear_interpolator.jl | 52 +++++++++++-------- 3 files changed, 35 insertions(+), 23 deletions(-) diff --git a/CondaPkg.toml b/CondaPkg.toml index 67292270b..c472d036c 100644 --- a/CondaPkg.toml +++ b/CondaPkg.toml @@ -6,7 +6,9 @@ copernicusmarine = ">=2.0.0" numpy = ">=2.0.0" tensorflow = ">=2.17" [deps] -xesmf = "" xarray = "" numpy = "" scipy = "==1.16.0" + + [deps.xesmf] + channel = "conda-forge" diff --git a/Project.toml b/Project.toml index 3a1a34f26..505782508 100644 --- a/Project.toml +++ b/Project.toml @@ -80,4 +80,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "SpeedyWeather", "Reactant", "CondaPkg", "SparseArrays", "LinearAlgebra"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll"] diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl index 2e2adbf2d..19097333e 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -19,10 +19,21 @@ end regrid!(dst, weights, scr) = LinearAlgebra.mul!(vec(dst), weights, vec(scr)) +function install_package(package) + @info "Installing the $package package..." + CondaPkg.add(package; channel = "conda-forge") + return nothing +end + # Import and store as constants for submodules -get_numpy() = pyimport("numpy") -get_xesmf() = pyimport("xesmf") -get_xarray() = pyimport("xarray") +function get_package(package) + try + return pyimport(package) + catch + install_package(package) + return pyimport(package) + end +end two_dimensionalize(lat::Matrix, lon::Matrix) = lat, lon @@ -37,7 +48,7 @@ end coordinate_dataset(grid::SpeedyWeather.SpectralGrid) = coordinate_dataset(grid.grid) function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractGrid) - numpy = get_numpy() + numpy = get_package("numpy") lon, lat = RingGrids.get_londlatds(grid) return numpy.array(PyObject(Dict("lat" => lat, "lon" => lon))) end @@ -70,8 +81,8 @@ function coordinate_dataset(grid::AbstractGrid) end function structured_coordinate_dataset(lat, lon, lat_b, lon_b) - numpy = get_numpy() - xarray = get_xarray() + numpy = get_package("numpy") + xarray = get_package("xarray") lat = Array(lat') lon = Array(lon') @@ -87,47 +98,46 @@ function structured_coordinate_dataset(lat, lon, lat_b, lon_b) ds_lat = xarray.DataArray( lat, dims=["y", "x"], - coords= PyObject(Dict( + coords=Dict( "lat" => (["y", "x"], lat), "lon" => (["y", "x"], lon) - )), + ), name="latitude" ) ds_lon = xarray.DataArray( lon, dims=["y", "x"], - coords= PyObject(Dict( + coords=Dict( "lat" => (["y", "x"], lat), "lon" => (["y", "x"], lon) - )), + ), name="longitude" ) ds_lat_b = xarray.DataArray( lat_b, dims=["y_b", "x_b"], - coords= PyObject(Dict( + coords=Dict( "lat_b" => (["y_b", "x_b"], lat_b), "lon_b" => (["y_b", "x_b"], lon_b) - )), + ), ) ds_lon_b = xarray.DataArray( lon_b, dims=["y_b", "x_b"], - coords= PyObject(Dict( + coords=Dict( "lat_b" => (["y_b", "x_b"], lat_b), "lon_b" => (["y_b", "x_b"], lon_b) - )), + ), ) - return xarray.Dataset( - PyObject( - Dict("lat" => ds_lat, - "lon" => ds_lon, - "lat_b" => ds_lat_b, - "lon_b" => ds_lon_b)) + return xarray.Dataset( + Dict("lat" => ds_lat, + "lon" => ds_lon, + "lat_b" => ds_lat_b, + "lon_b" => ds_lon_b) ) end @@ -136,7 +146,7 @@ function regridder_weights(dst::Grids, src::Grids; method::String="bilinear") src_ds = coordinate_dataset(src) dst_ds = coordinate_dataset(dst) - regridder = get_xesmf().Regridder(src_ds, dst_ds, method, periodic=PyObject(true)) + regridder = get_package("xesmf").Regridder(src_ds, dst_ds, method, periodic=PyObject(true)) # Move back to Julia # Convert the regridder weights to a Julia sparse matrix From dd41c8405cea0749e525e1fa651a5926a394e0cf Mon Sep 17 00:00:00 2001 From: simone-silvestri Date: Mon, 1 Sep 2025 12:03:29 -0400 Subject: [PATCH 44/56] this should work now --- .../bilinear_interpolator.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl index 19097333e..e5d5eaaf9 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -50,7 +50,7 @@ coordinate_dataset(grid::SpeedyWeather.SpectralGrid) = coordinate_dataset(grid.g function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractGrid) numpy = get_package("numpy") lon, lat = RingGrids.get_londlatds(grid) - return numpy.array(PyObject(Dict("lat" => lat, "lon" => lon))) + return numpy.array(Dict("lat" => lat, "lon" => lon)) end function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractFullGrid) @@ -146,14 +146,14 @@ function regridder_weights(dst::Grids, src::Grids; method::String="bilinear") src_ds = coordinate_dataset(src) dst_ds = coordinate_dataset(dst) - regridder = get_package("xesmf").Regridder(src_ds, dst_ds, method, periodic=PyObject(true)) + regridder = get_package("xesmf").Regridder(src_ds, dst_ds, method, periodic=true) # Move back to Julia # Convert the regridder weights to a Julia sparse matrix coords = regridder.weights.data - shape = Tuple(Int.(coords[:shape])) - vals = Float64.(coords[:data]) - coords = coords[:coords] + shape = pyconvert(Tuple{Int, Int}, coords.shape) + vals = pyconvert(Array{Float64}, coords.data) + coords = pyconvert(Array{Float64}, coords.coords) rows = coords[1,:].+1 cols = coords[2,:].+1 From c52000d0645051f22a8ff514ce8f756878500203 Mon Sep 17 00:00:00 2001 From: simone-silvestri Date: Mon, 1 Sep 2025 12:26:43 -0400 Subject: [PATCH 45/56] molmass ratio --- .../similarity_theory_turbulent_fluxes.jl | 4 ++-- src/OceanSeaIceModels/PrescribedAtmospheres.jl | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl index 401e49fb1..58fed9884 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl @@ -14,7 +14,7 @@ using KernelAbstractions.Extras.LoopInfo: @unroll using Statistics: norm import Thermodynamics as AtmosphericThermodynamics -import Thermodynamics.Parameters: molmass_ratio +import Thermodynamics.Parameters: Rv_over_Rd ##### ##### Bulk turbulent fluxes based on similarity theory @@ -259,7 +259,7 @@ L_★ = u_★² / ϰ b_★ . @inline function buoyancy_scale(θ★, q★, ℂ, 𝒬, g) 𝒯ₐ = AtmosphericThermodynamics.virtual_temperature(ℂ, 𝒬) qₐ = AtmosphericThermodynamics.vapor_specific_humidity(ℂ, 𝒬) - ε = AtmosphericThermodynamics.Parameters.molmass_ratio(ℂ) + ε = AtmosphericThermodynamics.Parameters.Rv_over_Rd(ℂ) δ = ε - 1 # typically equal to 0.608 b★ = g / 𝒯ₐ * (θ★ * (1 + δ * qₐ) + δ * 𝒯ₐ * q★) diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index 596e3d011..521631bd8 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -16,7 +16,7 @@ import Thermodynamics.Parameters: gas_constant, # molmass_dryair, # Molar mass of dry air (without moisture) molmass_water, # Molar mass of gaseous water vapor - molmass_ratio, # Ratio of the molar masses of dry air to water vapor + Rv_over_Rd, # Ratio of the molar masses of dry air to water vapor R_v, # Specific gas constant for water vapor R_d, # Specific gas constant for dry air kappa_d, # Ideal gas adiabatic exponent for dry air @@ -102,7 +102,7 @@ const CP{FT} = ConstitutiveParameters{FT} where FT @inline gas_constant(p::CP) = p.gas_constant @inline molmass_dryair(p::CP) = p.dry_air_molar_mass @inline molmass_water(p::CP) = p.water_molar_mass -@inline molmass_ratio(p::CP) = molmass_dryair(p) / molmass_water(p) +@inline Rv_over_Rd(p::CP) = molmass_dryair(p) / molmass_water(p) @inline R_v(p::CP) = gas_constant(p) / molmass_water(p) @inline R_d(p::CP) = gas_constant(p) / molmass_dryair(p) @@ -261,7 +261,7 @@ const ATP = AtmosphereThermodynamicsParameters @inline gas_constant(p::ATP) = gas_constant(p.constitutive) @inline molmass_dryair(p::ATP) = molmass_dryair(p.constitutive) @inline molmass_water(p::ATP) = molmass_water(p.constitutive) -@inline molmass_ratio(p::ATP) = molmass_ratio(p.constitutive) +@inline Rv_over_Rd(p::ATP) = Rv_over_Rd(p.constitutive) @inline LH_v0(p::ATP) = LH_v0(p.phase_transitions) @inline LH_s0(p::ATP) = LH_s0(p.phase_transitions) @inline LH_f0(p::ATP) = LH_f0(p.phase_transitions) From bbc7bcb9b6944ec8c759dd826975570152822e58 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Wed, 3 Sep 2025 14:26:48 +0200 Subject: [PATCH 46/56] add --- CondaPkg.toml | 14 +---- Project.toml | 4 +- experiments/coupled_simulation.jl | 63 ++++++++++--------- .../speedy_weather_exchanger.jl | 2 +- 4 files changed, 36 insertions(+), 47 deletions(-) diff --git a/CondaPkg.toml b/CondaPkg.toml index c472d036c..e214d0360 100644 --- a/CondaPkg.toml +++ b/CondaPkg.toml @@ -1,14 +1,2 @@ - -[pip.deps] -xarray = ">=2024.7.0" -jax = ">=0.6" -copernicusmarine = ">=2.0.0" -numpy = ">=2.0.0" -tensorflow = ">=2.17" [deps] -xarray = "" -numpy = "" -scipy = "==1.16.0" - - [deps.xesmf] - channel = "conda-forge" +xesmf = "0.8.10" diff --git a/Project.toml b/Project.toml index 505782508..9d7501714 100644 --- a/Project.toml +++ b/Project.toml @@ -48,7 +48,7 @@ ClimaOceanSpeedyWeatherExt = "SpeedyWeather" Adapt = "4" CFTime = "0.1, 0.2" CUDA = "4, 5" -ClimaSeaIce = "0.3.5" +ClimaSeaIce = "0.3.6" CondaPkg = "0.2.28" CubicSplines = "0.2" DataDeps = "0.7" @@ -59,7 +59,7 @@ JLD2 = "0.4, 0.5" KernelAbstractions = "0.9" MPI = "0.20" NCDatasets = "0.12, 0.13, 0.14" -Oceananigans = "0.97.6" +Oceananigans = "0.98.0" OffsetArrays = "1.14" PrecompileTools = "1" PythonCall = "0.9" diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index a6c123234..8101a3d70 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -10,20 +10,20 @@ using Printf, Statistics, Dates ##### The Ocean!!! ##### -Nx = 2160 # Longitudinal direction -Ny = 1080 # Meridional direction -Nz = 10 # Vertical levels +Nx = 360 # Longitudinal direction +Ny = 180 # Meridional direction +Nz = 40 # Vertical levels -r_faces = ExponentialCoordinate(Nz, -3000, 0) -grid = TripolarGrid(GPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(7, 7, 6)) +r_faces = ExponentialCoordinate(Nz, -6000, 0) +grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(7, 7, 6)) # Regridding the bathymetry... bottom_height = regrid_bathymetry(grid; major_basins=1, interpolation_passes=15) grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) # Advection -momentum_advection = WENOVectorInvariant() -tracer_advection = WENO(order=7) +momentum_advection = WENOVectorInvariant(order=5) +tracer_advection = WENO(order=5) # Free Surface free_surface = SplitExplicitFreeSurface(grid; substeps=80) @@ -31,7 +31,7 @@ free_surface = SplitExplicitFreeSurface(grid; substeps=80) # Parameterizations catke_closure = RiBasedVerticalDiffusivity() eddy_closure = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=1e3, κ_symmetric=1e3) -closures = (catke_closure, VerticalScalarDiffusivity(κ=1e-5, ν=1e-4)) +closures = (catke_closure, eddy_closure, VerticalScalarDiffusivity(κ=1e-5, ν=1e-4)) # The ocean simulation ocean = ocean_simulation(grid; @@ -48,7 +48,7 @@ Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()) ##### The Atmosphere!!! ##### -spectral_grid = SpectralGrid(trunc=255, nlayers=8, Grid=FullClenshawGrid) +spectral_grid = SpectralGrid(trunc=63, nlayers=8, Grid=FullClenshawGrid) humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) @@ -65,7 +65,7 @@ atmosphere_model = PrimitiveWetModel(spectral_grid; sea_ice = nothing) # This is provided by ClimaSeaIce atmosphere = initialize!(atmosphere_model) -initialize!(atmosphere) +initialize!(atmosphere; output=false) function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) progn, diagn, model = SpeedyWeather.unpack(simulation) @@ -87,7 +87,6 @@ initialize_atmospheric_state!(atmosphere) # atmosphere.model.feedback.verbose = false atmosphere.model.output.output_dt = Hour(3) -atmosphere.model.output.active = true add!(atmosphere.model, SpeedyWeather.SurfaceFluxesOutput()...) add!(atmosphere.model, SpeedyWeather.PrecipitationOutput()...) @@ -112,7 +111,7 @@ Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Mo # absorbed by clouds. At the moment, speedy does not provide longwave down, until this # is provided we just estimate a lower emissivity # *** THIS IS A HACK TO BE REMOVED *** -radiation = Radiation(ocean_emissivity=0.1, sea_ice_emissivity=0.1) +radiation = Radiation(ocean_emissivity=0.0, sea_ice_emissivity=0.0) earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) earth = Oceananigans.Simulation(earth_model; Δt, stop_time=360days) @@ -144,22 +143,24 @@ function progress(sim) end outputs = merge(ocean.model.velocities, ocean.model.tracers) +sea_ice_fields = merge(sea_ice.model.velocities, sea_ice.model.dynamics.auxiliaries.fields, + (; h=sea_ice.model.ice_thickness, ℵ=sea_ice.model.ice_concentration)) -ocean.output_writers[:free_surf] = JLD2OutputWriter(ocean.model, (; η=ocean.model.free_surface.η); - overwrite_exisiting=true, - schedule=TimeInterval(3600 * 3), - filename="ocean_free_surface.jld2") +ocean.output_writers[:free_surf] = JLD2Writer(ocean.model, (; η=ocean.model.free_surface.η); + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="ocean_free_surface.jld2") -ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; - overwrite_exisiting=true, - schedule=TimeInterval(3600 * 3), - filename="ocean_surface_fields.jld2", - indices=(:, :, grid.Nz)) +ocean.output_writers[:surface] = JLD2Writer(ocean.model, outputs; + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="ocean_surface_fields.jld2", + indices=(:, :, grid.Nz)) -sea_ice.output_writers[:fields] = JLD2OutputWriter(sea_ice.model, Oceananigans.fields(sea_ice.model); - overwrite_exisiting=true, - schedule=TimeInterval(3600 * 3), - filename="sea_ice_fields.jld2") +sea_ice.output_writers[:fields] = JLD2Writer(sea_ice.model, sea_ice_fields; + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="sea_ice_fields.jld2") Qcao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat Qvao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.latent_heat @@ -170,15 +171,15 @@ Qvai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.latent_heat τxai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.x_momentum τyai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.y_momentum Qoi = earth.model.interfaces.net_fluxes.sea_ice_bottom.heat -Soi = earth.model.interfaces.net_fluxes.sea_ice_bottom.salt +Soi = earth.model.interfaces.sea_ice_ocean_interface.fluxes.salt fluxes = (; Qcao, Qvao, τxao, τyao, Qcai, Qvai, τxai, τyai, Qoi, Soi) -earth.output_writers[:fluxes] = JLD2OutputWriter(earth.model, fluxes; - overwrite_exisiting=true, - schedule=TimeInterval(3600 * 3), - filename="intercomponent_fluxes.jld2") +earth.output_writers[:fluxes] = JLD2Writer(earth.model.ocean.model, fluxes; + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="intercomponent_fluxes.jld2") # And add it as a callback to the simulation. -add_callback!(earth, progress, IterationInterval(10)) +# add_callback!(earth, progress, IterationInterval(10)) Oceananigans.run!(earth) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl index 5051d21e4..2812354ca 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -129,7 +129,7 @@ function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) # All the location of these fluxes will change Qca = atmos.prognostic_variables.ocean.sensible_heat_flux.data Mva = atmos.prognostic_variables.ocean.surface_humidity_flux.data - sst = atmos.prognostic_variables.ocean.sea_surface_temperature + sst = atmos.prognostic_variables.ocean.sea_surface_temperature.data To = coupled_model.interfaces.atmosphere_ocean_interface.temperature Ti = coupled_model.interfaces.atmosphere_sea_ice_interface.temperature From 0db655fcea2979a905f03491f0965e32bd0dd882 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Wed, 3 Sep 2025 15:09:44 +0200 Subject: [PATCH 47/56] add a coupled simulation --- experiments/coupled_simulation.jl | 42 +---------------- .../speedy_atmosphere_simulations.jl | 45 +++++++++++++++++++ src/ClimaOcean.jl | 2 + 3 files changed, 49 insertions(+), 40 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index 8101a3d70..f9521b3cd 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -49,46 +49,8 @@ Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()) ##### spectral_grid = SpectralGrid(trunc=63, nlayers=8, Grid=FullClenshawGrid) - -humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) -humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) -surface_humidity_flux = SurfaceHumidityFlux(ocean=humidity_flux_ocean, land=humidity_flux_land) - -ocean_heat_flux = PrescribedOceanHeatFlux(spectral_grid) -land_heat_flux = SurfaceLandHeatFlux(spectral_grid) -surface_heat_flux = SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) - -atmosphere_model = PrimitiveWetModel(spectral_grid; - surface_heat_flux, - surface_humidity_flux, - ocean = nothing, - sea_ice = nothing) # This is provided by ClimaSeaIce - -atmosphere = initialize!(atmosphere_model) -initialize!(atmosphere; output=false) - -function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) - progn, diagn, model = SpeedyWeather.unpack(simulation) - - (; time) = progn.clock # current time - - # set the tendencies back to zero for accumulation - fill!(diagn.tendencies, 0, typeof(model)) - - if model.physics - # calculate all parameterizations - SpeedyWeather.parameterization_tendencies!(diagn, progn, time, model) - end - - return nothing -end - -initialize_atmospheric_state!(atmosphere) -# atmosphere.model.feedback.verbose = false - +atmosphere = atmosphere_simulation(spectral_grid; output=true) atmosphere.model.output.output_dt = Hour(3) -add!(atmosphere.model, SpeedyWeather.SurfaceFluxesOutput()...) -add!(atmosphere.model, SpeedyWeather.PrecipitationOutput()...) ##### ##### The Sea-Ice!!! @@ -180,6 +142,6 @@ earth.output_writers[:fluxes] = JLD2Writer(earth.model.ocean.model, fluxes; filename="intercomponent_fluxes.jld2") # And add it as a callback to the simulation. -# add_callback!(earth, progress, IterationInterval(10)) +add_callback!(earth, progress, IterationInterval(100)) Oceananigans.run!(earth) diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl index 3c493b72a..f8f4c70f9 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -1,3 +1,5 @@ +import ClimaOcean: atmosphere_simulation + # Make sure the atmospheric parameters from SpeedyWeather can be used in the compute fluxes function import ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: thermodynamics_parameters, @@ -26,3 +28,46 @@ boundary_layer_height(atmos::SpeedySimulation) = 600 # This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather thermodynamics_parameters(atmos::SpeedySimulation) = ClimaOcean.OceanSeaIceModels.AtmosphereThermodynamicsParameters(Float32) + +function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) + progn, diagn, model = SpeedyWeather.unpack(simulation) + (; time) = progn.clock # current time + + # set the tendencies back to zero for accumulation + fill!(diagn.tendencies, 0, typeof(model)) + + if model.physics + SpeedyWeather.parameterization_tendencies!(diagn, progn, time, model) + end + + return nothing +end + +function atmosphere_simulation(spectral_grid::SpectralGrid; output=false) + # Surface fluxes + humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) + humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) + surface_humidity_flux = SurfaceHumidityFlux(ocean=humidity_flux_ocean, land=humidity_flux_land) + + ocean_heat_flux = PrescribedOceanHeatFlux(spectral_grid) + land_heat_flux = SurfaceLandHeatFlux(spectral_grid) + surface_heat_flux = SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) + + # The atmospheric model + atmosphere_model = PrimitiveWetModel(spectral_grid; + surface_heat_flux, + surface_humidity_flux, + ocean = nothing, + sea_ice = nothing) # This is provided by ClimaSeaIce + + # Construct the simulation + atmosphere = initialize!(atmosphere_model) + + # Initialize the simulation + initialize!(atmosphere; output) + + # Fill in prognostic fields + initialize_atmospheric_state!(atmosphere) + + return atmosphere +end diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 7f4edb40e..2f2afd38b 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -77,6 +77,8 @@ end return NamedTuple{names}(vals) end +function atmosphere_simulation end + include("OceanSimulations/OceanSimulations.jl") include("SeaIceSimulations.jl") include("OceanSeaIceModels/OceanSeaIceModels.jl") From 6391632e63a75675fef8ee42aa2cad39d050e0c5 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Wed, 3 Sep 2025 15:15:37 +0200 Subject: [PATCH 48/56] go ahead --- experiments/coupled_simulation.jl | 3 ++- .../speedy_atmosphere_simulations.jl | 17 ++++++++++++++--- src/ClimaOcean.jl | 1 + .../time_step_ocean_sea_ice_model.jl | 4 ++-- 4 files changed, 19 insertions(+), 6 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index f9521b3cd..5ef2b3cf0 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -66,7 +66,8 @@ Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Mo ##### Coupled model ##### -Δt = convert(eltype(grid), atmosphere_model.time_stepping.Δt_sec) +# The ocean goes twice as slow as the atmosphere! +Δt = 2 * convert(eltype(grid), atmosphere_model.time_stepping.Δt_sec) # Remember in the future that reflected radiation is computed independently by speedy # so we need to communicate albedo in some way if this reflected radiation is to be diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl index f8f4c70f9..c2b0db987 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -10,8 +10,19 @@ const SpeedySimulation = SpeedyWeather.Simulation const SpeedyCoupledModel = ClimaOcean.OceanSeaIceModel{<:Any, <:SpeedySimulation} Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" -# Take one time-step -Oceananigans.TimeSteppers.time_step!(atmos::SpeedySimulation, Δt) = SpeedyWeather.timestep!(atmos) +# Take one time-step or more depending on the global timestep +function Oceananigans.TimeSteppers.time_step!(atmos::SpeedySimulation, Δt) + Δt_atmos = atmos.model.time_stepping.Δt_sec + nsteps = ceil(Int, Δt / Δt_atmos) + + if (Δt / Δt_atmos) % 1 != 0 + @warn "The only atmosphere timesteps supported are intiger divisors of the global timestepping" + end + + for _ in 1:nsteps + SpeedyWeather.timestep!(atmos) + end +end # The height of near-surface variables used in the turbulent flux solver function surface_layer_height(s::SpeedySimulation) @@ -43,7 +54,7 @@ function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) return nothing end -function atmosphere_simulation(spectral_grid::SpectralGrid; output=false) +function atmosphere_simulation(spectral_grid::SpeedyWeather.SpectralGrid; output=false) # Surface fluxes humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 2f2afd38b..bd6da806c 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -41,6 +41,7 @@ export DatasetRestoring, ocean_simulation, sea_ice_simulation, + atmosphere_simulation, initialize! using Oceananigans diff --git a/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl b/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl index 0c8d3cc47..0ea66d4f5 100644 --- a/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl @@ -20,10 +20,10 @@ function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_ # TODO after ice time-step: # - Adjust ocean heat flux if the ice completely melts? - time_step!(ocean, Δt) + !isnothing(sea_ice) && time_step!(ocean, Δt) # Time step the atmosphere - time_step!(atmosphere, Δt) + !isnothing(sea_ice) && time_step!(atmosphere, Δt) # TODO: # - Store fractional ice-free / ice-covered _time_ for more From aebdd57736490ad68226736337913229bf9402ac Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 4 Sep 2025 11:00:45 +0200 Subject: [PATCH 49/56] add this --- experiments/coupled_simulation.jl | 17 +++++++++++ .../speedy_atmosphere_simulations.jl | 28 +++++++++---------- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/experiments/coupled_simulation.jl b/experiments/coupled_simulation.jl index 5ef2b3cf0..66dd61b57 100644 --- a/experiments/coupled_simulation.jl +++ b/experiments/coupled_simulation.jl @@ -146,3 +146,20 @@ earth.output_writers[:fluxes] = JLD2Writer(earth.model.ocean.model, fluxes; add_callback!(earth, progress, IterationInterval(100)) Oceananigans.run!(earth) + +#### +#### Visualize!!! +#### + +using NCDatasets, GLMakie + +SWO = Dataset("run_0002/output.nc") + +Ta = SWO["temp"][:, :, 8, :] +qa = SWO["humid"][:, :, 8, :] +ua = SWO["u"][:, :, 8, :] +va = SWO["v"][:, :, 8, :] +sp = sqrt.(ua.^2 + va.^2) +# Surface fields +SST = FieldTimeSeries("ocean_surface_fields.jld2", "T") +SSS = FieldTimeSeries("ocean_surface_fields.jld2", "S") diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl index c2b0db987..5713db86a 100644 --- a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -16,7 +16,7 @@ function Oceananigans.TimeSteppers.time_step!(atmos::SpeedySimulation, Δt) nsteps = ceil(Int, Δt / Δt_atmos) if (Δt / Δt_atmos) % 1 != 0 - @warn "The only atmosphere timesteps supported are intiger divisors of the global timestepping" + @warn "ClimaOcean only supports atmosphere timesteps that are integer divisors of the ESM timesteps" end for _ in 1:nsteps @@ -56,26 +56,26 @@ end function atmosphere_simulation(spectral_grid::SpeedyWeather.SpectralGrid; output=false) # Surface fluxes - humidity_flux_ocean = PrescribedOceanHumidityFlux(spectral_grid) - humidity_flux_land = SurfaceLandHumidityFlux(spectral_grid) - surface_humidity_flux = SurfaceHumidityFlux(ocean=humidity_flux_ocean, land=humidity_flux_land) + humidity_flux_ocean = SpeedyWeather.PrescribedOceanHumidityFlux(spectral_grid) + humidity_flux_land = SpeedyWeather.SurfaceLandHumidityFlux(spectral_grid) + surface_humidity_flux = SpeedyWeather.SurfaceHumidityFlux(ocean=humidity_flux_ocean, land=humidity_flux_land) - ocean_heat_flux = PrescribedOceanHeatFlux(spectral_grid) - land_heat_flux = SurfaceLandHeatFlux(spectral_grid) - surface_heat_flux = SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) + ocean_heat_flux = SpeedyWeather.PrescribedOceanHeatFlux(spectral_grid) + land_heat_flux = SpeedyWeather.SurfaceLandHeatFlux(spectral_grid) + surface_heat_flux = SpeedyWeather.SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) # The atmospheric model - atmosphere_model = PrimitiveWetModel(spectral_grid; - surface_heat_flux, - surface_humidity_flux, - ocean = nothing, - sea_ice = nothing) # This is provided by ClimaSeaIce + atmosphere_model = SpeedyWeather.PrimitiveWetModel(spectral_grid; + surface_heat_flux, + surface_humidity_flux, + ocean = nothing, + sea_ice = nothing) # This is provided by ClimaSeaIce # Construct the simulation - atmosphere = initialize!(atmosphere_model) + atmosphere = SpeedyWeather.initialize!(atmosphere_model) # Initialize the simulation - initialize!(atmosphere; output) + SpeedyWeather.initialize!(atmosphere; output) # Fill in prognostic fields initialize_atmospheric_state!(atmosphere) From fb09b3734f45f8a4ce7c83d523bece8fef0230d0 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 8 Sep 2025 11:57:41 +0200 Subject: [PATCH 50/56] add it as an example --- .buildkite/examples_build.yml | 2 +- docs/make.jl | 4 +- .../coupled_simulation.jl | 169 +++++++++++++----- 3 files changed, 129 insertions(+), 46 deletions(-) rename {experiments => examples}/coupled_simulation.jl (53%) diff --git a/.buildkite/examples_build.yml b/.buildkite/examples_build.yml index 86c282d22..a78ec7d20 100644 --- a/.buildkite/examples_build.yml +++ b/.buildkite/examples_build.yml @@ -30,7 +30,7 @@ steps: key: "build_documentation" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia --color=yes -O0 --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate(); Pkg.precompile(; strict=true)'" - - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia --color=yes -O0 --project=docs/ docs/make.jl" + - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia --color=yes -t auto --check-bounds=no --project=docs/ docs/make.jl" agents: queue: ClimaOcean-docs diff --git a/docs/make.jl b/docs/make.jl index 4103341e6..f2e0001fc 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -19,7 +19,8 @@ const OUTPUT_DIR = joinpath(@__DIR__, "src/literated") to_be_literated = [ "single_column_os_papa_simulation.jl", "one_degree_simulation.jl", - "near_global_ocean_simulation.jl" + "near_global_ocean_simulation.jl", + "coupled_simulation.jl", ] for file in to_be_literated @@ -44,6 +45,7 @@ pages = [ "Single-column ocean simulation" => "literated/single_column_os_papa_simulation.md", "One-degree ocean--sea ice simulation" => "literated/one_degree_simulation.md", "Near-global ocean simulation" => "literated/near_global_ocean_simulation.md", + "Coupled atmosphere--ocean--sea ice simulation" => "literated/coupled_simulation.md", ], "Vertical grids" => "vertical_grids.md", diff --git a/experiments/coupled_simulation.jl b/examples/coupled_simulation.jl similarity index 53% rename from experiments/coupled_simulation.jl rename to examples/coupled_simulation.jl index 66dd61b57..a5f5d34a3 100644 --- a/experiments/coupled_simulation.jl +++ b/examples/coupled_simulation.jl @@ -1,39 +1,48 @@ -##### -##### A one degree coupled Atmosphere - Ocean - Sea Ice model -##### +# # Global coupled atmosphere and ocean--sea ice simulation +# +# This example configures a global ocean--sea ice simulation at 1ᵒ horizontal resolution with +# realistic bathymetry and a few closures including the "Gent-McWilliams" `IsopycnalSkewSymmetricDiffusivity`. +# The atmosphere is represented by a SpeecdyWeather model at T63 resolution (approximately 1.875ᵒ). +# and initialized by temperature, salinity, sea ice concentration, and sea ice thickness +# from the ECCO state estimate. +# +# For this example, we need Oceananigans (the ocean), ClimaSeaIce (the sea ice) and +# SpeedyWeather (the atmosphere), coupled and orchestrated by ClimaOcean (the coupled model). using Oceananigans, ClimaSeaIce, SpeedyWeather, ClimaOcean -using Oceananigans, Oceananigans.Units +using NCDatasets, CairoMakie +using Oceananigans.Units using Printf, Statistics, Dates -##### -##### The Ocean!!! -##### +# ## Ocean and sea-ice model configuration +# The ocean and sea-ice are configured in accordance with the "one_degree_simulation" example +# https://clima.github.io/ClimaOceanDocumentation/dev/literated/one_degree_simulation/ +# The first step is to create the grid with realistic bathymetry. -Nx = 360 # Longitudinal direction -Ny = 180 # Meridional direction -Nz = 40 # Vertical levels +Nx = 360 +Ny = 180 +Nz = 40 r_faces = ExponentialCoordinate(Nz, -6000, 0) grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(7, 7, 6)) # Regridding the bathymetry... + bottom_height = regrid_bathymetry(grid; major_basins=1, interpolation_passes=15) grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) -# Advection +# Now we can specify the numerical details and closures for the ocean simulation. + momentum_advection = WENOVectorInvariant(order=5) tracer_advection = WENO(order=5) - -# Free Surface free_surface = SplitExplicitFreeSurface(grid; substeps=80) -# Parameterizations -catke_closure = RiBasedVerticalDiffusivity() +catke_closure = ClimaOcean.OceanSimulations.default_ocean_closure() eddy_closure = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=1e3, κ_symmetric=1e3) closures = (catke_closure, eddy_closure, VerticalScalarDiffusivity(κ=1e-5, ν=1e-4)) -# The ocean simulation +# The ocean simulation, complete with initial conditions for temperature and salinity from ECCO. + ocean = ocean_simulation(grid; momentum_advection, tracer_advection, @@ -44,17 +53,8 @@ ocean = ocean_simulation(grid; Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()), S=Metadatum(:salinity, dataset=ECCO4Monthly())) -##### -##### The Atmosphere!!! -##### - -spectral_grid = SpectralGrid(trunc=63, nlayers=8, Grid=FullClenshawGrid) -atmosphere = atmosphere_simulation(spectral_grid; output=true) -atmosphere.model.output.output_dt = Hour(3) - -##### -##### The Sea-Ice!!! -##### + +# The sea-ice simulation, complete with initial conditions for sea-ice thickness and concentration from ECCO. dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean; solver=ClimaSeaIce.SeaIceDynamics.SplitExplicitSolver(100)) sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) @@ -62,22 +62,32 @@ sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Monthly()), ℵ=Metadatum(:sea_ice_concentration, dataset=ECCO4Monthly())) -##### -##### Coupled model -##### +# ## Atmosphere model configuration +# The atmosphere is provided by SpeedyWeather.jl. Here we configure a T63L8 model with a 3 hour output interval. +# The `atmosphere_simulation` function takes care of building an atmosphere model with appropriate +# hooks for ClimaOcean to compute intercomponent fluxes. We also set the output interval to 3 hours. + +spectral_grid = SpectralGrid(trunc=63, nlayers=8, Grid=FullClenshawGrid) +atmosphere = atmosphere_simulation(spectral_grid; output=true) +atmosphere.model.output.output_dt = Hour(3) + +# ## The coupled model +# Now we can build the coupled model. We need to specify the time step for the coupled model. +# We decide to step the global model every 2 atmosphere time steps. (i.e. the ocean and the +# sea-ice will be stepped every two atmosphere time steps). -# The ocean goes twice as slow as the atmosphere! Δt = 2 * convert(eltype(grid), atmosphere_model.time_stepping.Δt_sec) -# Remember in the future that reflected radiation is computed independently by speedy -# so we need to communicate albedo in some way if this reflected radiation is to be -# absorbed by clouds. At the moment, speedy does not provide longwave down, until this -# is provided we just estimate a lower emissivity -# *** THIS IS A HACK TO BE REMOVED *** +# We build the complete model. Since radiation is idealized in this example, we set the emissivities to zero. + radiation = Radiation(ocean_emissivity=0.0, sea_ice_emissivity=0.0) earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) earth = Oceananigans.Simulation(earth_model; Δt, stop_time=360days) +# ## Running the simulation +# We can now run the simulation. We add a callback to monitor the progress of the simulation +# and write outputs to disk every 3 hours. + wall_time = Ref(time_ns()) function progress(sim) @@ -142,24 +152,95 @@ earth.output_writers[:fluxes] = JLD2Writer(earth.model.ocean.model, fluxes; schedule=TimeInterval(3600 * 3), filename="intercomponent_fluxes.jld2") -# And add it as a callback to the simulation. add_callback!(earth, progress, IterationInterval(100)) Oceananigans.run!(earth) -#### -#### Visualize!!! -#### - -using NCDatasets, GLMakie +# ## Visualizing the results +# We can visualize some of the results. Here we plot the surface speeds in the atmosphere, ocean, and sea-ice +# as well as the 2m temperature in the atmosphere, the sea surface temperature, and the sensible and latent heat +# fluxes at the atmosphere-ocean interface. -SWO = Dataset("run_0002/output.nc") +SWO = Dataset("run_0001/output.nc") Ta = SWO["temp"][:, :, 8, :] qa = SWO["humid"][:, :, 8, :] ua = SWO["u"][:, :, 8, :] va = SWO["v"][:, :, 8, :] sp = sqrt.(ua.^2 + va.^2) -# Surface fields + SST = FieldTimeSeries("ocean_surface_fields.jld2", "T") SSS = FieldTimeSeries("ocean_surface_fields.jld2", "S") +SSU = FieldTimeSeries("ocean_surface_fields.jld2", "u") +SSV = FieldTimeSeries("ocean_surface_fields.jld2", "v") + +SIU = FieldTimeSeries("sea_ice_fields.jld2", "u") +SIV = FieldTimeSeries("sea_ice_fields.jld2", "v") +SIA = FieldTimeSeries("sea_ice_fields.jld2", "ℵ") + +Qcao = FieldTimeSeries("intercomponent_fluxes.jld2", "Qcao") +Qvao = FieldTimeSeries("intercomponent_fluxes.jld2", "Qvao") + +uotmp = Field{Face, Center, Nothing}(SST.grid) +votmp = Field{Center, Face, Nothing}(SST.grid) + +uitmp = Field{Face, Center, Nothing}(SST.grid) +vitmp = Field{Center, Face, Nothing}(SST.grid) +atmp = Field{Center, Center, Nothing}(SST.grid) + +sotmp = Field(sqrt(uotmp^2 + votmp^2)) +sitmp = Field(sqrt(uitmp^2 + vitmp^2) * atmp) + +iter = Observable(1) +san = @lift sp[:, :, $iter] +son = @lift begin + set!(uotmp, SSU[$iter * 2]) + set!(votmp, SSV[$iter * 2]) + compute!(sotmp) + interior(sotmp, :, :, 1) +end + +ssn = @lift begin + set!(uitmp, SIU[$iter * 2]) + set!(vitmp, SIV[$iter * 2]) + set!(atmp, SIA[$iter * 2]) + compute!(sitmp) + interior(sitmp, :, :, 1) +end + +fig = Figure(size = (1200, 800)) +ax2 = Axis(fig[1, 1], title = "Surface speed, atmosphere (m/s)") +hm2 = heatmap!(ax2, san; colormap = :deep) +ax1 = Axis(fig[1, 2], title = "Surface speed, ocean (m/s)") +hm = heatmap!(ax1, son; colormap = :deep) +ax3 = Axis(fig[1, 3], title = "Surface speed, sea-ice (m/s)") +hm = heatmap!(ax3, ssn; colormap = :deep) + +record(fig, "surface_speeds.mp4", iter, framerate = 15) do i + iter[] = i +endnothing #hide + +# ![](surface_speeds.mp4) + + +Tan = @lift Ta[:, :, $iter] +Ton = @lift interior(SST[$iter * 2], :, :, 1) +Qcn = @lift interior(Qcao[$iter * 2], :, :, 1) +Qvn = @lift interior(Qvao[$iter * 2], :, :, 1) + +fig = Figure(size = (1200, 800)) +ax1 = Axis(fig[1, 1], title = "2m Temperature, atmosphere (K)") +hm = heatmap!(ax1, Tan; colormap = :plasma) +ax2 = Axis(fig[1, 2], title = "Sea Surface Temperature (C)") +hm2 = heatmap!(ax2, Ton; colormap = :plasma) +ax3 = Axis(fig[2, 1], title = "Sensible heat flux (W/m²)") +hm3 = heatmap!(ax3, Qcn; colormap = :balance, colorrange = (-200, 200)) +ax4 = Axis(fig[2, 2], title = "Latent heat flux (W/m²)") +hm4 = heatmap!(ax4, Qvn; colormap = :balance, colorrange = (-200, 200)) + +record(fig, "surface_temperature_and_heat_flux.mp4", 1:length(sp[1, 1, :])) do i + iter[] = i +end +nothing #hide + +# ![](surface_temperature_and_heat_flux.mp4) From 5452faf47c11cc564fb7ea447c60abeb75cffd8c Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 8 Sep 2025 12:02:13 +0200 Subject: [PATCH 51/56] bugfix --- examples/coupled_simulation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/coupled_simulation.jl b/examples/coupled_simulation.jl index a5f5d34a3..eccd8c3a4 100644 --- a/examples/coupled_simulation.jl +++ b/examples/coupled_simulation.jl @@ -76,7 +76,7 @@ atmosphere.model.output.output_dt = Hour(3) # We decide to step the global model every 2 atmosphere time steps. (i.e. the ocean and the # sea-ice will be stepped every two atmosphere time steps). -Δt = 2 * convert(eltype(grid), atmosphere_model.time_stepping.Δt_sec) +Δt = 2 * convert(eltype(grid), atmosphere.model.time_stepping.Δt_sec) # We build the complete model. Since radiation is idealized in this example, we set the emissivities to zero. From a153dd9e9a197a71b44d953d35a7b9a5ed318a32 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 9 Sep 2025 09:37:22 +0200 Subject: [PATCH 52/56] Update examples/coupled_simulation.jl Co-authored-by: Gregory L. Wagner --- examples/coupled_simulation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/coupled_simulation.jl b/examples/coupled_simulation.jl index eccd8c3a4..d268773a2 100644 --- a/examples/coupled_simulation.jl +++ b/examples/coupled_simulation.jl @@ -6,7 +6,7 @@ # and initialized by temperature, salinity, sea ice concentration, and sea ice thickness # from the ECCO state estimate. # -# For this example, we need Oceananigans (the ocean), ClimaSeaIce (the sea ice) and +# For this example, we need Oceananigans.HydrostaticFreeSurfaceModel (the ocean), ClimaSeaIce.SeaIceModel (the sea ice) and # SpeedyWeather (the atmosphere), coupled and orchestrated by ClimaOcean (the coupled model). using Oceananigans, ClimaSeaIce, SpeedyWeather, ClimaOcean From 3ba6b53dacb36f3d3f115e74a9661a674fb19079 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 9 Sep 2025 09:37:30 +0200 Subject: [PATCH 53/56] Update examples/coupled_simulation.jl Co-authored-by: Gregory L. Wagner --- examples/coupled_simulation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/coupled_simulation.jl b/examples/coupled_simulation.jl index d268773a2..fa7d9bce5 100644 --- a/examples/coupled_simulation.jl +++ b/examples/coupled_simulation.jl @@ -7,7 +7,7 @@ # from the ECCO state estimate. # # For this example, we need Oceananigans.HydrostaticFreeSurfaceModel (the ocean), ClimaSeaIce.SeaIceModel (the sea ice) and -# SpeedyWeather (the atmosphere), coupled and orchestrated by ClimaOcean (the coupled model). +# SpeedyWeather (the atmosphere), coupled and orchestrated by ClimaOcean.OceanSeaIceModel (the coupled system). using Oceananigans, ClimaSeaIce, SpeedyWeather, ClimaOcean using NCDatasets, CairoMakie From ba689852c09295522f1ab851fe42cc3a3f8ca4f1 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 9 Sep 2025 09:37:37 +0200 Subject: [PATCH 54/56] Update docs/make.jl Co-authored-by: Gregory L. Wagner --- docs/make.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/make.jl b/docs/make.jl index a959dae62..c19715d34 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -21,7 +21,7 @@ to_be_literated = [ "single_column_os_papa_simulation.jl", "one_degree_simulation.jl", "near_global_ocean_simulation.jl", - "coupled_simulation.jl", + "atmosphere_ocean_simulation.jl", ] for file in to_be_literated From 255f44c48cb6ec0cecfc0cfa4d17d0f58eb055b7 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 9 Sep 2025 09:42:22 +0200 Subject: [PATCH 55/56] corrections --- Project.toml | 3 +-- ...lation.jl => atmosphere_ocean_simulation.jl} | 17 +++++++---------- 2 files changed, 8 insertions(+), 12 deletions(-) rename examples/{coupled_simulation.jl => atmosphere_ocean_simulation.jl} (96%) diff --git a/Project.toml b/Project.toml index d7993a69f..36225a9f8 100644 --- a/Project.toml +++ b/Project.toml @@ -62,7 +62,6 @@ NCDatasets = "0.12, 0.13, 0.14" Oceananigans = "0.97.6, 0.98" OffsetArrays = "1.14" PrecompileTools = "1" -PythonCall = "0.9" Reactant = "0.2.45" Scratch = "1" SeawaterPolynomials = "0.3.5" @@ -80,4 +79,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "CondaPkg", "SpeedyWeather", "Reactant"] diff --git a/examples/coupled_simulation.jl b/examples/atmosphere_ocean_simulation.jl similarity index 96% rename from examples/coupled_simulation.jl rename to examples/atmosphere_ocean_simulation.jl index fa7d9bce5..691264345 100644 --- a/examples/coupled_simulation.jl +++ b/examples/atmosphere_ocean_simulation.jl @@ -7,7 +7,7 @@ # from the ECCO state estimate. # # For this example, we need Oceananigans.HydrostaticFreeSurfaceModel (the ocean), ClimaSeaIce.SeaIceModel (the sea ice) and -# SpeedyWeather (the atmosphere), coupled and orchestrated by ClimaOcean.OceanSeaIceModel (the coupled system). +# SpeedyWeather.PrimitiveWetModel (the atmosphere), coupled and orchestrated by ClimaOcean.OceanSeaIceModel (the coupled system). using Oceananigans, ClimaSeaIce, SpeedyWeather, ClimaOcean using NCDatasets, CairoMakie @@ -21,10 +21,10 @@ using Printf, Statistics, Dates Nx = 360 Ny = 180 -Nz = 40 +Nz = 30 -r_faces = ExponentialCoordinate(Nz, -6000, 0) -grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(7, 7, 6)) +r_faces = ExponentialCoordinate(Nz, -4000, 0) +grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(6, 6, 5)) # Regridding the bathymetry... @@ -163,14 +163,12 @@ Oceananigans.run!(earth) SWO = Dataset("run_0001/output.nc") -Ta = SWO["temp"][:, :, 8, :] -qa = SWO["humid"][:, :, 8, :] -ua = SWO["u"][:, :, 8, :] -va = SWO["v"][:, :, 8, :] +Ta = reverse(SWO["temp"][:, :, 8, :], dims=2) +ua = reverse(SWO["u"][:, :, 8, :], dims=2) +va = reverse(SWO["v"][:, :, 8, :], dims=2) sp = sqrt.(ua.^2 + va.^2) SST = FieldTimeSeries("ocean_surface_fields.jld2", "T") -SSS = FieldTimeSeries("ocean_surface_fields.jld2", "S") SSU = FieldTimeSeries("ocean_surface_fields.jld2", "u") SSV = FieldTimeSeries("ocean_surface_fields.jld2", "v") @@ -222,7 +220,6 @@ endnothing #hide # ![](surface_speeds.mp4) - Tan = @lift Ta[:, :, $iter] Ton = @lift interior(SST[$iter * 2], :, :, 1) Qcn = @lift interior(Qcao[$iter * 2], :, :, 1) From 6aef9ba4140f0345f6d44e9a1fbe76dd016c4602 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 29 Sep 2025 11:44:48 +0200 Subject: [PATCH 56/56] new versions --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index d7993a69f..9ab9b1bee 100644 --- a/Project.toml +++ b/Project.toml @@ -48,7 +48,7 @@ ClimaOceanSpeedyWeatherExt = "SpeedyWeather" Adapt = "4" CFTime = "0.1, 0.2" CUDA = "4, 5" -ClimaSeaIce = "0.3.6" +ClimaSeaIce = "0.3.7" CondaPkg = "0.2.28" CubicSplines = "0.2" DataDeps = "0.7" @@ -59,7 +59,7 @@ JLD2 = "0.4, 0.5, 0.6" KernelAbstractions = "0.9" MPI = "0.20" NCDatasets = "0.12, 0.13, 0.14" -Oceananigans = "0.97.6, 0.98" +Oceananigans = "0.99" OffsetArrays = "1.14" PrecompileTools = "1" PythonCall = "0.9"