From c0751aea303ad04db235936ee036570a34638bf1 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 18 Oct 2023 17:37:03 -0400 Subject: [PATCH] diagnostic variables with Base.@kwdef --- src/dynamics/diagnostic_variables.jl | 238 +++++++++++---------------- src/physics/column_variables.jl | 1 + 2 files changed, 94 insertions(+), 145 deletions(-) diff --git a/src/dynamics/diagnostic_variables.jl b/src/dynamics/diagnostic_variables.jl index e8ed390b7..e24bb1593 100644 --- a/src/dynamics/diagnostic_variables.jl +++ b/src/dynamics/diagnostic_variables.jl @@ -1,85 +1,61 @@ +# const LTM = LowerTriangularMatrix # already defined in prognostic_variables + """ - Tendencies{Grid<:AbstractGrid,NF<:AbstractFloat} - -Struct holding the tendencies of the prognostic spectral variables for a given layer.""" -struct Tendencies{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} - vor_tend ::LowerTriangularMatrix{Complex{NF}} # Vorticity of horizontal wind field [1/s] - div_tend ::LowerTriangularMatrix{Complex{NF}} # Divergence of horizontal wind field [1/s] - temp_tend ::LowerTriangularMatrix{Complex{NF}} # Absolute temperature [K] - humid_tend::LowerTriangularMatrix{Complex{NF}} # Specific humidity [g/kg] - u_tend ::LowerTriangularMatrix{Complex{NF}} # zonal velocity (spectral) - v_tend ::LowerTriangularMatrix{Complex{NF}} # meridional velocity (spectral) - u_tend_grid ::Grid # zonal velocity (grid) - v_tend_grid ::Grid # meridinoal velocity (grid) - temp_tend_grid ::Grid # temperature - humid_tend_grid ::Grid # specific humidity +Tendencies of the prognostic spectral variables for a given layer. +$(TYPEDFIELDS)""" +Base.@kwdef struct Tendencies{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} + nlat_half::Int + trunc::Int + vor_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # Vorticity of horizontal wind field [1/s] + div_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # Divergence of horizontal wind field [1/s] + temp_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # Absolute temperature [K] + humid_tend::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # Specific humidity [g/kg] + u_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # zonal velocity (spectral) + v_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # meridional velocity (spectral) + u_tend_grid ::Grid = zeros(Grid,nlat_half) # zonal velocity (grid) + v_tend_grid ::Grid = zeros(Grid,nlat_half) # meridinoal velocity (grid) + temp_tend_grid ::Grid = zeros(Grid,nlat_half) # temperature + humid_tend_grid ::Grid = zeros(Grid,nlat_half) # specific humidity end +# generator function based on a SpectralGrid function Base.zeros(::Type{Tendencies}, SG::SpectralGrid) - - (;NF, trunc,Grid, nlat_half) = SG - LTM = LowerTriangularMatrix - - # use one more l for size compat with vector quantities - vor_tend = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # vorticity - div_tend = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # divergence - temp_tend = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # absolute Temperature - humid_tend = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # specific humidity - u_tend = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # zonal velocity - v_tend = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # meridional velocity - u_tend_grid = zeros(Grid{NF},nlat_half) # zonal velocity - v_tend_grid = zeros(Grid{NF},nlat_half) # meridional velocity - temp_tend_grid = zeros(Grid{NF},nlat_half) # temperature - humid_tend_grid = zeros(Grid{NF},nlat_half) # specific humidity - - return Tendencies( vor_tend,div_tend,temp_tend,humid_tend, - u_tend,v_tend,u_tend_grid,v_tend_grid, - temp_tend_grid,humid_tend_grid) + (;NF, trunc, Grid, nlat_half) = SG + return Tendencies{NF,Grid{NF}}(;nlat_half,trunc) end """ - GridVariables{NF<:AbstractFloat} - -Struct holding the prognostic spectral variables of a given layer in grid point space.""" -struct GridVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} - vor_grid ::Grid # vorticity - div_grid ::Grid # divergence - temp_grid ::Grid # absolute temperature [K] - temp_grid_prev ::Grid # absolute temperature of previous time step [K] - temp_virt_grid ::Grid # virtual tempereature [K] - humid_grid ::Grid # specific_humidity - geopot_grid ::Grid # geopotential (is that needed?) - u_grid ::Grid # zonal velocity *coslat [m/s] - v_grid ::Grid # meridional velocity *coslat [m/s] - u_grid_prev ::Grid # zonal velocity *coslat of previous time step [m/s] - v_grid_prev ::Grid # meridional velocity *coslat of previous time step [m/s] +Transformed prognostic variables (plus a few others) into grid-point space. +$TYPEDFIELDS.""" +Base.@kwdef struct GridVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} + nlat_half::Int # resolution parameter for any grid + vor_grid ::Grid = zeros(Grid,nlat_half) # vorticity + div_grid ::Grid = zeros(Grid,nlat_half) # divergence + temp_grid ::Grid = zeros(Grid,nlat_half) # absolute temperature [K] + temp_grid_prev ::Grid = zeros(Grid,nlat_half) # absolute temperature of previous time step [K] + temp_virt_grid ::Grid = zeros(Grid,nlat_half) # virtual tempereature [K] + humid_grid ::Grid = zeros(Grid,nlat_half) # specific_humidity + geopot_grid ::Grid = zeros(Grid,nlat_half) # geopotential (is that needed?) + u_grid ::Grid = zeros(Grid,nlat_half) # zonal velocity *coslat [m/s] + v_grid ::Grid = zeros(Grid,nlat_half) # meridional velocity *coslat [m/s] + u_grid_prev ::Grid = zeros(Grid,nlat_half) # zonal velocity *coslat of previous time step [m/s] + v_grid_prev ::Grid = zeros(Grid,nlat_half) # meridional velocity *coslat of previous time step [m/s] end +# generator function based on a SpectralGrid function Base.zeros(::Type{GridVariables},SG::SpectralGrid) - (;NF, Grid, nlat_half) = SG - vor_grid = zeros(Grid{NF},nlat_half) # vorticity - div_grid = zeros(Grid{NF},nlat_half) # divergence - temp_grid = zeros(Grid{NF},nlat_half) # absolute temperature - temp_grid_prev = zeros(Grid{NF},nlat_half) # absolute temperature - temp_virt_grid = zeros(Grid{NF},nlat_half) # virtual temperature - humid_grid = zeros(Grid{NF},nlat_half) # specific humidity - geopot_grid = zeros(Grid{NF},nlat_half) # geopotential - u_grid = zeros(Grid{NF},nlat_half) # zonal velocity *coslat - v_grid = zeros(Grid{NF},nlat_half) # meridonal velocity *coslat - u_grid_prev = zeros(Grid{NF},nlat_half) # zonal velocity *coslat - v_grid_prev = zeros(Grid{NF},nlat_half) # meridonal velocity *coslat - - return GridVariables( vor_grid,div_grid,temp_grid,temp_grid_prev,temp_virt_grid,humid_grid,geopot_grid, - u_grid,v_grid,u_grid_prev,v_grid_prev) + return GridVariables{NF,Grid{NF}}(;nlat_half) end """ - DynamicsVariables{Grid<:AbstractGrid,NF<:AbstractFloat} +Intermediate quantities for the dynamics of a given layer. +$(TYPEDFIELDS)""" +Base.@kwdef struct DynamicsVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} -Struct holding intermediate quantities for the dynamics of a given layer.""" -struct DynamicsVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} + nlat_half::Int + trunc::Int # MULTI-PURPOSE VECTOR (a,b), work array to be reused in various places, examples: # uω_coslat⁻¹, vω_coslat⁻¹ = a,b (all models) @@ -87,58 +63,46 @@ struct DynamicsVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} # uh_coslat⁻¹, vh_coslat⁻¹ = a,b (ShallowWaterModel) # uh_coslat⁻¹_grid, vh_coslat⁻¹_grid = a_grid, b_grid (ShallowWaterModel) # Bernoulli potential: 1/2*(u^2+v^2) + Φ = a,a_grid (ShallowWater + PrimitiveEquation) - a::LowerTriangularMatrix{Complex{NF}} - b::LowerTriangularMatrix{Complex{NF}} - a_grid::Grid - b_grid::Grid + a::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) + b::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) + a_grid::Grid = zeros(Grid,nlat_half) + b_grid::Grid = zeros(Grid,nlat_half) # VERTICAL INTEGRATION - uv∇lnp ::Grid # = (uₖ,vₖ)⋅∇ln(pₛ), pressure flux - uv∇lnp_sum_above::Grid # sum of Δσₖ-weighted uv∇lnp above - div_sum_above ::Grid # sum of div_weighted from top to k - temp_virt ::LowerTriangularMatrix{Complex{NF}} # virtual temperature spectral for geopot - geopot ::LowerTriangularMatrix{Complex{NF}} # geopotential on full layers + uv∇lnp ::Grid = zeros(Grid,nlat_half) # = (uₖ,vₖ)⋅∇ln(pₛ), pressure flux + uv∇lnp_sum_above::Grid = zeros(Grid,nlat_half) # sum of Δσₖ-weighted uv∇lnp above + div_sum_above ::Grid = zeros(Grid,nlat_half) # sum of div_weighted from top to k + + # virtual temperature spectral for geopot, geopotential on full layers + temp_virt ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) + geopot ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # VERTICAL VELOCITY (̇̇dσ/dt) - σ_tend ::Grid # = dσ/dt, on half levels below, at k+1/2 + σ_tend::Grid = zeros(Grid,nlat_half) # = dσ/dt, on half levels below, at k+1/2 end +# generator function based on a SpectralGrid function Base.zeros(::Type{DynamicsVariables}, SG::SpectralGrid) - (;NF, trunc, Grid, nlat_half) = SG - - # MULTI-PURPOSE VECTOR (a,b), work array to be reused in various places - a = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) - b = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) - a_grid = zeros(Grid{NF},nlat_half) - b_grid = zeros(Grid{NF},nlat_half) - - # VERTICAL INTEGRATION - uv∇lnp = zeros(Grid{NF},nlat_half) # = (uₖ,vₖ)⋅∇ln(pₛ), pressure flux - uv∇lnp_sum_above= zeros(Grid{NF},nlat_half) # sum of Δσₖ-weighted uv∇lnp from 1:k-1 - div_sum_above = zeros(Grid{NF},nlat_half) # sum of Δσₖ-weighted div from 1:k-1 - temp_virt = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) - geopot = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) - - # VERTICAL VELOCITY (̇̇dσ/dt) - σ_tend = zeros(Grid{NF},nlat_half) # = dσ/dt, on half levels below, at k+1/2 - - return DynamicsVariables( a,b,a_grid,b_grid, - uv∇lnp,uv∇lnp_sum_above,div_sum_above, - temp_virt,geopot, - σ_tend) + return DynamicsVariables{NF,Grid{NF}}(;nlat_half,trunc) end +""" +All diagnostic variables for a given layer: tendencies, prognostic varibles on the grid, +and intermediate dynamics variables. +$(TYPEDFIELDS)""" struct DiagnosticVariablesLayer{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} + npoints ::Int # number of grid points + k ::Int # which vertical model level? + tendencies ::Tendencies{NF,Grid} grid_variables ::GridVariables{NF,Grid} dynamics_variables ::DynamicsVariables{NF,Grid} - npoints ::Int # number of grid points - k ::Int # which vertical model level? temp_average ::Base.RefValue{NF} # average temperature for this level end +# generator function based on a SpectralGrid function Base.zeros(::Type{DiagnosticVariablesLayer}, SG::SpectralGrid, k::Integer=0) # use k=0 (i.e. unspecified) as default @@ -147,63 +111,46 @@ function Base.zeros(::Type{DiagnosticVariablesLayer}, grid_variables = zeros(GridVariables,SG) dynamics_variables = zeros(DynamicsVariables,SG) temp_average = Ref(zero(SG.NF)) - return DiagnosticVariablesLayer(tendencies,grid_variables,dynamics_variables,npoints,k,temp_average) + return DiagnosticVariablesLayer(npoints,k,tendencies,grid_variables,dynamics_variables,temp_average) end -struct SurfaceVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} - pres_grid::Grid # log surface pressure - pres_tend::LowerTriangularMatrix{Complex{NF}} # tendency of it - pres_tend_grid::Grid # gridded tendency - - ∇lnp_x::Grid # zonal gradient of log surf pressure - ∇lnp_y::Grid # meridional gradient of log surf pres - - u_mean_grid::Grid # vertical average of: zonal velocity *coslat - v_mean_grid::Grid # meridional velocity *coslat - div_mean_grid::Grid # divergence - div_mean::LowerTriangularMatrix{Complex{NF}} # divergence (in spectral though) +""" +Diagnostic variables for the surface layer. +$(TYPEDFIELDS)""" +Base.@kwdef struct SurfaceVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} + + nlat_half::Int + trunc::Int + npoints::Int + + # log surface pressure, tendency of it and gridded tendency + pres_grid::Grid = zeros(Grid,nlat_half) + pres_tend::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) + pres_tend_grid::Grid = zeros(Grid,nlat_half) + + ∇lnp_x::Grid = zeros(Grid,nlat_half) # zonal gradient of log surf pressure + ∇lnp_y::Grid = zeros(Grid,nlat_half) # meridional gradient of log surf pres + + u_mean_grid::Grid = zeros(Grid,nlat_half) # vertical average of: zonal velocity *coslat + v_mean_grid::Grid = zeros(Grid,nlat_half) # meridional velocity *coslat + div_mean_grid::Grid = zeros(Grid,nlat_half) # divergence + div_mean::LTM{Complex{NF}} = zeros(LTM{Complex{NF}},trunc+2,trunc+1) # divergence (in spectral though) - precip_large_scale::Grid - precip_convection::Grid - - npoints::Int # number of grid points + precip_large_scale::Grid = zeros(Grid,nlat_half) # large scale precipitation (for output) + precip_convection::Grid = zeros(Grid,nlat_half) # convective precipitation (for output) end +# generator function based on a SpectralGrid function Base.zeros(::Type{SurfaceVariables}, SG::SpectralGrid) (;NF, trunc, Grid, nlat_half, npoints) = SG - - # log of surface pressure and tendency thereof - pres_grid = zeros(Grid{NF},nlat_half) - pres_tend = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) - pres_tend_grid = zeros(Grid{NF},nlat_half) - - # gradients of log surface pressure - ∇lnp_x = zeros(Grid{NF},nlat_half) # zonal gradient of log surf pressure - ∇lnp_y = zeros(Grid{NF},nlat_half) # meridional gradient of log surf pres - - # vertical averaged (weighted by σ level thickness) velocities (*coslat) and divergence - u_mean_grid = zeros(Grid{NF},nlat_half) - v_mean_grid = zeros(Grid{NF},nlat_half) - div_mean_grid = zeros(Grid{NF},nlat_half) - div_mean = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) - - # precipitation fields - precip_large_scale = zeros(Grid{NF},nlat_half) - precip_convection = zeros(Grid{NF},nlat_half) - - return SurfaceVariables(pres_grid,pres_tend,pres_tend_grid, - ∇lnp_x,∇lnp_y, - u_mean_grid,v_mean_grid,div_mean_grid,div_mean, - precip_large_scale,precip_convection, - npoints) + return SurfaceVariables{NF,Grid{NF}}(;nlat_half,trunc,npoints) end """ - DiagnosticVariables{Grid<:AbstractGrid,NF<:AbstractFloat} - -Struct holding the diagnostic variables.""" +All diagnostic variables. +$(TYPEDFIELDS)""" struct DiagnosticVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF},Model<:ModelSetup} layers ::Vector{DiagnosticVariablesLayer{NF,Grid}} surface ::SurfaceVariables{NF,Grid} @@ -216,6 +163,7 @@ struct DiagnosticVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF},Model<:Model scale::Base.RefValue{NF} # vorticity and divergence are scaled by radius end +# generator function based on a SpectralGrid function Base.zeros( ::Type{DiagnosticVariables}, SG::SpectralGrid, diff --git a/src/physics/column_variables.jl b/src/physics/column_variables.jl index 257757cc5..d4a485e66 100644 --- a/src/physics/column_variables.jl +++ b/src/physics/column_variables.jl @@ -37,6 +37,7 @@ function get_column!( C::ColumnVariables, C.geopot[k] = layer.grid_variables.geopot_grid[ij] end + # TODO skin = surface approximation for now C.skin_temperature_sea = P.ocean.sea_surface_temperature[ij] C.skin_temperature_land = P.land.land_surface_temperature[ij] end