Add SpeedyWeather extension

[glwagner]

This will facilitate coupled simulations with SpeedyWeather using OceanSeaIceModel (whose name may have to change once this extension goes in...)

[glwagner]

@milankl

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 36 lines in your changes missing coverage. Please review.

Project coverage is 24.09%. Comparing base (5a7ea23) to head (36a6bdf).

Files with missing lines	Patch %	Lines
ext/ClimaOceanSpeedyWeatherExt.jl	0.00%	36 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #463      +/-   ##
==========================================
- Coverage   24.44%   24.09%   -0.36%     
==========================================
  Files          43       44       +1     
  Lines        2454     2490      +36     
==========================================
  Hits          600      600              
- Misses       1854     1890      +36     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[simone-silvestri]

here is the preliminary work: glwagner/ClimAdventures.jl#1

[simone-silvestri]

This was working (simple interpolation atmos -> ocean grid)

    # Get the atmospheric state on the ocean grid
    ua = on_architecture(CPU(), surface_atmosphere_state.u)
    va = on_architecture(CPU(), surface_atmosphere_state.v)
    Ta = on_architecture(CPU(), surface_atmosphere_state.T)
    qa = on_architecture(CPU(), surface_atmosphere_state.q)
    pa = on_architecture(CPU(), surface_atmosphere_state.p)
    Qs = on_architecture(CPU(), surface_atmosphere_state.Qs)
    Qℓ = on_architecture(CPU(), surface_atmosphere_state.Qℓ)
    Mp = on_architecture(CPU(), interpolated_prescribed_freshwater_flux)
    ρf = fluxes.freshwater_density

    λ,  φ,  _ = Oceananigans.Grids.nodes(grid, Center(), Center(), Center(), with_halos=true) 

    λ = Array(vec(on_architecture(CPU(), λ)))
    φ = Array(vec(on_architecture(CPU(), φ)))

    spectral_grid = atmos.model.spectral_grid
    interpolator = RingGrids.AnvilInterpolator(Float32, spectral_grid.Grid, spectral_grid.nlat_half, length(λ))
    RingGrids.update_locator!(interpolator, λ, φ)

    atmosphere_precipitation = (atmos.diagnostic_variables.physics.precip_rate_large_scale .+ 
                                atmos.diagnostic_variables.physics.precip_rate_convection) .* ρf

    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)

    surface_atmosphere_state.u  .= ua 
    surface_atmosphere_state.v  .= va 
    surface_atmosphere_state.T  .= Ta 
    surface_atmosphere_state.q  .= qa 
    surface_atmosphere_state.p  .= pa 
    surface_atmosphere_state.Qs .= Qs
    surface_atmosphere_state.Qℓ .= Qℓ

And fluxes

# Regrid the fluxes from the ocean/sea-ice grid to the atmospheric model grid
# This is a _hack_!! (We are performing a double interpolation)
function regrid_fluxes_to_atmospheric_model!(atmos::SpeedyWeather.Simulation, similarity_theory_fields, total_fluxes)
    
    spectral_grid = atmos.model.spectral_grid

    # 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 
    interpolator = RingGrids.AnvilInterpolator(Float32, FullClenshawGrid, 90, spectral_grid.npoints)
    londs, latds = RingGrids.get_londlatds(spectral_grid.Grid, spectral_grid.nlat_half)
    RingGrids.update_locator!(interpolator, londs, latds)

    regrid_flux!(atmos.diagnostic_variables.physics.sensible_heat_flux,  Qs, interpolator, weights)
    regrid_flux!(atmos.diagnostic_variables.physics.evaporative_flux,    Mv, interpolator, weights)
    regrid_flux!(atmos.diagnostic_variables.physics.surface_longwave_up, Qu, interpolator, weights)

    return nothing
end

function regrid_flux!(speedy_flux, climaocean_flux, interpolator, weights)
    interpolate!(tmp_field, climaocean_flux, weights) 
    tmp_field = on_architecture(CPU(), interior(tmp_field, :, :, 1))
    
    # This should be solved with conservative remapping?
    tmp_field[isnan.(tmp_field)] .= 0 # We do not have antarctic

    flux = FullClenshawGrid(vec(reverse(tmp_field, dims=2)))
    RingGrids.interpolate!(speedy_flux, flux, interpolator)
    return nothing
end

[glwagner]

Should this really be regridded?

[simone-silvestri]

I am not still not sure what we want to do with the upwelling longwave. We can either compute it in ClimaOcean or pass the temperature to the atmosphere which computes its upwelling radiation.
If we separate radiation it might be easier, but for the moment we can compute it in ClimaOcean and pass it to the atmosphere

[glwagner]

I guess with radiation computed in EarthSystemModel, we will indeed have to regrid the radiative fluxes back to the atmosphere?

[simone-silvestri]

yeah, let's start with regridding also the longwave up which might be propedeutic for a radiation computed in the EarthSystemModel

[simone-silvestri]

An example of a coupled simulation (1 degree ocean and sea-ice, trunc 63 in the atmosphere)

heat.mp4

surface_speed.mp4

[glwagner]

Very nice! I wonder if we should put an example in. We could do 2 deg ocean perhaps, then it adds something relative to the existing 1 degree and might be cheap.

[simone-silvestri]

I have added it as an example, let's see how long it takes (maybe only 60 days of evolution is enough?) I will also try to reduce the resolution

[giordano]

--check-bounds=no 😱

[simone-silvestri]

😅 just checking the speed of the solution. In any case, this is an example build, not a testing pipeline, so we really want to have it as fast as possible.