SpeedyWeather / SpeedyWeather.jl

Play atmospheric modelling like it's LEGO.
https://speedyweather.github.io/SpeedyWeather.jl/dev
MIT License
400 stars 24 forks source link

set! initial conditions #557

Open milankl opened 2 days ago

milankl commented 2 days ago

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

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