Merge pull request #154 from maximilian-gelbrecht/set-vars

set_var! functions to easily set variables in PrognosticVariables

src/initial_conditions.jl

@@ -75,27 +75,14 @@ function initialize_from_file!(progn_new::PrognosticVariables{NF},M::ModelSetup)
     version = restart_file["version"]   # currently unused, TODO check for compat with version
     time = restart_file["time"]         # currently unused
 
-    # SPECTRAL TRUNCATION/INTERPOLATION to new resolution and conversion to NF
-    for (layer_old,layer_new) in zip(progn_old.layers,progn_new.layers)
+    var_names = [propertynames(progn_old.layers[1].leapfrog[1])..., :pres]
 
-        layer_old_lf1 = layer_old.leapfrog[1]
-        layer_new_lf1 = layer_new.leapfrog[1]
-
-        vars_old = (getproperty(layer_old_lf1,prop) for prop in propertynames(layer_old_lf1))
-        vars_new = (getproperty(layer_new_lf1,prop) for prop in propertynames(layer_new_lf1))
-
-        for (var_old,var_new) in zip(vars_old,vars_new)
-            lmax,mmax = size(var_new) .- (2,1)
-            var_old_trunc = spectral_truncation(var_old,lmax+1,mmax)
-            var_new .= convert(LowerTriangularMatrix{Complex{NF}},var_old_trunc)
+    for var_name in var_names
+        if has(progn_new, var_name) 
+            var = get_var(progn_old, var_name) 
+            set_var!(progn_new, var_name, var)
         end
-    end
-
-    # same for surface pressure
-    pres = progn_new.pres.leapfrog[1]
-    lmax,mmax = size(pres) .- (2,1)
-    pres .= convert(LowerTriangularMatrix{Complex{NF}},
-                                    spectral_truncation(progn_old.pres.leapfrog[1],lmax+1,mmax))
+    end 
 
     return progn_new
 end

src/models.jl

@@ -20,6 +20,8 @@ struct BarotropicModel{NF<:AbstractFloat, D<:AbstractDevice} <: ModelSetup{D}
     device_setup::DeviceSetup{D}
 end
 
+has(::Type{BarotropicModel{NF,D}}, var_name::Symbol) where {NF,D} = var_name in [:vor]
+
 """
     M = ShallowWaterModel(  ::Parameters,
                             ::Constants,
@@ -41,6 +43,9 @@ struct ShallowWaterModel{NF<:AbstractFloat, D<:AbstractDevice} <: ModelSetup{D}
     device_setup::DeviceSetup{D}
 end
 
+has(::Type{ShallowWaterModel{NF,D}}, var_name::Symbol) where {NF,D} = var_name in [:vor, :div, :pres]
+
+
 """
     M = PrimitiveEquationModel( ::Parameters,
                                 ::Constants,
@@ -62,3 +67,13 @@ struct PrimitiveEquationModel{NF<:AbstractFloat,D<:AbstractDevice} <: ModelSetup
     implicit::Implicit{NF}
     device_setup::DeviceSetup{D}
 end
+
+has(::Type{PrimitiveEquationModel{NF,D}}, var_name::Symbol) where {NF,D} = var_name in [:vor, :div, :temp, :humid, :pres]
+
+"""
+    has(::ModelSetup, var_name::Symbol)
+
+Returns true if the model `M` has a prognostic variable `var_name`, false otherwise. The default fallback is that all variables are included. 
+"""
+has(::Type{ModelSetup{D}}, var_name::Symbol) where D = var_name in [:vor, :div, :temp, :humid, :pres] 
+has(M::ModelSetup, var_name) = has(typeof(M), var_name)

src/prognostic_variables.jl

@@ -18,7 +18,7 @@ struct SurfaceLeapfrog{NF<:AbstractFloat}
     leapfrog::Vector{LowerTriangularMatrix{Complex{NF}}}     # 2-element vector for two leapfrog time steps
 end
 
-struct PrognosticVariables{NF<:AbstractFloat}
+struct PrognosticVariables{NF<:AbstractFloat, M<:ModelSetup}
     # data
     layers::Vector{PrognosticVariablesLeapfrog{NF}} # each element = 1 vertical layer (incl leapfrog dim)    
     pres::SurfaceLeapfrog{NF}                       # 2-element leapfrog vec of log of surface pressure [log(hPa)]
@@ -44,12 +44,11 @@ end
 # reduce size of unneeded variables if ModelSetup is provided
 function Base.zeros(::Type{PrognosticVariablesLayer{NF}},model::ModelSetup,lmax::Integer,mmax::Integer) where NF
     # use one more l for size compat with vector quantities
-    vor   = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1)
-    lmax,mmax = model isa BarotropicModel   ? (-2,-1) : (lmax,mmax)   # other variables not needed for BarotropicModel
-    div   = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1)
-    lmax,mmax = model isa ShallowWaterModel ? (-2,-1) : (lmax,mmax)   # other variables not needed for ShallowWaterModel
-    temp  = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1)
-    humid = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1)
+    vor = has(model, :vor) ? zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) : LowerTriangularMatrix{Complex{NF}}(undef, 0, 0)
+    div = has(model, :div) ? zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) : LowerTriangularMatrix{Complex{NF}}(undef, 0, 0)
+    temp = has(model, :temp) ? zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) : LowerTriangularMatrix{Complex{NF}}(undef, 0, 0)
+    humid = has(model, :humid) ? zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) : LowerTriangularMatrix{Complex{NF}}(undef, 0, 0)
+
     return PrognosticVariablesLayer(vor,div,temp,humid)
 end
 
