Merge branch 'main' into ncc/run-example-in-docs

src/RingGrids/RingGrids.jl

@@ -1,7 +1,10 @@
 module RingGrids
 
-using DocStringExtensions
+# DOCUMENTATION AND VISUALISATION
+using  DocStringExtensions
+import UnicodePlots
 
+# NUMERICS
 import Statistics: mean
 import FastGaussQuadrature
 
@@ -55,14 +58,16 @@ export  interpolate,
         update_locator,
         update_locator!
 
+export  plot
+
 include("utility_functions.jl")
 
 include("grids_general.jl")
-include("show.jl")
 include("full_grids.jl")
 include("octahedral.jl")
 include("healpix.jl")
 include("octahealpix.jl")
 include("quadrature_weights.jl")
 include("interpolation.jl")
+include("show.jl")
 end

src/RingGrids/grids_general.jl

@@ -133,6 +133,7 @@ function get_latdlonds(Grid::Type{<:AbstractGrid},nlat_half::Integer)
 end
 
 get_latd(grid::Grid) where {Grid<:AbstractGrid} = get_latd(Grid,grid.nlat_half)
+get_lond(grid::Grid) where {Grid<:AbstractGrid} = get_lond(Grid,grid.nlat_half)
 
 function get_latd(Grid::Type{<:AbstractGrid},nlat_half::Integer)
     colat = get_colat(Grid,nlat_half)

src/RingGrids/interpolation.jl

