SpeedyWeather ensemble online training

[milankl]

Just throwing in some first code snippets for ensemble online training with SpeedyWeather, will convert to a proper PR soon

We need

• A callback that interpolates any variable at a given schedule (says once a day, we have a Schedule now, Scheduler for callbacks SpeedyWeather/SpeedyWeather.jl#488) onto some full grid that the NN can treat like an image, see MyInterpolatedVorticity below
• A callback that afterwards takes that and writes it into some shared array (a column of it) visible to the NN caller too.
• A distributed ensemble script, like illustrated here Ensemble simulations SpeedyWeather/SpeedyWeather.jl#380

using SpeedyWeather

struct MyInterpolatedVorticity{Grid, NF_model, Grid_model} <: SpeedyWeather.AbstractCallback
    "Interpolate only when scheduled"
    schedule::Schedule

    "From which layer to interpolate"
    layer::Int

    "Variable on interpolated grid"
    var::Grid

    "Interpolator to interpolate model grid into var"
    interpolator::RingGrids.AnvilInterpolator{NF_model, Grid_model}
end

function MyInterpolatedVorticity(
    SG::SpectralGrid;
    layer::Integer = 1,
    nlat_half::Integer = 32,
    NF::Type{<:AbstractFloat} = Float32,
    Grid::Type{<:RingGrids.AbstractGrid} = FullGaussianGrid{NF},
    schedule::Schedule = Schedule(every=Day(1))
)
    n_points = RingGrids.get_npoints(Grid, nlat_half)
    var = zeros(Grid, nlat_half) 
    interpolator = RingGrids.AnvilInterpolator(SG.NF, SG.Grid, SG.nlat_half, n_points)
    RingGrids.update_locator!(interpolator, RingGrids.get_latdlonds(var)...)
    MyInterpolatedVorticity{Grid, SG.NF, SG.Grid}(schedule, layer, var, interpolator)
end

function SpeedyWeather.initialize!(
    callback::MyInterpolatedVorticity,
    progn::PrognosticVariables,
    diagn::DiagnosticVariables,
    model::ModelSetup,
)
    initialize!(callback.schedule, progn.clock)
end

function SpeedyWeather.callback!(
    callback::MyInterpolatedVorticity,
    progn::PrognosticVariables,
    diagn::DiagnosticVariables,
    model::ModelSetup,
)
    isscheduled(callbacks.schedule, progn.clock) || return nothing
    k = callback.layer
    (;vor_grid) = diagn.layers[k].grid_variables
    (;var, interpolator) = callback
    RingGrids.interpolate!(var, vor_grid, interpolator)
end

SpeedyWeather.finish!(::MyInterpolatedVorticity, args...) = nothing

@sandreza