@@ -81,7 +80,7 @@ function Base.zeros(::Type{PrognosticVariables{NF}},
 
     layers = [zeros(PrognosticVariablesLeapfrog{NF},lmax,mmax) for _ in 1:nlev]     # k layers
     pres = zeros(SurfaceLeapfrog{NF},lmax,mmax)                                     # 2 leapfrog time steps for pres
-    return PrognosticVariables(layers,pres,lmax,mmax,N_LEAPFROG,nlev)
+    return PrognosticVariables{NF,ModelSetup}(layers,pres,lmax,mmax,N_LEAPFROG,nlev)
 end
 
 # pass on model to reduce size
@@ -92,7 +91,219 @@ function Base.zeros(::Type{PrognosticVariables{NF}},
                     nlev::Integer) where NF
 
     layers = [zeros(PrognosticVariablesLeapfrog{NF},model,lmax,mmax) for _ in 1:nlev]   # k layers
-    _lmax,_mmax = model isa BarotropicModel ? (-2,-1) : (lmax,mmax)   # pressure not needed for BarotropicModel
-    pres = zeros(SurfaceLeapfrog{NF},_lmax,_mmax)                     # 2 leapfrog time steps for pres
-    return PrognosticVariables(layers,pres,lmax,mmax,N_LEAPFROG,nlev)
+    pres = has(model, :pres) ? zeros(SurfaceLeapfrog{NF},lmax,mmax) : zeros(SurfaceLeapfrog{NF},-2,-1)  # 2 leapfrog time steps for pres       
+    return PrognosticVariables{NF,typeof(model)}(layers,pres,lmax,mmax,N_LEAPFROG,nlev)
+end
+
+has(progn::PrognosticVariables{NF,M}, var_name::Symbol) where {NF,M} = has(M, var_name)
+
+# SET_VAR FUNCTIONS TO ASSIGN NEW VALUES TO PrognosticVariables
+
+"""
+    set_var!(progn::PrognosticVariables{NF},        
+             varname::Symbol, 
+             var::Vector{<:LowerTriangularMatrix};
+             lf::Integer=1) where NF
+
+Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` 
+with values given in `var` a vector with all information for all layers in spectral space.
+"""
+function set_var!(progn::PrognosticVariables{NF}, 
+                  varname::Symbol, 
+                  var::Vector{<:LowerTriangularMatrix};
+                  lf::Integer=1) where NF
+
+    @assert length(var) == length(progn.layers)
+    @assert has(progn, varname)
+
+    for (progn_layer, var_layer) in zip(progn.layers, var)
+        _set_var_core!(getfield(progn_layer.leapfrog[lf], varname), var_layer)
+    end 
+
+    return progn 
+end 
+
+function _set_var_core!(var_old::LowerTriangularMatrix{T}, var_new::LowerTriangularMatrix{R}) where {T,R}
+    lmax,mmax = size(var_old) .- (2,1)
+    var_new_trunc = spectral_truncation!(var_new, lmax+1, mmax)
+    copyto!(var_old, var_new_trunc)
+end 
+
+"""
+    set_var!(progn::PrognosticVariables{NF},        
+             varname::Symbol, 
+             var::Vector{<:AbstractGrid};
+             lf::Integer=1) where NF
+
+Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` 
+with values given in `var` a vector with all information for all layers in grid space.
+"""
+function set_var!(progn::PrognosticVariables{NF}, 
+                  varname::Symbol, 
+                  var::Vector{<:AbstractGrid};
+                  lf::Integer=1) where NF
+
+    @assert length(var) == length(progn.layers)
+
+    var_sph = [spectral(var_layer) for var_layer in var]
+
+    return set_var!(progn, varname, var_sph; lf=lf)
+end 
+
+"""
+    set_var!(progn::PrognosticVariables{NF}, 
+             varname::Symbol, 
+             var::Vector{<:AbstractGrid}, 
+             M::ModelSetup;
+             lf::Integer=1) where NF
+
+Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` 
+with values given in `var` a vector with all information for all layers in grid space.
+"""
+function set_var!(progn::PrognosticVariables{NF}, 
+                  varname::Symbol, 
+                  var::Vector{<:AbstractGrid}, 
+                  M::ModelSetup;
+                  lf::Integer=1) where NF
+
+    @assert length(var) == length(progn.layers)
+
+    var_sph = [spectral(var_layer, M.spectral_transform) for var_layer in var]
+
+    return set_var!(progn, varname, var_sph; lf=lf)
+end 
+
+"""
+    set_var!(progn::PrognosticVariables{NF}, 
+             varname::Symbol, 
+             var::Vector{<:AbstractMatrix}, 
+             Grid::Type{<:AbstractGrid}=FullGaussianGrid;
+             lf::Integer=1) where NF
+
+Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` 
+with values given in `var` a vector with all information for all layers in grid space.
+"""
+function set_var!(progn::PrognosticVariables{NF}, 
+                  varname::Symbol, 
+                  var::Vector{<:AbstractMatrix}, 
+                  Grid::Type{<:AbstractGrid}=FullGaussianGrid;
+                  lf::Integer=1) where NF
+
+    @assert length(var) == length(progn.layers)
+
+    var_grid = [spectral(var_layer, Grid) for var_layer in var]
+
+    return set_var!(progn, varname, var_grid; lf=lf)
+end 
+
+"""
+    function set_var!(progn::PrognosticVariables{NF}, 
+                      varname::Symbol, 
+                      s::Number;
+                      lf::Integer=1) where NF
+
+Sets all values of prognostic variable `varname` at leapfrog index `lf` to the scalar `s`.
+"""
+function set_var!(progn::PrognosticVariables{NF}, 
+                  varname::Symbol, 
+                  s::Number;
+                  lf::Integer=1) where NF
+
+    for progn_layer in progn.layers
+        fill!(getfield(progn_layer.leapfrog[lf], varname), s)
+    end 
+
+    return progn 
+end 
+
+"""
+    set_vorticity!(progn::PrognosticVariables, varargs...; kwargs...)
+
+See [`set_var!`](@ref)
+"""
+set_vorticity!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :vor, varargs...; kwargs...)
+
+"""
+    set_divergence!(progn::PrognosticVariables, varargs...; kwargs...)
+
+See [`set_var!`](@ref)
+"""
+set_divergence!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :div, varargs...; kwargs...)
+
+"""
+    set_temperature!(progn::PrognosticVariables, varargs...; kwargs...)
+
+See [`set_var!`](@ref)
+"""
+set_temperature!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :temp, varargs...; kwargs...)
+
+"""
+    set_humidity!(progn::PrognosticVariables, varargs...; kwargs...)
+
+See [`set_var!`](@ref)
+"""
+set_humidity!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :humid, varargs...; kwargs...)
+
+"""
+    set_pressure!(progn::PrognosticVariables{NF}, 
+                  pressure::LowerTriangularMatrix;
+                  lf::Integer=1) where NF
+
+Sets the prognostic variable with the surface pressure in spectral space at leapfrog index `lf`.
+"""
+function set_pressure!(progn::PrognosticVariables{NF}, 
+                       pressure::LowerTriangularMatrix;
+                       lf::Integer=1) where NF
+
+    _set_var_core!(progn.pres.leapfrog[lf], pressure)
+
+    return progn
 end
