Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

set_vars! function to easily set variables in PrognosticVariables #154

Merged
merged 6 commits into from
Sep 29, 2022
Merged

set_vars! function to easily set variables in PrognosticVariables #154

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9486640
initial var_args commit
maximilian-gelbrecht Sep 27, 2022
c6900ca
update, get_... and has funcs
maximilian-gelbrecht Sep 28, 2022
49ad052
doc strings and initialize_from_file
maximilian-gelbrecht Sep 28, 2022
f96ab02
bug fix in initialize_from_file!
maximilian-gelbrecht Sep 28, 2022
a52cbc6
forgot about pressure in the initialize_from_file!
maximilian-gelbrecht Sep 28, 2022
6660f07
docs strings and set_var! for scalar
maximilian-gelbrecht Sep 29, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 6 additions & 19 deletions src/initial_conditions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment on lines +80 to +83
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Neat. I like it.

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
Expand Down
15 changes: 15 additions & 0 deletions src/models.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe the ,D} can be dropped as it's not used (I had to remind myself what D stands for again...)? As far as I understand multi-parametric types in Julia they work like type{T1,T2,T3,T4} to be like type{T1}{T2}{T3}

julia> struct ABC{A,B,C}
           a::Int
           b::Int
       end

julia> abc = ABC{Float64,Float32,Float16}(1,2)
ABC{Float64, Float32, Float16}(1, 2)

julia> t(abc::ABC{T}) where T = T(abc.a)
julia> t(abc)
1.0

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So dispatch is only performed following the first type if the others are not specified.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand why exactly, but it is required, otherwise it won't work. I did try it before. Just again:

 MethodError: no method matching has(::Type{SpeedyWeather.PrimitiveEquationModel{Float64, SpeedyWeather.CPUDevice}}, ::Symbol)
  Closest candidates are:
    has(::Type{SpeedyWeather.BarotropicModel{NF}}, ::Symbol) where NF at ~/Nextcloud/SpeedyFork/SpeedyWeather.jl/src/models.jl:23
    has(::Type{SpeedyWeather.ShallowWaterModel{NF}}, ::Symbol) where NF at ~/Nextcloud/SpeedyFork/SpeedyWeather.jl/src/models.jl:46
    has(::Type{SpeedyWeather.PrimitiveEquationModel{NF}}, ::Symbol) where NF at ~/Nextcloud/SpeedyFork/SpeedyWeather.jl/src/models.jl:71


"""
M = ShallowWaterModel( ::Parameters,
::Constants,
Expand All @@ -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]
maximilian-gelbrecht marked this conversation as resolved.
Show resolved Hide resolved


"""
M = PrimitiveEquationModel( ::Parameters,
::Constants,
Expand All @@ -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]
maximilian-gelbrecht marked this conversation as resolved.
Show resolved Hide resolved

"""
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]
maximilian-gelbrecht marked this conversation as resolved.
Show resolved Hide resolved
has(M::ModelSetup, var_name) = has(typeof(M), var_name)
213 changes: 202 additions & 11 deletions src/prognostic_variables.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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)
maximilian-gelbrecht marked this conversation as resolved.
Show resolved Hide resolved

return PrognosticVariablesLayer(vor,div,temp,humid)
end

Expand Down Expand Up @@ -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
Expand All @@ -92,7 +91,199 @@ 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)
maximilian-gelbrecht marked this conversation as resolved.
Show resolved Hide resolved
@assert has(progn, varname)
Comment on lines +116 to +117
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reminder to self that we should do something about error messages in general. I like the @assert here as they actually print what's wrong. However, often I ended up using throw(BoundsError) which is kind of clear if there's only one @boundscheck before a loop in a function but if there's several ones, you don't know which array sizes don't match. Maybe we can make them somehow more telling...

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't know that before, but @assert let's you specify error messages:
@assert has(progn, varname) "No variable with this name found."


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}
maximilian-gelbrecht marked this conversation as resolved.
Show resolved Hide resolved
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{<:AbstractGrid},
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

"""
set_vorticity!(progn::PrognosticVariables, varargs...; kwargs...)

See [`get_var`](@ref)
"""
set_vorticity!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :vor, varargs...; kwargs...)

"""
set_divergence!(progn::PrognosticVariables, varargs...; kwargs...)

See [`get_var`](@ref)
"""
set_divergence!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :div, varargs...; kwargs...)

"""
set_temperature!(progn::PrognosticVariables, varargs...; kwargs...)

See [`get_var`](@ref)
"""
set_temperature!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :temp, varargs...; kwargs...)

"""
set_humidity!(progn::PrognosticVariables, varargs...; kwargs...)

See [`get_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]



6 changes: 5 additions & 1 deletion test/run_speedy_with_output.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
1 change: 1 addition & 0 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Loading