@@ -174,7 +174,10 @@ function interpolator(  Aout::AbstractGrid,
 end
 
 ## FUNCTIONS
-interpolate(latd::Real,lond::Real,A::AbstractGrid) = interpolate([latd],[lond],A)
+function interpolate(latd::Real,lond::Real,A::AbstractGrid)
+    Ai = interpolate([latd],[lond],A)
+    return Ai[1]
+end
 
 function interpolate(   latds::Vector{NF},     # latitudes to interpolate onto (90˚N...-90˚N)
                         londs::Vector{NF},     # longitudes to interpolate into (0˚...360˚E)
@@ -407,15 +410,11 @@ function find_lon_indices(  λ::NF,      # longitude to find incides for (0˚...
 end
 
 """
-    N,S = average_on_poles( A::AbstractGrid,
-                            rings::Vector{<:UnitRange})
-
-N,S are the interpolated values of A onto the north/south pole, by averaging all values on the
-northern/southern-most rings respectively."""
-average_on_poles(A::AbstractGrid{NF},rings::Vector{<:UnitRange}) where NF = average_on_poles(NF,A,rings)
-
-function average_on_poles(  ::Type{NF},
-                            A::AbstractGrid,
+$(TYPEDSIGNATURES)
+Computes the average at the North and South pole from a given grid `A` and it's precomputed
+ring indices `rings`. The North pole average is an equally weighted average of all grid points
+on the northern-most ring. Similar for the South pole."""
+function average_on_poles(  A::AbstractGrid{NF},
                             rings::Vector{<:UnitRange{<:Integer}}
                             ) where {NF<:AbstractFloat}
 
@@ -424,6 +423,19 @@ function average_on_poles(  ::Type{NF},
     return convert(NF,A_northpole), convert(NF,A_southpole)
 end
 
+"""
+$(TYPEDSIGNATURES)
+Method for `A::Abstract{T<:Integer}` which rounds the averaged values
+to return the same number format `NF`."""
+function average_on_poles(  A::AbstractGrid{NF},
+                            rings::Vector{<:UnitRange{<:Integer}}
+                            ) where {NF<:Integer}
+
+    A_northpole = mean(view(A,rings[1]))    # average of all grid points around the north pole
+    A_southpole = mean(view(A,rings[end]))  # same for south pole
+    return round(NF,A_northpole), round(NF,A_southpole)
+end
+
 """
 $(TYPEDSIGNATURES)
 The bilinear average of a,b,c,d which are values at grid points
@@ -446,19 +458,20 @@ longitude/x-coordinate. See schematic:
               0...............1 # fraction of distance Δcd between c,d
 ```
 ^ fraction of distance Δy between a-b and c-d."""
-function anvil_average( a::NF,      # top left value
-                        b::NF,      # top right value
-                        c::NF,      # bottom left value
-                        d::NF,      # bottom right value
-                        Δab::Real,  # fraction of distance between a,b ∈ [0,1)
-                        Δcd::Real,  # fraction of distance between c,d ∈ [0,1)
-                        Δy::Real,   # fraction of distance between ab,cd ∈ [0,1)
-                        ) where NF
+function anvil_average(
+    a,      # top left value
+    b,      # top right value
+    c,      # bottom left value
+    d,      # bottom right value
+    Δab,    # fraction of distance between a,b ∈ [0,1)
+    Δcd,    # fraction of distance between c,d ∈ [0,1)
+    Δy,     # fraction of distance between ab,cd ∈ [0,1)
+    )
 
     # the type of the weights is ::Real, but the following is written such that
     # always NF (the type of the data values a,b,c,d) is returned
-    ab_average = a + (b-a)*convert(NF,Δab)      # a for Δab=0, b for Δab=1
-    cd_average = c + (d-c)*convert(NF,Δcd)      # c for Δab=0, b for Δab=1
-    abcd_average = ab_average + (cd_average-ab_average)*convert(NF,Δy)
+    ab_average = a + (b-a)*Δab      # a for Δab=0, b for Δab=1
+    cd_average = c + (d-c)*Δcd      # c for Δab=0, b for Δab=1
+    abcd_average = ab_average + (cd_average-ab_average)*Δy
     return abcd_average
 end

src/RingGrids/show.jl

@@ -4,4 +4,30 @@
 function Base.array_summary(io::IO, grid::AbstractGrid, inds::Tuple{Vararg{Base.OneTo}})
     print(io, Base.dims2string(length.(inds)), ", $(get_nlat(grid))-ring ")
     Base.showarg(io, grid, true)
+end
+
+function plot(A::AbstractGrid;title::String="$(get_nlat(A))-ring $(typeof(A))")
+    A_full = interpolate(full_grid(typeof(A)),A.nlat_half,A)
+    plot(A_full;title)
+end
+
+function plot(A::AbstractFullGrid;title::String="$(get_nlat(A))-ring $(typeof(A))")
+
+    A_matrix = Matrix(A)
+    nlon,nlat = size(A_matrix)
+    A_view = view(A_matrix,:,nlat:-1:1)
+
+    plot_kwargs = pairs((   xlabel="˚E",
+                            xfact=360/(nlon-1),
+                            ylabel="˚N",
+                            yfact=180/(nlat-1),
+                            yoffset=-90,
+                            title=title,
+                            colormap=:viridis,
+                            compact=true,
+                            colorbar=true,
+                            width=60,
+                            height=30))
+
+    UnicodePlots.heatmap(A_view';plot_kwargs...)
 end

src/SpeedyWeather.jl

@@ -1,7 +1,6 @@
 module SpeedyWeather
 
 # STRUCTURE
-import Parameters: @with_kw
 using DocStringExtensions
 
 # NUMERICS
@@ -67,7 +66,8 @@ export  LowerTriangularMatrix,
         OctahedralGaussianGrid,
         OctahedralClenshawGrid,
         HEALPixGrid,
-        OctaHEALPixGrid
+        OctaHEALPixGrid,
+        plot
 
 export  Leapfrog
 
@@ -98,6 +98,9 @@ export  NoBoundaryLayer,
 export  NoVerticalDiffusion,
         VerticalLaplacian
 
+# Large scale condensation
+export  SpeedyCondensation
+
 # EXPORT STRUCTS
 export  DynamicsConstants,
         SpectralTransform,
@@ -166,14 +169,14 @@ include("physics/thermodynamics.jl")
 include("physics/boundary_layer.jl")
 include("physics/temperature_relaxation.jl")
 include("physics/vertical_diffusion.jl")
+include("physics/large_scale_condensation.jl")
 include("physics/pretty_printing.jl")
 
 # MODELS
 include("dynamics/models.jl")
 
 # # PHYSICS
 # include("physics/convection.jl")
-# include("physics/large_scale_condensation.jl")
 # include("physics/longwave_radiation.jl")
 # include("physics/shortwave_radiation.jl")
 

src/abstract_types.jl

@@ -32,6 +32,7 @@ abstract type BoundaryLayerDrag{NF} <: AbstractParameterization{NF} end
 abstract type TemperatureRelaxation{NF} <: AbstractParameterization{NF} end
 abstract type VerticalDiffusion{NF} <: AbstractParameterization{NF} end
 abstract type AbstractThermodynamics{NF} <: AbstractParameterization{NF} end
+abstract type AbstractCondensation{NF} <: AbstractParameterization{NF} end
 
 # INPUT/OUTPUT
 abstract type AbstractFeedback end

src/dynamics/atmospheres.jl

@@ -4,7 +4,7 @@ Create a struct `EarthAtmosphere<:AbstractPlanet`, with the following physical/c
 characteristics. Note that `radius` is not part of it as this should be chosen
 in `SpectralGrid`. Keyword arguments are
 $(TYPEDFIELDS)"""
-@with_kw struct EarthAtmosphere <: AbstractAtmosphere
+Base.@kwdef struct EarthAtmosphere <: AbstractAtmosphere
     # ATMOSPHERE
     "molar mass of dry air [g/mol]"
     mol_mass_dry_air::Float64 = 28.9649
@@ -24,6 +24,9 @@ $(TYPEDFIELDS)"""
     "specific gas constant for water vapour [J/kg/K]"
     R_vapour::Float64 = 1000*R_gas/mol_mass_vapour
 
+    "water density [kg/m³]"
+    water_density::Float64 = 1000
+
     "latent heat of condensation [J/g] for consistency with specific humidity [g/Kg], also called alhc"
     latent_heat_condensation::Float64 = 2501
 
@@ -50,6 +53,9 @@ $(TYPEDFIELDS)"""
     "start of the stratosphere in sigma coordinates"
     σ_tropopause::Float64 = 0.2
 
+    "top of the planetary boundary layer in sigma coordinates"
+    σ_boundary_layer::Float64 = 0.95
+
     "scale height for pressure [km]"
     scale_height::Float64 = 7.5
 

src/dynamics/constants.jl

@@ -1,7 +1,7 @@
 """
 Struct holding constants needed at runtime for the dynamical core in number format NF.
 $(TYPEDFIELDS)"""
-@with_kw struct DynamicsConstants{NF<:AbstractFloat} <: AbstractDynamicsConstants{NF}
+Base.@kwdef struct DynamicsConstants{NF<:AbstractFloat} <: AbstractDynamicsConstants{NF}
     # PHYSICAL CONSTANTS
     "Radius of Planet [m]"
     radius::NF
@@ -31,6 +31,9 @@ $(TYPEDFIELDS)"""
     "= R_dry/cₚ, gas const for air over heat capacity"
     κ::NF
 
+    "water density [kg/m³]"
+    water_density::NF
+
     "coriolis frequency [1/s] (scaled by radius as is vorticity) = 2Ω*sin(lat)*radius"
     f_coriolis::Vector{NF}
 
@@ -63,7 +66,7 @@ function DynamicsConstants( spectral_grid::SpectralGrid,
                             geometry::Geometry)
 
     # PHYSICAL CONSTANTS
-    (;R_dry, R_vapour, lapse_rate, cₚ) = atmosphere
+    (;R_dry, R_vapour, lapse_rate, cₚ, water_density) = atmosphere
     (;ΔT_stratosphere, σ_tropopause, temp_ref) = atmosphere
     (;NF, radius) = spectral_grid
     (;rotation, gravity) = planet
@@ -102,7 +105,7 @@ function DynamicsConstants( spectral_grid::SpectralGrid,
 
     # This implies conversion to NF
     return DynamicsConstants{NF}(;  radius,rotation,gravity,layer_thickness,
-                                    R_dry,R_vapour,μ_virt_temp,cₚ,κ,
+                                    R_dry,R_vapour,μ_virt_temp,cₚ,κ,water_density,
                                     f_coriolis,
                                     σ_lnp_A,σ_lnp_B,
                                     Δp_geopot_half, Δp_geopot_full,

src/dynamics/horizontal_diffusion.jl

@@ -7,7 +7,7 @@ layers. Furthermore the power can be decreased above the `tapering_σ` to
 `power_stratosphere` (default 2). For Barotropic, ShallowWater,
 the default non-adaptive constant-time scale hyper diffusion is used. Options are
 $(TYPEDFIELDS)"""
-@with_kw struct HyperDiffusion{NF} <: HorizontalDiffusion{NF}
+Base.@kwdef struct HyperDiffusion{NF} <: HorizontalDiffusion{NF}
     # DIMENSIONS
     "spectral resolution"
     trunc::Int
@@ -260,13 +260,14 @@ function horizontal_diffusion!( progn::PrognosticLayerTimesteps,
     ∇²ⁿ_implicit = HD.∇²ⁿ_implicit[k]
 
     # Primitive equation models diffuse vor and divergence more selective/adaptive
-    (;vor,div,temp) = progn.timesteps[lf]
-    (;vor_tend,div_tend,temp_tend) = diagn.tendencies
+    (;vor,div,temp,humid) = progn.timesteps[lf]
+    (;vor_tend,div_tend,temp_tend,humid_tend) = diagn.tendencies
     horizontal_diffusion!(vor_tend,vor,∇²ⁿ,∇²ⁿ_implicit)
     horizontal_diffusion!(div_tend,div,∇²ⁿ,∇²ⁿ_implicit)
 
-    # but use the weaker normal diffusion for temperature
+    # but use the weaker normal diffusion for temperature, humidity
     ∇²ⁿ = HD.∇²ⁿ_2D
     ∇²ⁿ_implicit = HD.∇²ⁿ_2D_implicit
     horizontal_diffusion!(temp_tend,temp,∇²ⁿ,∇²ⁿ_implicit)
+    model isa PrimitiveWet && horizontal_diffusion!(humid_tend,humid,∇²ⁿ,∇²ⁿ_implicit)
 end

src/dynamics/implicit.jl

@@ -8,7 +8,7 @@ initialize!(I::NoImplicit,dt::Real,::DiagnosticVariables,::ModelSetup) = nothing
 Struct that holds various precomputed arrays for the semi-implicit correction to
 prevent gravity waves from amplifying in the shallow water model.
 $(TYPEDFIELDS)"""
-@with_kw struct ImplicitShallowWater{NF<:AbstractFloat} <: AbstractImplicit{NF}
+Base.@kwdef struct ImplicitShallowWater{NF<:AbstractFloat} <: AbstractImplicit{NF}
 
     # DIMENSIONS
     trunc::Int
@@ -130,7 +130,7 @@ end
 Struct that holds various precomputed arrays for the semi-implicit correction to
 prevent gravity waves from amplifying in the primitive equation model.
 $(TYPEDFIELDS)"""
-@with_kw struct ImplicitPrimitiveEq{NF<:AbstractFloat} <: AbstractImplicit{NF}
+Base.@kwdef struct ImplicitPrimitiveEq{NF<:AbstractFloat} <: AbstractImplicit{NF}
 
     # DIMENSIONS
     "spectral resolution"