+
+"""
+    set_pressure!(progn::PrognosticVariables{NF}, 
+                  pressure::AbstractGrid, 
+                  M::ModelSetup;
+                  lf::Integer=1) where NF
+
+Sets the prognostic variable with the surface pressure in grid space at leapfrog index `lf`.
+"""
+set_pressure!(progn::PrognosticVariables, pressure::AbstractGrid, M::ModelSetup; lf::Integer=1) = set_pressure!(progn, spectral(pressure, M.spectral_transform); lf=lf)
+
+"""
+    set_pressure!(progn::PrognosticVariables{NF}, 
+                  pressure::AbstractGrid, 
+                  lf::Integer=1) where NF
+
+Sets the prognostic variable with the surface pressure in grid space at leapfrog index `lf`.
+"""
+set_pressure!(progn::PrognosticVariables, pressure::AbstractGrid, lf::Integer=1) = set_pressure!(progn, spectral(pressure); lf=lf)
+
+"""
+    set_pressure!(progn::PrognosticVariables{NF}, 
+                  pressure::AbstractMatrix, 
+                  Grid::Type{<:AbstractGrid}, 
+                  lf::Integer=1) where NF
+
+Sets the prognostic variable with the surface pressure in grid space at leapfrog index `lf`.
+"""
+set_pressure!(progn::PrognosticVariables, pressure::AbstractMatrix, Grid::Type{<:AbstractGrid}, lf::Integer=1) = set_pressure!(progn, spectral(pressure, Grid); lf=lf)
+
+"""
+    get_var(progn::PrognosticVariables, var_name::Symbol; lf::Integer=1)
+
+Returns the prognostic variable `var_name` at leapfrog index `lf` as a `Vector{LowerTriangularMatrices}`.
+"""
+function get_var(progn::PrognosticVariables, var_name::Symbol; lf::Integer=1)
+    @assert has(progn, var_name)
+    return [getfield(layer.leapfrog[lf], var_name) for layer in progn.layers]
+end 
+
+get_vorticity(progn::PrognosticVariables; kwargs...) = get_var(progn, :vor; kwargs...)
+get_divergence(progn::PrognosticVariables; kwargs...) = get_var(progn, :div; kwargs...)
+get_temperature(progn::PrognosticVariables; kwargs...) = get_var(progn, :temp; kwargs...)
+get_humidity(progn::PrognosticVariables; kwargs...) = get_var(progn, :humid; kwargs...)
+get_pressure(progn::PrognosticVariables; lf::Integer=1) = progn.pres.leapfrog[lf]
+
+
+

test/run_speedy_with_output.jl

@@ -19,4 +19,8 @@
 
     p = run_speedy(Float64,Grid=OctahedralClenshawGrid,output_grid=:matrix,output_NF=Float32,output=true)
     @test all(isfinite.(p.layers[1].leapfrog[1].vor))
-end
+end
+
+@testset "Restart from output file" begin 
+    progn, diagn, model = initialize_speedy(initial_conditions=:restart, restart_id=1)
+end 

test/runtests.jl

@@ -5,6 +5,7 @@ using Test
 include("utility_functions.jl")
 include("lower_triangular_matrix.jl")
 include("grids.jl")
+include("set_vars.jl")
 
 # GPU/KERNELABSTRACTIONS
 include("kernelabstractions.jl")