set! initial conditions

[milankl]

With the new variable structure in #525 we also need to redesign a convenient way of setting the initial conditions of the prognostic variables. I think we want

• set!(simulation, vor=f) with a function like f(x,y,z) = ... where x,y,z would be coordinates lon, lat, height (probably lond 0-360˚, latd -90˚ to 90˚, sigma 0 to 1)
• set! with a scalar, e.g. set!(simulation, humid=0)
• other ideas?

I've drafted that a bit and just throw this in here to collect this

set!(S::AbstractSimulation; kwargs...) = set!(S.prognostic_variables, S.model.geometry; kwargs...)

function set!(
    progn::PrognosticVariables,
    geometry::Geometry;
    u = nothing,
    v = nothing,
    vor = nothing,
    div = nothing,
    temp = nothing,
    humid = nothing,
    pres = nothing,
    lf::Integer = 1,
    add::Bool = false,
)
    isnothing(vor)   || set!(progn.vor[lf],     vor, geometry; add)
    isnothing(div)   || set!(progn.div[lf],     div, geometry; add)
    isnothing(temp)  || set!(progn.temp[lf],   temp, geometry; add)
    isnothing(humid) || set!(progn.humid[lf], humid, geometry; add)
    isnothing(pres)  || set!(progn.pres[lf],   pres, geometry; add)

    isnothing(u) | isnothing(v) || set_vordiv!(progn.vor[lf], progn.div[lf], vor, div; add)
end

function set!(var::LowerTriangularArray, f::Function, geometry::Geometry, S::SpectralTransform; add)
    grid =
    set!(grid, f, geometry; add)
    transform!(var, grid, S)
end

function set!(var::AbstractGridArray, f::Function, geometry::Geometry; add)
    (; londs, latds, σ_levels_full) = geometry
    kernel(a, b) = add ? a+b : b
    for k in eachgrid(var)
        for ij in eachgridpoint(var)
            var[ij, k] = kernel(var[ij, k], f(londs[ij], latds[ij], σ_levels_full[k]))
        end
    end
end

function set!(var::LowerTriangularArray, f::Real, geometry::Geometry; add)
end

[maximilian-gelbrecht]

I like the idea and the general design.

I'd probably just add a method to set! with an instance of AbstractGrid.

Oh, and I think explicitly denoting Function is not ideal because not all callable objects are subtypes of Function (e.g. overloaded structs are not).

[milankl]

Oceananigans implemented set! here

https://github.com/CliMA/Oceananigans.jl/blob/738d1724e5128502efe5579dd210fc17aacb9089/src/Fields/set!.jl

[maximilian-gelbrecht]

Implemented in #525 , but with one broken test remaining, so this issue stays open.

[milankl]

More set! functionality in #579 for tracers