Add SpeedyWeather extension

Project.toml

@@ -9,10 +9,14 @@ 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"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -24,6 +28,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 +39,7 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [extensions]
 ClimaOceanReactantExt = "Reactant"
+ClimaOceanSpeedyWeatherExt = "SpeedyWeather"
 
 [compat]
 Adapt = "4"
@@ -54,6 +60,7 @@ PrecompileTools = "1"
 Reactant = "0.2.45"
 Scratch = "1"
 SeawaterPolynomials = "0.3.5"
+SpeedyWeather = "0.14.0"
 StaticArrays = "1"
 Statistics = "1.9"
 SurfaceFluxes = "0.11, 0.12"
@@ -67,4 +74,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"]

ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl

@@ -0,0 +1,65 @@
+module ClimaOceanSpeedyWeatherExt
+
+using OffsetArrays
+using KernelAbstractions
+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...)
+
+include("speedy_atmosphere_simulations.jl")
+include("speedy_weather_exchanger.jl")
+# include("speedy_weather_fluxes.jl")
+
+end # module ClimaOceanSpeedyWeatherExt

ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl

@@ -0,0 +1,52 @@
+# 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
+
+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
+    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)
+
+
+

ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl

@@ -0,0 +1,126 @@
+using ConservativeRegridding
+using GeometryOps: CutAtAntimeridianAndPoles, ClosedRing
+using GeoInterface: Polygon, LinearRing
+using Oceananigans
+using Oceananigans.Grids: architecture
+
+import ClimaOcean.OceanSeaIceModels:
+    compute_net_atmosphere_fluxes!
+
+import ClimaOcean.OceanSeaIceModels.InterfaceComputations:
+    atmosphere_exchanger,
+    initialize!,
+    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
+# 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
+    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 = (
+            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(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
+
+    regridder = ConservativeRegridding.Regridder(exchange_cell_matrix, atmosphere_cell_matrix)
+
+    exchanger = (; cpu_surface_state, regridder)
+
+    return exchanger
+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)
+    atmosphere_exchanger = interfaces.exchanger.atmosphere_exchanger
+    regridder = atmosphere_exchanger.regridder
+    exchange_grid = interfaces.exchanger.exchange_grid
+    exchange_state = interfaces.exchanger.exchange_atmosphere_state
+
+    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 
+
+    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)
+    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, 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
+    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(Qc)))
+    ConservativeRegridding.regrid!(atmos.diagnostic_variables.physics.evaporative_flux,    regridder, vec(interior(Mv)))
+
+    return nothing
+end

ext/ClimaOceanSpeedyWeatherExt/speedy_weather_fluxes.jl

@@ -0,0 +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