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,25 @@
+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
+
+include("conservative_regridding.jl")
+include("speedy_atmosphere_simulations.jl")
+include("speedy_weather_exchanger.jl")
+# include("speedy_weather_fluxes.jl")
+
+end # module ClimaOceanSpeedyWeatherExt

ext/ClimaOceanSpeedyWeatherExt/conservative_regridding.jl

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

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,124 @@
+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) = 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
+# 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)
+
+    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(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
+
+    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)
+    exchanger = interfaces.exchanger
+    regridder = exchanger.regridder
+    exchange_grid = interfaces.exchanger.exchange_grid
+    exchange_state = 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  = 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ℓ
+
+    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
+
+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