diagnostic variables with Base.@kwdef

src/dynamics/diagnostic_variables.jl

@@ -1,144 +1,108 @@
+# 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)
     # uω_coslat⁻¹_grid, vω_coslat⁻¹_grid = a_grid,b_grid    (all models)
     # 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,

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