diff --git a/.gitignore b/.gitignore index 104211769..71244b451 100644 --- a/.gitignore +++ b/.gitignore @@ -30,4 +30,5 @@ Manifest.toml # no model output folders run????/ -run-*/ \ No newline at end of file +run-*/ +run_*/ \ No newline at end of file diff --git a/Project.toml b/Project.toml index c961ae645..52131f95f 100644 --- a/Project.toml +++ b/Project.toml @@ -21,6 +21,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NetCDF = "30363a11-5582-574a-97bb-aa9a979735b9" +Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a" Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca" diff --git a/README.md b/README.md index 39ce585bb..ea5302c9d 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ the shallow water equations. The arguments for `run_speedy` are described in To run the primitive equation dry core you can for example do ```julia -julia> run_speedy(Float32, PrimitiveDryCore, physics=true, diffusion=HyperDiffusion(power=2)) +julia> run_speedy(Float32, PrimitiveDry, physics=true, diffusion=HyperDiffusion(power=2)) Weather is speedy: 100%|███████████████████████████████████| Time: 0:00:03 (753.71 years/day) ``` diff --git a/docs/src/how_to_run_speedy.md b/docs/src/how_to_run_speedy.md index 2174764b8..ca889f1cf 100644 --- a/docs/src/how_to_run_speedy.md +++ b/docs/src/how_to_run_speedy.md @@ -14,11 +14,11 @@ at higher resolution (`trunc`, the triangular spectral truncation), slow down th of the Earth (`rotation` in ``s^{-1}``), and create some netCDF ouput, do ```julia -run_speedy(Float64,PrimitiveDryCore,trunc=42,planet=Earth(rotation=1e-5),output=true) +run_speedy(Float64,PrimitiveDry,trunc=42,planet=Earth(rotation=1e-5),output=true) ``` If provided, the number format has to be the first argument, the model (`Barotropic`, `ShallowWater`, -`PrimitiveDryCore`, `PrimitiveWetCore` are available) the second, and all other arguments are keyword +`PrimitiveDry`, `PrimitiveWet` are available) the second, and all other arguments are keyword arguments. ## The `run_speedy` interface diff --git a/src/SpeedyTransforms/SpeedyTransforms.jl b/src/SpeedyTransforms/SpeedyTransforms.jl index bd41cb197..dee32081d 100644 --- a/src/SpeedyTransforms/SpeedyTransforms.jl +++ b/src/SpeedyTransforms/SpeedyTransforms.jl @@ -43,5 +43,6 @@ include("spectral_transform.jl") include("spectral_gradients.jl") include("spectral_truncation.jl") include("spectrum.jl") +include("show.jl") end \ No newline at end of file diff --git a/src/SpeedyTransforms/aliasing.jl b/src/SpeedyTransforms/aliasing.jl index 42cc33d46..e67040a21 100644 --- a/src/SpeedyTransforms/aliasing.jl +++ b/src/SpeedyTransforms/aliasing.jl @@ -14,7 +14,7 @@ end """ m = roundup_fft(n::Int; - small_primes::Vector{Int}=[2,3]) + small_primes::Vector{Int}=[2,3,5]) Returns an integer `m >= n` with only small prime factors 2, 3 (default, others can be specified with the keyword argument `small_primes`) to obtain an efficiently fourier-transformable number of diff --git a/src/SpeedyTransforms/show.jl b/src/SpeedyTransforms/show.jl new file mode 100644 index 000000000..033916cd1 --- /dev/null +++ b/src/SpeedyTransforms/show.jl @@ -0,0 +1,7 @@ +function Base.show(io::IO,S::SpectralTransform{NF}) where NF + (;lmax,Grid,nlat_half) = S + println(io,"$(typeof(S))(") + println(io," Spectral: T$lmax LowerTriangularMatrix{Complex{$NF}}") + println(io," Grid: $(RingGrids.get_nlat(Grid,nlat_half))-ring $Grid{$NF}") + print(io," Legendre: recompute polynomials $(S.recompute_legendre))") +end diff --git a/src/SpeedyTransforms/spectral_transform.jl b/src/SpeedyTransforms/spectral_transform.jl index 0ecd5895a..c8ecfa25d 100644 --- a/src/SpeedyTransforms/spectral_transform.jl +++ b/src/SpeedyTransforms/spectral_transform.jl @@ -65,7 +65,7 @@ struct SpectralTransform{NF<:AbstractFloat} end """ - S = SpectralTransform(NF,Grid,trunc,recompute_legendre) + S = SpectralTransform(NF,Grid,trunc) Generator function for a SpectralTransform struct. With `NF` the number format, `Grid` the grid type `<:AbstractGrid` and spectral truncation `trunc` this function sets up @@ -73,9 +73,9 @@ necessary constants for the spetral transform. Also plans the Fourier transforms and preallocates the Legendre polynomials (if recompute_legendre == false) and quadrature weights.""" function SpectralTransform( ::Type{NF}, # Number format NF Grid::Type{<:AbstractGrid}, # type of spatial grid used - trunc::Int, # Spectral truncation - recompute_legendre::Bool; # re or precompute legendre polynomials? - legendre_shortcut::Symbol=:linear, # shorten Legendre loop over order m + trunc::Int; # Spectral truncation + recompute_legendre::Bool = true, # re or precompute legendre polynomials? + legendre_shortcut::Symbol = :linear, # shorten Legendre loop over order m dealiasing::Real=DEFAULT_DEALIASING ) where NF @@ -264,12 +264,12 @@ end Generator function for a `SpectralTransform` struct based on the size of the spectral coefficients `alms` and the grid `Grid`. Recomputes the Legendre polynomials by default.""" function SpectralTransform( alms::AbstractMatrix{Complex{NF}}; # spectral coefficients - recompute_legendre::Bool=true, # saves memory - Grid::Type{<:AbstractGrid}=DEFAULT_GRID, + recompute_legendre::Bool = true, # saves memory + Grid::Type{<:AbstractGrid} = DEFAULT_GRID, ) where NF # number format NF _, mmax = size(alms) .- 1 # -1 for 0-based degree l, order m - return SpectralTransform(NF,Grid,mmax,recompute_legendre) + return SpectralTransform(NF,Grid,mmax;recompute_legendre) end """ @@ -284,7 +284,7 @@ function SpectralTransform( map::AbstractGrid{NF}; # gridded field Grid = typeof(map) trunc = get_truncation(map) - return SpectralTransform(NF,Grid,trunc,recompute_legendre) + return SpectralTransform(NF,Grid,trunc;recompute_legendre) end """ @@ -549,12 +549,12 @@ Spectral transform (spectral to grid space) from spherical coefficients `alms` t field `map`. Based on the size of `alms` the grid type `grid`, the spatial resolution is retrieved based on the truncation defined for `grid`. SpectralTransform struct `S` is allocated to execute `gridded(alms,S)`.""" function gridded( alms::AbstractMatrix{T}; # spectral coefficients - recompute_legendre::Bool=true, # saves memory - Grid::Type{<:AbstractGrid}=DEFAULT_GRID, + recompute_legendre::Bool = true, # saves memory + Grid::Type{<:AbstractGrid} = DEFAULT_GRID, ) where {NF,T<:Complex{NF}} # number format NF _, mmax = size(alms) .- 1 # -1 for 0-based degree l, order m - S = SpectralTransform(NF,Grid,mmax,recompute_legendre) + S = SpectralTransform(NF,Grid,mmax;recompute_legendre) return gridded(alms,S) end @@ -595,12 +595,12 @@ end Converts `map` to `Grid(map)` to execute `spectral(map::AbstractGrid;kwargs...)`.""" function spectral( map::AbstractGrid{NF}; # gridded field - recompute_legendre::Bool=true, # saves memory + recompute_legendre::Bool = true,# saves memory ) where NF # number format NF Grid = typeof(map) trunc = get_truncation(map.nlat_half) - S = SpectralTransform(NF,Grid,trunc,recompute_legendre) + S = SpectralTransform(NF,Grid,trunc;recompute_legendre) return spectral(map,S) end diff --git a/src/SpeedyWeather.jl b/src/SpeedyWeather.jl index 1ddc70550..d68e33222 100644 --- a/src/SpeedyWeather.jl +++ b/src/SpeedyWeather.jl @@ -1,6 +1,7 @@ module SpeedyWeather # STRUCTURE +import Parameters: @with_kw using DocStringExtensions # NUMERICS @@ -27,19 +28,34 @@ import BitInformation: round, round! import UnicodePlots import ProgressMeter -# EXPORT MAIN INTERFACE TO SPEEDY +# EXPORT MONOLITHIC INTERFACE TO SPEEDY export run_speedy, run_speedy!, - initialize_speedy + initialize_speedy, + initialize!, + run! + +export NoVerticalCoordinates, + SigmaCoordinates, + SigmaPressureCoordinates # EXPORT MODELS -export Barotropic, +export Barotropic, # abstract ShallowWater, PrimitiveEquation, - PrimitiveDryCore, - PrimitiveWetCore + PrimitiveDry, + PrimitiveWet + +export Model, + BarotropicModel, # concrete + ShallowWaterModel, + PrimitiveDryModel, + PrimitiveWetModel # EXPORT GRIDS +export SpectralGrid, + Geometry + export LowerTriangularMatrix, FullClenshawGrid, FullGaussianGrid, @@ -50,6 +66,8 @@ export LowerTriangularMatrix, HEALPixGrid, OctaHEALPixGrid +export Leapfrog + # EXPORT OROGRAPHIES export NoOrography, EarthOrography, @@ -78,10 +96,7 @@ export NoVerticalDiffusion, VerticalLaplacian # EXPORT STRUCTS -export Parameters, - DynamicsConstants, - ParameterizationConstants, - Geometry, +export DynamicsConstants, SpectralTransform, Boundaries, PrognosticVariables, @@ -89,12 +104,14 @@ export Parameters, ColumnVariables # EXPORT SPECTRAL FUNCTIONS -export spectral, +export SpectralTransform, + spectral, gridded, spectral_truncation + +export OutputWriter, Feedback include("utility_functions.jl") -include("abstract_types.jl") # LowerTriangularMatrices for spherical harmonics export LowerTriangularMatrices @@ -110,49 +127,56 @@ using .RingGrids export SpeedyTransforms include("SpeedyTransforms/SpeedyTransforms.jl") using .SpeedyTransforms - -include("gpu.jl") # defines utility for GPU / KernelAbstractions -include("default_parameters.jl") # defines Parameters - -# SOME DEFINITIONS FIRST -include("dynamics/constants.jl") # defines DynamicsConstants -include("physics/constants.jl") # defines ParameterizationConstants -include("physics/define_column.jl") # define ColumnVariables - -# DYNAMICS -include("dynamics/geometry.jl") # defines Geometry -include("dynamics/boundaries.jl") # defines Boundaries -include("dynamics/define_diffusion.jl") # defines HorizontalDiffusion -include("dynamics/define_implicit.jl") # defines ImplicitShallowWater, ImplicitPrimitiveEq # defines GenLogisticCoefs -include("dynamics/planets.jl") # defines Earth -include("dynamics/models.jl") # defines ModelSetups -include("dynamics/prognostic_variables.jl") # defines PrognosticVariables -include("dynamics/diagnostic_variables.jl") # defines DiagnosticVariables + +# Utility for GPU / KernelAbstractions +include("gpu.jl") + +# GEOMETRY CONSTANTS ETC +include("abstract_types.jl") +include("dynamics/vertical_coordinates.jl") +include("dynamics/spectral_grid.jl") +include("dynamics/planets.jl") +include("dynamics/atmospheres.jl") +include("dynamics/constants.jl") +include("dynamics/orography.jl") + +# VARIABLES +include("dynamics/prognostic_variables.jl") +include("physics/define_column.jl") +include("dynamics/diagnostic_variables.jl") + +# MODEL COMPONENTS +include("dynamics/time_integration.jl") +include("dynamics/forcing.jl") +include("dynamics/geopotential.jl") include("dynamics/initial_conditions.jl") +include("dynamics/horizontal_diffusion.jl") +include("dynamics/implicit.jl") include("dynamics/scaling.jl") -include("dynamics/geopotential.jl") -include("dynamics/tendencies_dynamics.jl") include("dynamics/tendencies.jl") -include("dynamics/implicit.jl") -include("dynamics/diffusion.jl") -include("dynamics/time_integration.jl") +include("dynamics/tendencies_dynamics.jl") -# PHYSICS +# PARAMETERIZATIONS +include("physics/tendencies.jl") include("physics/column_variables.jl") include("physics/thermodynamics.jl") -include("physics/tendencies.jl") -include("physics/convection.jl") -include("physics/large_scale_condensation.jl") -include("physics/longwave_radiation.jl") -include("physics/shortwave_radiation.jl") include("physics/boundary_layer.jl") include("physics/temperature_relaxation.jl") include("physics/vertical_diffusion.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") # OUTPUT include("output/output.jl") # defines Output include("output/feedback.jl") # defines Feedback -include("output/pretty_printing.jl") # INTERFACE include("run_speedy.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index a85ea75c4..55a4208b2 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -3,45 +3,41 @@ abstract type ModelSetup end abstract type Barotropic <: ModelSetup end abstract type ShallowWater <: ModelSetup end abstract type PrimitiveEquation <: ModelSetup end -abstract type PrimitiveDryCore <: PrimitiveEquation end -abstract type PrimitiveWetCore <: PrimitiveEquation end +abstract type PrimitiveDry <: PrimitiveEquation end +abstract type PrimitiveWet <: PrimitiveEquation end -# GEOMETRY -abstract type AbstractGeometry{NF} end - -# PARAMETERS (to be chosen by user) -abstract type AbstractParameters{M} end +abstract type AbstractPlanet end +abstract type AbstractAtmosphere end -# COEFFICIENTS (bundled parameters) -abstract type Coefficients end +# GEOMETRY, GRID +abstract type AbstractGeometry{NF} end +abstract type VerticalCoordinates end -# CONSTANTS (to be defined from parameters & coefficients, -# face the dynamical core/parameterizations and not the user) -abstract type AbstractParameterizationConstants{NF} end +# CONSTANTS (face the dynamical core and not the user) abstract type AbstractDynamicsConstants{NF} end # INITIAL CONDITIONS AND OROGRAPHY/BOUNDARIES -abstract type InitialConditions end # subtypes defined in initial_conditions.jl -abstract type AbstractOrography end # subtypes defined in boundaries.jl -abstract type AbstractBoundaries{NF} end +abstract type InitialConditions end +abstract type AbstractOrography{NF,Grid} end # ATMOSPHERIC COLUMN FOR PARAMETERIZATIONS abstract type AbstractColumnVariables{NF} end +# FORCING (Barotropic and ShallowWaterModel) +abstract type AbstractForcing{NF} end + # PARAMETERIZATIONS -abstract type BoundaryLayer{NF} end -abstract type TemperatureRelaxation{NF} end -abstract type VerticalDiffusion{NF} end +abstract type AbstractParameterization{NF} end +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 # INPUT/OUTPUT abstract type AbstractFeedback end -abstract type AbstractOutput end - -# IMPLICIT PRECOMPUTED TERMS -abstract type AbstractImplicit{NF} end - -# PLANETS -abstract type Planet end +abstract type AbstractOutputWriter end # NUMERICS -abstract type DiffusionParameters end \ No newline at end of file +abstract type HorizontalDiffusion{NF} end +abstract type AbstractImplicit{NF} end +abstract type TimeStepper{NF} end \ No newline at end of file diff --git a/src/default_parameters.jl b/src/default_parameters.jl deleted file mode 100644 index 771fe3458..000000000 --- a/src/default_parameters.jl +++ /dev/null @@ -1,378 +0,0 @@ -const DEFAULT_NF = Float32 # number format -const DEFAULT_MODEL = PrimitiveDryCore # abstract model type - -""" - P = Parameters{M<:ModelSetup}(kwargs...) <: AbstractParameters{M} - -A struct to hold all model parameters that may be changed by the user. -The struct uses keywords such that default values can be changed at creation. -The default values of the keywords define the default model setup. - -$(TYPEDFIELDS) -""" -Base.@kwdef struct Parameters{Model<:ModelSetup} <: AbstractParameters{Model} - - "number format" - NF::DataType = DEFAULT_NF - - - # RESOLUTION AND GRID - - "spectral truncation" - trunc::Int = 31 - - "grid in use" - Grid::Type{<:AbstractGrid} = OctahedralGaussianGrid - - "dealiasing factor, 1=linear, 2=quadratic, 3=cubic grid" - dealiasing::Float64 = 2 - - - # PLANET'S PROPERTIES - - "planet" - planet::Planet = Earth() - - - # ATMOSPHERE - - "molar mass of dry air [g/mol]" - mol_mass_dry_air = 28.9649 - - "molar mass of water vapour [g/mol]" - mol_mass_vapour = 18.0153 - - "specific heat at constant pressure [J/K/kg]" - cₚ::Float64 = 1004 - - "universal gas constant [J/K/mol]" - R_gas::Float64 = 8.3145 - - "specific gas constant for dry air [J/kg/K]" - R_dry::Float64 = 1000*R_gas/mol_mass_dry_air - - "specific gas constant for water vapour [J/kg/K]" - R_vapour::Float64 = 1000*R_gas/mol_mass_vapour - - "latent heat of condensation [J/g] for consistency with specific humidity [g/Kg]" - alhc::Float64 = 2501 - - "latent heat of sublimation [?]" - alhs::Float64 = 2801 - - "stefan-Boltzmann constant [W/m²/K⁴]" - sbc::Float64 = 5.67e-8 - - - # STANDARD ATMOSPHERE (reference values) - - "moist adiabatic temperature lapse rate ``-dT/dz`` [K/km]" - lapse_rate::Float64 = 5 - - "absolute temperature at surface ``z=0`` [K]" - temp_ref::Float64 = 288 - - "absolute temperature in stratosphere [K]" - temp_top::Float64 = 216 - - "for stratospheric lapse rate [K] after Jablonowski" - ΔT_stratosphere::Float64 = 4.8e5 - - "scale height for pressure [km]" - scale_height::Float64 = 7.5 - - "surface pressure [hPa]" - pres_ref::Float64 = 1000 - - "scale height for specific humidity [km]" - scale_height_humid::Float64 = 2.5 - - "relative humidity of near-surface air [1]" - relhumid_ref::Float64 = 0.7 - - "saturation water vapour pressure [Pa]" - water_pres_ref::Float64 = 17 - - "layer thickness for the shallow water model [km]" - layer_thickness::Float64 = 8.5 - - - # VERTICAL COORDINATES - - "vertical coordinates of the nlev vertical levels, defined by a generalised logistic function, interpolating ECMWF's L31 configuration" - GLcoefs::Coefficients = GenLogisticCoefs() - - "σ coordinate where the tropopause starts" - σ_tropopause::Float64 = 0.2 - - "only used if set manually, otherwise empty" - σ_levels_half::Vector{Float64} = [] - - "number of vertical levels" - nlev::Int = nlev_default(Model, σ_levels_half) - - - # DIFFUSION AND DRAG - - "horizontal (hyper)-diffusion" - diffusion::DiffusionParameters = HyperDiffusion() - - "vertical diffusion" - vertical_diffusion::VerticalDiffusion = NoVerticalDiffusion() - - "static energy diffusion" - static_energy_diffusion::VerticalDiffusion = StaticEnergyDiffusion() - - - # FORCING - - "turn on interface relaxation for shallow water?" - interface_relaxation::Bool = false - - "time scale [hrs] of interface relaxation" - interface_relax_time::Float64 = 96 - - "Amplitude [m] of interface relaxation" - interface_relax_amplitude::Float64 = 300 - - - # PARAMETRIZATIONS - - "en/disables the physics parameterizations" - physics::Bool = true - - - "Compute shortwave radiation every `n` steps" - n_shortwave::Int = 3 - - "Turn on SPPT?" - sppt_on::Bool = false - - "For computing saturation vapour pressure" - magnus_coefs::Coefficients = MagnusCoefs{NF}() - - # Large-Scale Condensation (from table B10) - "Index of atmospheric level at which large-scale condensation begins" - k_lsc::Int = 2 - - "Relative humidity threshold for boundary layer" - RH_thresh_pbl_lsc::Float64 = 0.95 - - "Vertical range of relative humidity threshold" - RH_thresh_range_lsc::Float64 = 0.1 - - "Maximum relative humidity threshold" - RH_thresh_max_lsc::Float64 = 0.9 - - "Relaxation time for humidity (hours)" - humid_relax_time_lsc::Float64 = 4.0 - - # Convection - "Minimum (normalised) surface pressure for the occurrence of convection" - pres_thresh_cnv::Float64 = 0.8 - - "Relative humidity threshold for convection in PBL" - RH_thresh_pbl_cnv::Float64 = 0.9 - - "Relative humidity threshold for convection in the troposphere" - RH_thresh_trop_cnv::Float64 = 0.7 - - "Relaxation time for PBL humidity (hours)" - humid_relax_time_cnv::Float64 = 6.0 - - "Maximum entrainment as a fraction of cloud-base mass flux" - max_entrainment::Float64 = 0.5 - - "Ratio between secondary and primary mass flux at cloud-base" - ratio_secondary_mass_flux::Float64 = 0.8 - - # Longwave radiation - "Number of bands used to compute `fband`" - nband::Int = 4 - - # Radiation - "radiation coefficients" - radiation_coefs::Coefficients = RadiationCoefs{NF}() - - # BOUNDARY LAYER - "boundary layer drag" - boundary_layer::BoundaryLayer{Float64} = LinearDrag() - - "temperature relaxation" - temperature_relaxation::TemperatureRelaxation{Float64} = HeldSuarez() - - - # TIME STEPPING - - "time at which the integration starts" - startdate::DateTime = DateTime(2000,1,1) - - "number of days to integrate for" - n_days::Float64 = 10 - - "time step in minutes for T31, scale linearly to `trunc`" - Δt_at_T31::Float64 = 30 - - - # NUMERICS - - "Robert (1966) time filter coefficeint to suppress comput. mode" - robert_filter::Float64 = 0.05 - - "William's time filter (Amezcua 2011) coefficient for 3rd order acc" - williams_filter::Float64 = 0.53 - - "coefficient for semi-implicit computations to filter gravity waves" - implicit_α::Float64 = 1 - - "recalculate implicit operators on temperature profile every n time steps" - recalculate_implicit::Int = 100 - - # LEGENDRE TRANSFORM AND POLYNOMIALS - - "recomputation is slower but requires less memory" - recompute_legendre::Bool = false - - "which format to use to calculate the Legendre polynomials" - legendre_NF::DataType = Float64 - - "`:linear`, `:quadratic`, `:cubic`, `:lincub_coslat`, `:linquad_coslat²`" - legendre_shortcut::Symbol = :linear - - - # BOUNDARY FILES - - "package location is default" - boundary_path::String = "" - - "orography" - orography::AbstractOrography = EarthOrography() - - "scale orography by a factor" - orography_scale::Float64 = 1 - - "path of orography" - orography_path::String = boundary_path - - "filename of orography" - orography_file::String = "orography_F512.nc" - - - # INITIAL CONDITIONS - - "random seed for the global random number generator" - seed::Int = 123456789 - - "initial conditions" - initial_conditions::InitialConditions = initial_conditions_default(Model) - - "calculate the initial surface pressure from orography" - pressure_on_orography::Bool = false - - - # OUTPUT - - "print dialog for feedback" - verbose::Bool = true - - "print debug info, NaR detection" - debug::Bool = true - - "Store data in netCDF?" - output::Bool = false - - "output time step [hours]" - output_dt::Float64 = 6 - - "path to output folder" - output_path::String = pwd() - - "name of the output folder, defaults to 4-digit number counting up from `run-0001`" - run_id::Union{String,Int} = get_run_id(output, output_path) - - "name of the output netcdf file" - output_filename::String = "output.nc" - - "variables to output: `:u`, `:v`, `:vor`, `:div`, `:temp`, `:humid`" - output_vars::Vector{Symbol} = output_vars_default(Model) - - "compression level; 1=low but fast, 9=high but slow" - compression_level::Int = 3 - - "mantissa bits to keep for every variable" - keepbits::NamedTuple = default_keepbits() - - "SpeedyWeather.jl version number" - version::VersionNumber = pkgversion(SpeedyWeather) - - - # OUTPUT GRID - - "number format used for output" - output_NF::DataType = Float32 - - "0 = reuse nlat_half from dynamical core" - output_nlat_half::Int = 0 - - "output grid" - output_Grid::Type{<:AbstractFullGrid} = RingGrids.full_grid(Grid) - - "output interpolator" - output_Interpolator::Type{<:AbstractInterpolator} = DEFAULT_INTERPOLATOR - - "if true sort gridpoints into a matrix" - output_matrix::Bool = false - - "rotation of output quadrant" - output_quadrant_rotation::NTuple{4,Int} = (0,1,2,3) - - "matrix of output quadrant" - output_matrix_quadrant::NTuple{4,Tuple{Int,Int}} = ((2,2),(1,2),(1,1),(2,1)) - - "missing value to be used in netcdf output" - missing_value::Float64 = NaN - - - # RESTART - - "also write restart file if output==true?" - write_restart::Bool = output - - "path for restart file" - restart_path::String = output_path - - "`run_id` of restart file in `run-????/restart.jld2`" - restart_id::Union{String,Int} = 1 -end - -Parameters(;kwargs...) = Parameters{default_concrete_model(DEFAULT_MODEL)}(;kwargs...) - -""" - nlev = nlev_default(Model::Type{<:ModelSetup}, σ_levels_half::AbstractVector) - -Number of vertical levels chosen either automatically based on `Model`, -or from the length of `σ_levels_half` if not a 0-length vector -(default if not specified parameter). -""" -function nlev_default(Model::Type{<:ModelSetup}, σ_levels_half::AbstractVector) - if length(σ_levels_half) == 0 # choose nlev automatically - Model <: Barotropic && return 1 - Model <: ShallowWater && return 1 - Model <: PrimitiveEquation && return 8 - else # use manually set levels - return length(σ_levels_half) - 1 - end -end - -# default variables to output by model -output_vars_default(::Type{<:Barotropic}) = [:vor,:u] -output_vars_default(::Type{<:ShallowWater}) = [:vor,:u] -output_vars_default(::Type{<:PrimitiveDryCore}) = [:vor,:u,:temp,:pres] -output_vars_default(::Type{<:PrimitiveWetCore}) = [:vor,:u,:temp,:humid,:pres] - -default_keepbits() = (u=7,v=7,vor=5,div=5,temp=10,pres=12,humid=7) - -# default initial conditions by model -initial_conditions_default(::Type{<:Barotropic}) = StartWithVorticity() -initial_conditions_default(::Type{<:ShallowWater}) = ZonalJet() -initial_conditions_default(::Type{<:PrimitiveEquation}) = ZonalWind() \ No newline at end of file diff --git a/src/dynamics/atmospheres.jl b/src/dynamics/atmospheres.jl new file mode 100644 index 000000000..5c543ef3c --- /dev/null +++ b/src/dynamics/atmospheres.jl @@ -0,0 +1,78 @@ +""" +$(TYPEDSIGNATURES) +Create a struct `EarthAtmosphere<:AbstractPlanet`, with the following physical/chemical +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 + # ATMOSPHERE + "molar mass of dry air [g/mol]" + mol_mass_dry_air::Float64 = 28.9649 + + "molar mass of water vapour [g/mol]" + mol_mass_vapour::Float64 = 18.0153 + + "specific heat at constant pressure [J/K/kg]" + cₚ::Float64 = 1004 + + "universal gas constant [J/K/mol]" + R_gas::Float64 = 8.3145 + + "specific gas constant for dry air [J/kg/K]" + R_dry::Float64 = 1000*R_gas/mol_mass_dry_air + + "specific gas constant for water vapour [J/kg/K]" + R_vapour::Float64 = 1000*R_gas/mol_mass_vapour + + "latent heat of condensation [J/g] for consistency with specific humidity [g/Kg], also called alhc" + latent_heat_condensation::Float64 = 2501 + + "latent heat of sublimation [J/g], also called alhs" + latent_heat_sublimation::Float64 = 2801 + + "stefan-Boltzmann constant [W/m²/K⁴]" + stefan_boltzmann::Float64 = 5.67e-8 + + + # STANDARD ATMOSPHERE (reference values) + "moist adiabatic temperature lapse rate ``-dT/dz`` [K/km]" + lapse_rate::Float64 = 5 + + "absolute temperature at surface ``z=0`` [K]" + temp_ref::Float64 = 288 + + "absolute temperature in stratosphere [K]" + temp_top::Float64 = 216 + + "for stratospheric lapse rate [K] after Jablonowski" + ΔT_stratosphere::Float64 = 4.8e5 + + "start of the stratosphere in sigma coordinates" + σ_tropopause::Float64 = 0.2 + + "scale height for pressure [km]" + scale_height::Float64 = 7.5 + + "surface pressure [hPa]" + pres_ref::Float64 = 1000 + + "scale height for specific humidity [km]" + scale_height_humid::Float64 = 2.5 + + "relative humidity of near-surface air [1]" + relhumid_ref::Float64 = 0.7 + + "saturation water vapour pressure [Pa]" + water_pres_ref::Float64 = 17 + + "layer thickness for the shallow water model [km]" + layer_thickness::Float64 = 8.5 +end + +function Base.show(io::IO,atm::AbstractAtmosphere) + print(io,"$(typeof(atm)):") + for key in propertynames(atm) + val = getfield(atm,key) + print(io,"\n $key::$(typeof(val)) = $val") + end +end \ No newline at end of file diff --git a/src/dynamics/boundaries.jl b/src/dynamics/boundaries.jl deleted file mode 100644 index d0c1fdc22..000000000 --- a/src/dynamics/boundaries.jl +++ /dev/null @@ -1,142 +0,0 @@ -"""Concrete struct that holds orography in grid-point space and surface geopotential in -spectral space.""" -struct Orography{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} - orography::Grid # height [m] - geopot_surf::LowerTriangularMatrix{Complex{NF}} # surface geopotential, height*gravity [m²/s²] -end - -"""Zonal ridge orography after Jablonowski and Williamson 2006.""" -Base.@kwdef struct ZonalRidge <: AbstractOrography - η₀::Float32 = 0.252 # conversion from σ to Jablonowski's ηᵥ-coordinates - u₀::Float32 = 35 # max amplitude of zonal wind [m/s] that scales orography height -end - -# empty structs (no parameters needed) for other orographies -Base.@kwdef struct EarthOrography <: AbstractOrography - smoothing::Bool = true # smooth the orography field? - smoothing_power::Float64 = 1.0 # power of Laplacian for smoothing - smoothing_strength::Float64 = 0.1 # highest degree l is multiplied by - smoothing_truncation::Int = 85 # resolution of orography -end - -struct NoOrography <: AbstractOrography end - -function Base.zeros(::Type{Orography},S::SpectralTransform{NF}) where NF - (;Grid, nlat_half, lmax, mmax) = S - orography = zeros(Grid{NF},nlat_half) - geopot_surf = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) - return Orography(orography,geopot_surf) -end - -struct Boundaries{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} <: AbstractBoundaries{NF} - orography::Orography{NF,Grid} - # landsea_mask::AbstractLandSeaMask{NF} - # albedo::AbstractAlbedo{NF} -end - -function Boundaries(P::Parameters, - S::SpectralTransform{NF}, - G::Geometry{NF}) where NF - - orography = zeros(Orography,S) # allocate orography arrays - initialize_orography!(orography,P.orography,P,S,G) # fill them with data - scale_orography!(orography,P) # make whatever mountains bigger/smaller - return Boundaries{NF,S.Grid{NF}}(orography) -end - -# recalculate spectral transform or geometry if not provided -Boundaries(P::Parameters,S::SpectralTransform) = Boundaries(P,S,Geometry(P)) -Boundaries(P::Parameters) = Boundaries(P,SpectralTransform(P)) - -function initialize_orography!( ::Orography, - ::NoOrography, - ::Parameters{<:ModelSetup}, - args...) - return nothing -end - -function initialize_orography!( ::Orography, - ::AbstractOrography, - ::Parameters{<:Barotropic}, - args...) - return nothing -end - -function initialize_orography!( orog::Orography, - EO::EarthOrography, - P::Parameters{M}, - S::SpectralTransform, - G::Geometry) where {M<:Union{ShallowWater,PrimitiveEquation}} - - (;orography, geopot_surf) = orog - (;orography_path, orography_file) = P - (;gravity) = P.planet - (;lmax, mmax) = S - - # LOAD NETCDF FILE - if orography_path == "" - path = joinpath(@__DIR__,"../../input_data",orography_file) - else - path = joinpath(orography_path,orography_file) - end - ncfile = NetCDF.open(path) - - orography_highres = ncfile.vars["orog"][:,:] # height [m] - - # Interpolate/coarsen to desired resolution - #TODO also read lat,lon from file and flip array in case it's not as expected - recompute_legendre = true # don't allocate large arrays as spectral transform is not reused - Grid = FullGaussianGrid # grid the orography file comes with - orography_spec = spectral(orography_highres;Grid,recompute_legendre) - - copyto!(geopot_surf,orography_spec) # truncates to the size of geopot_surf, no *gravity yet - if EO.smoothing # smooth orography in spectral space? - SpeedyTransforms.spectral_smoothing!(geopot_surf,EO.smoothing_strength, - power=EO.smoothing_power, - truncation=EO.smoothing_truncation) - end - - gridded!(orography,geopot_surf,S) # to grid-point space - geopot_surf .*= gravity # turn orography into surface geopotential - spectral_truncation!(geopot_surf) # set the lmax+1 harmonics to zero -end - -function initialize_orography!( orog::Orography, - coefs::ZonalRidge, - P::Parameters{M}, - S::SpectralTransform, - G::Geometry) where {M<:Union{ShallowWater,PrimitiveEquation}} - - (;gravity, rotation, radius) = P.planet - (;lmax, mmax) = S - - (;orography, geopot_surf) = orog - (;η₀, u₀) = coefs - - ηᵥ = (1-η₀)*π/2 # ηᵥ-coordinate of the surface [1] - A = u₀*cos(ηᵥ)^(3/2) # amplitude [m/s] - RΩ = radius*rotation # [m/s] - g⁻¹ = inv(gravity) # inverse gravity [s²/m] - φ = G.latds # latitude for each grid point [˚N] - - for ij in eachindex(φ,orography.data) - sinφ = sind(φ[ij]) - cosφ = cosd(φ[ij]) - - # Jablonowski & Williamson, 2006, eq. (7) - orography[ij] = g⁻¹*A*(A*(-2*sinφ^6*(cosφ^2 + 1/3) + 10/63) + (8/5*cosφ^3*(sinφ^2 + 2/3) - π/4)*RΩ) - end - - spectral!(geopot_surf,orography,S) # to grid-point space - geopot_surf .*= gravity # turn orography into surface geopotential - spectral_truncation!(geopot_surf) # set the lmax+1 harmonics to zero -end - -function scale_orography!( orog::Orography, - P::Parameters) - - (;orography, geopot_surf) = orog - orography .*= P.orography_scale - geopot_surf .*= P.orography_scale - return nothing -end \ No newline at end of file diff --git a/src/dynamics/constants.jl b/src/dynamics/constants.jl index 3d2b9b727..80da11c68 100644 --- a/src/dynamics/constants.jl +++ b/src/dynamics/constants.jl @@ -1,104 +1,110 @@ """ -Struct holding the parameters needed at runtime in number format NF. -""" -Base.@kwdef struct DynamicsConstants{NF<:AbstractFloat} <: AbstractDynamicsConstants{NF} - +Struct holding constants needed at runtime for the dynamical core in number format NF. +$(TYPEDFIELDS)""" +@with_kw struct DynamicsConstants{NF<:AbstractFloat} <: AbstractDynamicsConstants{NF} # PHYSICAL CONSTANTS - radius::NF # Radius of Planet - rotation::NF # Angular frequency of Planet's rotation - gravity::NF # Gravitational acceleration - R_dry::NF # specific gas constant for dry air [J/kg/K] - layer_thickness::NF # shallow water layer thickness [m] - μ_virt_temp::NF # used for virt temp calculation - κ::NF # = R_dry/cₚ, gas const for air over heat capacity - - # TIME STEPPING - Δt::NF # time step [s/m], use 2Δt for leapfrog, scaled by radius - Δt_unscaled::NF # time step [s], as Δt but not scaled with radius - Δt_sec::Int # time step [s] but encoded as 64-bit integer for rounding error-free accumulation - Δt_hrs::Float64 # time step [hrs] - robert_filter::NF # Robert (1966) time filter coefficient to suppress comput. mode - williams_filter::NF # Williams time filter (Amezcua 2011) coefficient for 3rd order acc - n_timesteps::Int # number of time steps to integrate for - - # OUTPUT TIME STEPPING - output_every_n_steps::Int # output every n time steps - n_outputsteps::Int # total number of output time steps - - # DIFFUSION AND DRAG - # drag_strat::NF # drag [1/s] for zonal wind in the stratosphere - - # INTERFACE FORCING - interface_relax_time::NF # time scale [1/s] for interface relaxation - - # PARAMETRIZATIONS - # Large-scale condensation (occurs when relative humidity exceeds a given threshold) - RH_thresh_pbl_lsc::NF # Relative humidity threshold for LSC in PBL - RH_thresh_range_lsc::NF # Vertical range of relative humidity threshold - RH_thresh_max_lsc ::NF # Maximum relative humidity threshold - humid_relax_time_lsc::NF # Relaxation time for humidity (hours) - - # Convection - pres_thresh_cnv::NF # Minimum (normalised) surface pressure for the occurrence of convection - RH_thresh_pbl_cnv::NF # Relative humidity threshold for convection in PBL - RH_thresh_trop_cnv::NF # Relative humidity threshold for convection in the troposphere - humid_relax_time_cnv::NF # Relaxation time for PBL humidity (hours) - max_entrainment::NF # Maximum entrainment as a fraction of cloud-base mass flux - ratio_secondary_mass_flux::NF # Ratio between secondary and primary mass flux at cloud-base + "Radius of Planet [m]" + radius::NF + + "Angular frequency of Planet's rotation [1/s]" + rotation::NF + + "Gravitational acceleration [m/s^2]" + gravity::NF + + "shallow water layer thickness [m]" + layer_thickness::NF + + # THERMODYNAMICS + "specific gas constant for dry air [J/kg/K]" + R_dry::NF + + "specific gas constant for water vapour [J/kg/K]" + R_vapour::NF + + "used in Tv = T(1+μq) for virt temp Tv(T,q) calculation" + μ_virt_temp::NF + + "specific heat at constant pressure [J/K/kg]" + cₚ::NF + + "= R_dry/cₚ, gas const for air over heat capacity" + κ::NF + + "coriolis frequency [1/s] (scaled by radius as is vorticity) = 2Ω*sin(lat)*radius" + f_coriolis::Vector{NF} + + # ADIABATIC TERM + "σ-related factor A needed for adiabatic terms" + σ_lnp_A::Vector{NF} + + "σ-related factor B needed for adiabatic terms" + σ_lnp_B::Vector{NF} + + # GEOPOTENTIAL INTEGRATION (on half/full levels) + "= R*(ln(p_k+1) - ln(p_k+1/2)), for half level geopotential" + Δp_geopot_half::Vector{NF} + + "= R*(ln(p_k+1/2) - ln(p_k)), for full level geopotential" + Δp_geopot_full::Vector{NF} + + # REFERENCE TEMPERATURE PROFILE for implicit + "reference temperature profile" + temp_ref_profile::Vector{NF} end """ +$(TYPEDSIGNATURES) Generator function for a DynamicsConstants struct. """ -function DynamicsConstants(P::Parameters) +function DynamicsConstants( spectral_grid::SpectralGrid, + planet::AbstractPlanet, + atmosphere::AbstractAtmosphere, + geometry::Geometry) # PHYSICAL CONSTANTS - (;R_dry, R_vapour, cₚ) = P - (;radius, rotation,gravity) = P.planet - (;layer_thickness) = P - H₀ = layer_thickness*1000 # ShallowWater: convert from [km]s to [m] + (;R_dry, R_vapour, lapse_rate, cₚ) = atmosphere + (;ΔT_stratosphere, σ_tropopause, temp_ref) = atmosphere + (;NF, radius) = spectral_grid + (;rotation, gravity) = planet + layer_thickness = atmosphere.layer_thickness*1000 # ShallowWater: convert from [km] to [m] ξ = R_dry/R_vapour # Ratio of gas constants: dry air / water vapour [1] μ_virt_temp = (1-ξ)/ξ # used in Tv = T(1+μq), for conversion from humidity q # and temperature T to virtual temperature Tv κ = R_dry/cₚ # = 2/7ish for diatomic gas - # TIME INTEGRATION CONSTANTS - (;robert_filter, williams_filter) = P - (;trunc, Δt_at_T31, n_days, output_dt) = P - - # PARAMETRIZATION CONSTANTS - (;RH_thresh_pbl_lsc, RH_thresh_range_lsc, RH_thresh_max_lsc, humid_relax_time_lsc) = P # Large-scale condensation - (;pres_thresh_cnv, RH_thresh_pbl_cnv, RH_thresh_trop_cnv, humid_relax_time_cnv, - max_entrainment, ratio_secondary_mass_flux) = P # Convection - - Δt = round(60*Δt_at_T31*(32/(trunc+1)))# scale time step Δt to specified resolution, [min] → [s] - Δt_sec = convert(Int,Δt) # encode time step Δt [s] as integer - Δt_hrs = Δt/3600 # convert time step Δt from minutes to hours - n_timesteps = ceil(Int,24*n_days/Δt_hrs) # number of time steps to integrate for - output_every_n_steps = max(1,floor(Int,output_dt/Δt_hrs)) # output every n time steps - n_outputsteps = (n_timesteps ÷ output_every_n_steps)+1 # total number of output time steps - - # stratospheric drag [1/s] from time scale [hrs] - # drag_strat = 1/(P.diffusion.time_scale_stratosphere*3600) - - # interface relaxation forcing - (;interface_relax_time) = P - interface_relax_time *= 3600/radius # convert from hours to seconds - - # SCALING - Δt_unscaled = Δt # [s] not scaled - Δt /= radius # [s/m] scale with radius + # CORIOLIS FREQUENCY (scaled by radius as is vorticity) + (;sinlat) = geometry + f_coriolis = 2rotation*sinlat*radius + + # ADIABATIC TERM, Simmons and Burridge, 1981, eq. 3.12 + (;σ_levels_half,σ_levels_full,σ_levels_thick) = geometry + # precompute ln(σ_k+1/2) - ln(σ_k-1/2) but swap sign, include 1/Δσₖ + σ_lnp_A = log.(σ_levels_half[1:end-1]./σ_levels_half[2:end]) ./ σ_levels_thick + σ_lnp_A[1] = 0 # the corresponding sum is 1:k-1 so 0 to replace log(0) from above + + # precompute the αₖ = 1 - p_k-1/2/Δpₖ*log(p_k+1/2/p_k-1/2) term in σ coordinates + σ_lnp_B = 1 .- σ_levels_half[1:end-1]./σ_levels_thick .* + log.(σ_levels_half[2:end]./σ_levels_half[1:end-1]) + σ_lnp_B[1] = σ_levels_half[1] <= 0 ? log(2) : σ_lnp_B[1] # set α₁ = log(2), eq. 3.19 + σ_lnp_B .*= -1 # absorb sign from -1/Δσₖ only, eq. 3.12 + + # GEOPOTENTIAL + Δp_geopot_half, Δp_geopot_full = initialize_geopotential(σ_levels_full,σ_levels_half,R_dry) + + # VERTICAL REFERENCE TEMPERATURE PROFILE (TODO: don't initialize here but in initalize! ?) + # integrate hydrostatic equation from pₛ to p, use ideal gas law p = ρRT and linear + # temperature decrease with height: T = Tₛ - ΔzΓ with lapse rate Γ + # for stratosphere (σ < σ_tropopause) increase temperature (Jablonowski & Williamson. 2006, eq. 5) + RΓg⁻¹ = R_dry*lapse_rate/(1000*gravity) # convert lapse rate from [K/km] to [K/m] + temp_ref_profile = [temp_ref*σ^RΓg⁻¹ for σ in σ_levels_full] + temp_ref_profile .+= [σ < σ_tropopause ? ΔT_stratosphere*(σ_tropopause-σ)^5 : 0 for σ in σ_levels_full] # This implies conversion to NF - return DynamicsConstants{P.NF}( radius,rotation,gravity,R_dry,H₀,μ_virt_temp,κ, - Δt,Δt_unscaled,Δt_sec,Δt_hrs, - robert_filter,williams_filter,n_timesteps, - output_every_n_steps, n_outputsteps, - # drag_strat, - interface_relax_time, - RH_thresh_pbl_lsc, RH_thresh_range_lsc, - RH_thresh_max_lsc, humid_relax_time_lsc, pres_thresh_cnv, - RH_thresh_pbl_cnv, RH_thresh_trop_cnv, humid_relax_time_cnv, - max_entrainment, ratio_secondary_mass_flux, - ) + return DynamicsConstants{NF}(; radius,rotation,gravity,layer_thickness, + R_dry,R_vapour,μ_virt_temp,cₚ,κ, + f_coriolis, + σ_lnp_A,σ_lnp_B, + Δp_geopot_half, Δp_geopot_full, + temp_ref_profile) end diff --git a/src/dynamics/define_diffusion.jl b/src/dynamics/define_diffusion.jl deleted file mode 100644 index f0f449c9a..000000000 --- a/src/dynamics/define_diffusion.jl +++ /dev/null @@ -1,123 +0,0 @@ -Base.@kwdef struct HyperDiffusion <: DiffusionParameters - # hyperdiffusion for temp, vor, div everywhere - # several powers of Laplacians, default 4 and 2, are added - # with respective time scales and scalings with resolution - power::Float64 = 4.0 # Power of Laplacian - time_scale::Float64 = 2.4 # Diffusion time scales [hrs] - resolution_scaling::Float64 = 0.5 # 0: constant with T - # 1: (inverse) linear with T - # 2: (inverse) quadratic, etc - - # incrased diffusion in stratosphere - power_stratosphere::Float64 = 2.0 # different power for stratosphere - tapering_σ::Float64 = 0.2 # scale towards that power linearly above this σ - - # increase diffusion based on high vorticity levels - adaptive::Bool = true # swith on/off - vor_max::Float64 = 1e-4 # [1/s] above this, diffusion is increased - adaptive_strength::Float64 = 2.0 # increase strength above vor_max by this factor - # times max(abs(vor))/vor_max -end - -""" - HD = HorizontalDiffusion(...) - -Horizontal Diffusion struct containing all the preallocated arrays for the calculation -of horizontal diffusion.""" -struct HorizontalDiffusion{NF<:AbstractFloat} # Number format NF - - lmax::Int # max degree l of spherical harmonics - time_scale::NF # strength of diffusion ~1/time_scale - - # (Hyper) diffusion, precalculated for each spherical harm degree - # and for each layer (to allow for varying orders/strength in the vertical) - ∇²ⁿ::Vector{Vector{NF}} # explicit part - ∇²ⁿ_implicit::Vector{Vector{NF}} # implicit part -end - -""" - HD = HorizontalDiffusion(::Parameters,::GeoSpectral,::Boundaries) - -Generator function for a HorizontalDiffusion struct `HD`. Precalculates damping matrices for -horizontal hyperdiffusion for temperature, vorticity and divergence, with an implicit term -and an explicit term. Also precalculates correction terms (horizontal and vertical) for -temperature and humidity. -""" -function HorizontalDiffusion( scheme::HyperDiffusion, - P::Parameters, - C::DynamicsConstants, - G::Geometry, - S::SpectralTransform{NF}) where NF - (;lmax) = S - (;radius) = P.planet - (;time_scale, resolution_scaling) = scheme - (;nlev) = G - - # Reduce diffusion time scale (=increase diffusion) with resolution - # times 1/radius because time step Δt is scaled with 1/radius - time_scale = 1/radius*3600*scheme.time_scale * (32/(lmax+1))^resolution_scaling - - # PREALLOCATE as vector as only dependend on degree l - # Damping coefficients for explicit part of the diffusion (=ν∇²ⁿ) - ∇²ⁿ = [zeros(NF,lmax+2) for _ in 1:nlev] # for temp, vor, div (explicit) - ∇²ⁿ_implicit = [zeros(NF,lmax+2) for _ in 1:nlev] # Implicit part (= 1/(1+2Δtν∇²ⁿ)) - HD = HorizontalDiffusion{NF}(lmax,time_scale,∇²ⁿ,∇²ⁿ_implicit) # initialize struct - - # PRECOMPUTE diffusion operators based on non-adaptive diffusion - vor_max = 0 # not increased diffusion initially - adapt_diffusion!(HD,scheme,vor_max,G,C) # fill diffusion arrays - return HD -end - -function adapt_diffusion!( HD::HorizontalDiffusion, - scheme::HyperDiffusion, - vor_max::Real, - G::Geometry, - C::DynamicsConstants) - nlev = length(HD.∇²ⁿ) - for k in 1:nlev - adapt_diffusion!(HD,scheme,vor_max,k,G,C) - end -end - -function adapt_diffusion!( HD::HorizontalDiffusion, - scheme::HyperDiffusion, - vor_max::Real, - k::Int, - G::Geometry, - C::DynamicsConstants) - - (;lmax,time_scale,∇²ⁿ,∇²ⁿ_implicit) = HD - (;power, power_stratosphere, tapering_σ) = scheme - (;Δt) = C - σ = G.σ_levels_full[k] - - # ADAPTIVE/FLOW AWARE - # increase diffusion if maximum vorticity per layer is larger than scheme.vor_max - if scheme.adaptive - # /= as 1/time_scale*∇²ʰ below - time_scale /= 1 + (scheme.adaptive_strength-1)*max(0,vor_max/scheme.vor_max - 1) - end - - # NORMALISATION - # Diffusion is applied by multiplication of the eigenvalues of the Laplacian -l*(l+1) - # normalise by the largest eigenvalue -lmax*(lmax+1) such that the highest wavenumber lmax - # is dampened to 0 at the given time scale raise to a power of the Laplacian for hyperdiffusion - # (=more scale-selective for smaller wavenumbers) - largest_eigenvalue = -lmax*(lmax+1) - - # VERTICAL TAPERING for the stratosphere - # go from 1 to 0 between σ=0 and tapering_σ - tapering = max(0,(tapering_σ-σ)/tapering_σ) # ∈ [0,1] - p = power + tapering*(power_stratosphere - power) - - @inbounds for l in 0:lmax+1 # diffusion for every degree l, but indendent of order m - eigenvalue_norm = -l*(l+1)/largest_eigenvalue # normalised diffusion ∇², power=1 - - # Explicit part (=-ν∇²ⁿ), time scales to damping frequencies [1/s] times norm. eigenvalue - ∇²ⁿ[k][l+1] = -eigenvalue_norm^p/time_scale - - # and implicit part of the diffusion (= 1/(1-2Δtν∇²ⁿ)) - ∇²ⁿ_implicit[k][l+1] = 1/(1-2Δt*∇²ⁿ[k][l+1]) - end -end \ No newline at end of file diff --git a/src/dynamics/define_implicit.jl b/src/dynamics/define_implicit.jl deleted file mode 100644 index fbf273d40..000000000 --- a/src/dynamics/define_implicit.jl +++ /dev/null @@ -1,47 +0,0 @@ -# SHALLOW WATER MODEL -""" - I = ImplicitShallowWater( ξH₀::Vector, - g∇²::Vector, - ξg∇²::Vector, - S⁻¹::Vector) - -Struct that holds various precomputed arrays for the semi-implicit correction to -prevent gravity waves from amplifying in the shallow water model.""" -struct ImplicitShallowWater{NF<:AbstractFloat} <: AbstractImplicit{NF} - ξH₀::Vector{NF} # = 2αΔt*layer_thickness, store in vec for mutability - g∇²::Vector{NF} # = gravity*eigenvalues - ξg∇²::Vector{NF} # = 2αΔt*gravity*eigenvalues - S⁻¹::Vector{NF} # = 1 / (1-ξH₀*ξg∇²), implicit operator -end - -# PRIMITIVE EQUATION MODEL -""" - I = ImplicitPrimitiveEq(ξ::Vector, - R::Matrix, - U::Vector, - L::Matrix, - W::Vector, - S⁻¹::Matrix) - -Struct that holds various precomputed arrays for the semi-implicit correction to -prevent gravity waves from amplifying in the primitive equation model.""" -struct ImplicitPrimitiveEq{NF<:AbstractFloat} <: AbstractImplicit{NF} - ξ::Base.RefValue{NF}# = 2α*Δt, packed in a RefValue for mutability - - # the following arrays all have ξ = 2αΔt absorbed L <- ξL, etc. - R::Matrix{NF} # divergence: used for δD = G_D + ξ(RδT + Uδlnpₛ) - U::Vector{NF} # divergence: see above - L::Matrix{NF} # tempereature: used for δT = G_T + ξLδD - W::Vector{NF} # surface pressure: used for δlnpₛ = G_lnpₛ + ξWδD - - # components of L - L0::Vector{NF} - L1::Matrix{NF} - L2::Vector{NF} - L3::Matrix{NF} - L4::Vector{NF} - - # combined implicit operator S - S::Matrix{NF} - S⁻¹::Array{NF,3} # used for δD = S⁻¹G, S = 1 - ξ²(RL + UW) -end \ No newline at end of file diff --git a/src/dynamics/diagnostic_variables.jl b/src/dynamics/diagnostic_variables.jl index 3eed96a97..dea598772 100644 --- a/src/dynamics/diagnostic_variables.jl +++ b/src/dynamics/diagnostic_variables.jl @@ -16,20 +16,18 @@ struct Tendencies{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} end function Base.zeros(::Type{Tendencies}, - G::Geometry{NF}, - S::SpectralTransform{NF}) where NF + SG::SpectralGrid) - (;Grid, nlat_half) = G - (;lmax, mmax) = S + (;NF, trunc,Grid, nlat_half) = SG LTM = LowerTriangularMatrix # use one more l for size compat with vector quantities - vor_tend = zeros(LTM{Complex{NF}},lmax+2,mmax+1) # vorticity - div_tend = zeros(LTM{Complex{NF}},lmax+2,mmax+1) # divergence - temp_tend = zeros(LTM{Complex{NF}},lmax+2,mmax+1) # absolute Temperature - humid_tend = zeros(LTM{Complex{NF}},lmax+2,mmax+1) # specific humidity - u_tend = zeros(LTM{Complex{NF}},lmax+2,mmax+1) # zonal velocity - v_tend = zeros(LTM{Complex{NF}},lmax+2,mmax+1) # meridional velocity + 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 @@ -55,9 +53,9 @@ struct GridVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} v_grid ::Grid # meridional velocity *coslat [m/s] end -function Base.zeros(::Type{GridVariables},G::Geometry{NF}) where NF +function Base.zeros(::Type{GridVariables},SG::SpectralGrid) - (;Grid, nlat_half) = G + (;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 @@ -100,15 +98,13 @@ struct DynamicsVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} end function Base.zeros(::Type{DynamicsVariables}, - G::Geometry{NF}, - S::SpectralTransform{NF}) where NF + SG::SpectralGrid) - (;lmax, mmax) = S - (;Grid, nlat_half) = G + (;NF, trunc, Grid, nlat_half) = SG # MULTI-PURPOSE VECTOR (a,b), work array to be reused in various places - a = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) - b = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) + 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) @@ -116,8 +112,8 @@ function Base.zeros(::Type{DynamicsVariables}, 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}},lmax+2,mmax+1) - geopot = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+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 @@ -132,22 +128,20 @@ struct DiagnosticVariablesLayer{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} 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? + npoints ::Int # number of grid points + k ::Int # which vertical model level? + temp_average ::Base.RefValue{NF} # average temperature for this level end function Base.zeros(::Type{DiagnosticVariablesLayer}, - G::Geometry{NF}, - S::SpectralTransform{NF}; - k::Integer=0) where NF # use k=0 (i.e. unspecified) as default - - (;npoints) = G - - tendencies = zeros(Tendencies,G,S) - grid_variables = zeros(GridVariables,G) - dynamics_variables = zeros(DynamicsVariables,G,S) - - return DiagnosticVariablesLayer(tendencies,grid_variables,dynamics_variables,npoints,k) + SG::SpectralGrid, + k::Integer=0) # use k=0 (i.e. unspecified) as default + (;npoints) = SG + tendencies = zeros(Tendencies,SG) + 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) end struct SurfaceVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} @@ -170,15 +164,13 @@ struct SurfaceVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} end function Base.zeros(::Type{SurfaceVariables}, - G::Geometry{NF}, - S::SpectralTransform{NF}) where NF + SG::SpectralGrid) - (;Grid, nlat_half, npoints) = G - (;lmax, mmax) = S + (;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}},lmax+2,mmax+1) + pres_tend = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) pres_tend_grid = zeros(Grid{NF},nlat_half) # gradients of log surface pressure @@ -189,7 +181,7 @@ function Base.zeros(::Type{SurfaceVariables}, 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}},lmax+2,mmax+1) + div_mean = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) # precipitation fields precip_large_scale = zeros(Grid{NF},nlat_half) @@ -206,34 +198,36 @@ end DiagnosticVariables{Grid<:AbstractGrid,NF<:AbstractFloat} Struct holding the diagnostic variables.""" -struct DiagnosticVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} +struct DiagnosticVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF},Model<:ModelSetup} layers ::Vector{DiagnosticVariablesLayer{NF,Grid}} surface ::SurfaceVariables{NF,Grid} columns ::Vector{ColumnVariables{NF}} - temp_profile::Vector{NF} + + nlat_half::Int # resolution parameter of any Grid nlev ::Int # number of vertical levels npoints ::Int # number of grid points + + scale::Base.RefValue{NF} # vorticity and divergence are scaled by radius end function Base.zeros(::Type{DiagnosticVariables}, - G::Geometry{NF}, - S::SpectralTransform{NF}) where NF + SG::SpectralGrid{Model}) where {Model<:ModelSetup} - (;nlev, npoints, n_stratosphere_levels) = G - layers = [zeros(DiagnosticVariablesLayer,G,S;k) for k in 1:nlev] - surface = zeros(SurfaceVariables,G,S) + (;NF,Grid,nlat_half, nlev, npoints) = SG + layers = [zeros(DiagnosticVariablesLayer,SG,k) for k in 1:nlev] + surface = zeros(SurfaceVariables,SG) # create one column variable per thread to avoid race conditions nthreads = Threads.nthreads() - columns = [ColumnVariables{NF}(;nlev,n_stratosphere_levels) for _ in 1:nthreads] + columns = [ColumnVariables{NF}(;nlev) for _ in 1:nthreads] - # global temperature profile, recalculated occasionally for the implicit solver - temp_profile = zeros(NF,nlev) + scale = Ref(convert(SG.NF,SG.radius)) - return DiagnosticVariables(layers,surface,columns,temp_profile,nlev,npoints) + return DiagnosticVariables{NF,Grid{NF},Model}( + layers,surface,columns,nlat_half,nlev,npoints,scale) end -DiagnosticVariables(G::Geometry{NF},S::SpectralTransform{NF}) where NF = zeros(DiagnosticVariables,G,S) +DiagnosticVariables(SG::SpectralGrid) = zeros(DiagnosticVariables,SG) # LOOP OVER ALL GRID POINTS (extend from RingGrids module) RingGrids.eachgridpoint(diagn::DiagnosticVariables) = Base.OneTo(diagn.npoints) diff --git a/src/dynamics/diffusion.jl b/src/dynamics/diffusion.jl deleted file mode 100644 index 6441a55f7..000000000 --- a/src/dynamics/diffusion.jl +++ /dev/null @@ -1,117 +0,0 @@ -""" - horizontal_diffusion!( tendency::LowerTriangularMatrix{Complex}, - A::LowerTriangularMatrix{Complex}, - ∇²ⁿ_expl::AbstractVector, - ∇²ⁿ_impl::AbstractVector) - -Apply horizontal diffusion to a 2D field `A` in spectral space by updating its tendency `tendency` -with an implicitly calculated diffusion term. The implicit diffusion of the next time step is split -into an explicit part `∇²ⁿ_expl` and an implicit part `∇²ⁿ_impl`, such that both can be calculated -in a single forward step by using `A` as well as its tendency `tendency`.""" -function horizontal_diffusion!( tendency::LowerTriangularMatrix{Complex{NF}}, # tendency of a - A::LowerTriangularMatrix{Complex{NF}}, # spectral horizontal field - ∇²ⁿ_expl::AbstractVector{NF}, # explicit spectral damping - ∇²ⁿ_impl::AbstractVector{NF} # implicit spectral damping - ) where {NF<:AbstractFloat} - lmax,mmax = size(tendency) # 1-based - @boundscheck size(tendency) == size(A) || throw(BoundsError) - @boundscheck lmax <= length(∇²ⁿ_expl) == length(∇²ⁿ_impl) || throw(BoundsError) - - lm = 0 - @inbounds for m in 1:mmax # loops over all columns/order m - for l in m:lmax-1 # but skips the lmax+2 degree (1-based) - lm += 1 # single index lm corresponding to harmonic l,m - tendency[lm] = (tendency[lm] + ∇²ⁿ_expl[l]*A[lm])*∇²ⁿ_impl[l] - end - lm += 1 # skip last row for scalar quantities - end -end - -# which variables to apply horizontal diffusion to -diffusion_vars(::Barotropic) = (:vor,) -diffusion_vars(::ShallowWater) = (:vor,:div) -diffusion_vars(::PrimitiveEquation) = (:vor,:div,:temp) - -function horizontal_diffusion!( progn::PrognosticLayerTimesteps, - diagn::DiagnosticVariablesLayer, - model::ModelSetup, - lf::Int=1) # leapfrog index used (2 is unstable) - HD = model.horizontal_diffusion - diffusion_scheme = model.parameters.diffusion - adapt_diffusion!(HD,diffusion_scheme,diagn,model.geometry,model.constants) - - # pick precalculated hyperdiffusion operator for layer k - k = diagn.k - ∇²ⁿ = HD.∇²ⁿ[k] - ∇²ⁿ_implicit = HD.∇²ⁿ_implicit[k] - - for varname in diffusion_vars(model) - var = getfield(progn.timesteps[lf],varname) - var_tend = getfield(diagn.tendencies,Symbol(varname,:_tend)) - horizontal_diffusion!(var_tend,var,∇²ⁿ,∇²ⁿ_implicit) - end -end - -function adapt_diffusion!( HD::HorizontalDiffusion, - scheme::HyperDiffusion, - diagn::DiagnosticVariablesLayer, - G::Geometry, - C::DynamicsConstants) - scheme.adaptive || return nothing - vor_min, vor_max = extrema(diagn.grid_variables.vor_grid) - vor_abs_max = max(abs(vor_min), abs(vor_max))/G.radius - adapt_diffusion!(HD,scheme,vor_abs_max,diagn.k,G,C) -end - -""" - stratospheric_zonal_drag!( tendency::AbstractArray{Complex{NF},3}, # tendency of - A::AbstractArray{Complex{NF},3}, # spectral vorticity or divergence - drag::Real # drag coefficient [1/s] - ) where {NF<:AbstractFloat} # number format NF - -Zonal drag in the uppermost layer of the stratosphere of 3D spectral field `A` (vorticity or divergence). -Drag is applied explicitly to the time step in `A` and its tendency `tendency` is changed in-place. -`drag` is the drag coefficient of unit 1/s. -""" -function stratospheric_zonal_drag!( tendency::AbstractArray{Complex{NF},3}, # tendency of - A::AbstractArray{Complex{NF},3}, # spectral vorticity or divergence - drag::Real # drag coefficient [1/s] - ) where {NF<:AbstractFloat} # number format NF - - lmax,mmax,nlev = size(A) # spherical harmonic degree l, order m, number of vertical levels nlev - lmax -= 1 # convert to 0-based l,m - mmax -= 1 - @boundscheck size(A) == size(tendency) || throw(BoundsError()) - - drag_NF = convert(NF,drag) - - @inbounds for l in 1:lmax+1 # loop over degree l, but 1-based - # size(A) = lmax x mmax x nlev, nlev = 1 is uppermost model level - # apply drag only to largest zonal wavenumber (m = 0) and in the uppermost model level (k=1) - tendency[l,1,1] = tendency[l,1,1] - drag_NF*A[l,1,1] - end -end - -"""Orographic temperature correction for absolute temperature to be applied before the horizontal diffusion.""" -function orographic_correction!(A_corrected::AbstractArray{Complex{NF},3}, # correction of - A::AbstractArray{Complex{NF},3}, # 3D spectral temperature or humidity - correction_horizontal::Matrix{Complex{NF}}, # horizontal correction matrix - correction_vertical::Vector{NF}, # vertical correction vector - ) where NF - - lmax,mmax,nlev = size(A) # degree l, order m of the spherical harmonics - lmax -= 1 # convert to 0-based - mmax -= 1 - - @boundscheck size(A) == size(A_corrected) || throw(BoundsError()) - @boundscheck (lmax+1,mmax+1) == size(correction_horizontal) || throw(BoundsError()) - @boundscheck (nlev,) == size(correction_vertical) || throw(BoundsError()) - - @inbounds for k in 1:nlev # vertical levels - for m in 1:mmax+1 # order of spherical harmonics - for l in m:lmax+1 # degree of spherical harmonics - A_corrected[l,m,k] = A[l,m,k] + hori_correction[l,m]*vert_correction[k] - end - end - end -end \ No newline at end of file diff --git a/src/dynamics/forcing.jl b/src/dynamics/forcing.jl new file mode 100644 index 000000000..6cb50d0cb --- /dev/null +++ b/src/dynamics/forcing.jl @@ -0,0 +1,42 @@ +struct NoForcing{NF} <: AbstractForcing{NF} end +NoForcing(SG::SpectralGrid) = NoForcing{SG.NF}() + +function initialize!( forcing::NoForcing, + model::ModelSetup) + return nothing +end + +function forcing!( diagn::DiagnosticVariablesLayer, + forcing::NoForcing, + time::DateTime) + return nothing +end + +# function interface_relaxation!( η::LowerTriangularMatrix{Complex{NF}}, +# surface::SurfaceVariables{NF}, +# time::DateTime, # time of relaxation +# M::ShallowWater, # contains η⁰, which η is relaxed to +# ) where NF + +# (; pres_tend ) = surface +# (; seasonal_cycle, equinox, axial_tilt ) = M.parameters.planet +# A = M.parameters.interface_relax_amplitude + +# s = 45/23.5 # heuristic conversion to Legendre polynomials +# θ = seasonal_cycle ? s*axial_tilt*sin(Dates.days(time - equinox)/365.25*2π) : 0 +# η2 = convert(NF,A*(2sind(θ))) # l=1,m=0 harmonic +# η3 = convert(NF,A*(0.2-1.5cosd(θ))) # l=2,m=0 harmonic + +# τ⁻¹ = inv(M.constants.interface_relax_time) +# pres_tend[2] += τ⁻¹*(η2-η[2]) +# pres_tend[3] += τ⁻¹*(η3-η[3]) +# end + +# "turn on interface relaxation for shallow water?" +# interface_relaxation::Bool = false + +# "time scale [hrs] of interface relaxation" +# interface_relax_time::Float64 = 96 + +# "Amplitude [m] of interface relaxation" +# interface_relax_amplitude::Float64 = 300 diff --git a/src/dynamics/geometry.jl b/src/dynamics/geometry.jl deleted file mode 100644 index e8abddc7b..000000000 --- a/src/dynamics/geometry.jl +++ /dev/null @@ -1,223 +0,0 @@ -""" - Geometry{NF<:AbstractFloat} <: AbstractGeometry - -Geometry struct containing parameters and arrays describing an iso-latitude grid <:AbstractGrid -and the vertical levels. NF is the number format used for the precomputed constants. -""" -struct Geometry{NF<:AbstractFloat} <: AbstractGeometry{NF} # NF: Number format - - # GRID TYPE AND RESOLUTION - Grid::Type{<:AbstractGrid} - nlat_half::Int # resolution parameter nlat_half, # of latitudes on one hemisphere (incl Equator) - - # GRID-POINT SPACE - nlon_max::Int # Maximum number of longitudes (at/around Equator) - nlon::Int # Same (used for compatibility) - nlat::Int # Number of latitudes - nlev::Int # Number of vertical levels - npoints::Int # total number of grid points - radius::Float64 # Planet's radius [m] - - # LATITUDES (either Gaussian, equi-angle, HEALPix or OctaHEALPix lats, depending on grid) - latd::Vector{Float64} # array of latitudes in degrees (90˚...-90˚) - - # LONGITUDES - lond::Vector{Float64} # array of longitudes in degrees (0...360˚), empty for non-full grids - - # COORDINATES - londs::Vector{NF} # longitude (-180˚...180˚) for each grid point in ring order - latds::Vector{NF} # latitude (-90˚...˚90) for each grid point in ring order - - # SINES AND COSINES OF LATITUDE - sinlat::Vector{NF} # sin of latitudes - coslat::Vector{NF} # cos of latitudes - coslat⁻¹::Vector{NF} # = 1/cos(lat) - coslat²::Vector{NF} # = cos²(lat) - coslat⁻²::Vector{NF} # = 1/cos²(lat) - - # CORIOLIS FREQUENCY (scaled by radius as is vorticity) - f_coriolis::Vector{NF} # = 2Ω*sin(lat)*radius - - # VERTICAL SIGMA COORDINATE σ = p/p0 (fraction of surface pressure) - n_stratosphere_levels::Int # number of upper levels for stratosphere - σ_levels_half::Vector{NF} # σ at half levels - σ_levels_full::Vector{NF} # σ at full levels - σ_levels_thick::Vector{NF} # σ level thicknesses - ln_σ_levels_full::Vector{NF} # log of σ at full levels - - σ_lnp_A::Vector{NF} # σ-related factors needed for adiabatic terms 1st - σ_lnp_B::Vector{NF} # and 2nd - - # VERTICAL REFERENCE TEMPERATURE PROFILE - temp_ref_profile::Vector{NF} - - # GEOPOTENTIAL INTEGRATION (on half/full levels) - Δp_geopot_half::Vector{NF} # = R*(ln(p_k+1) - ln(p_k+1/2)), for half level geopotential - Δp_geopot_full::Vector{NF} # = R*(ln(p_k+1/2) - ln(p_k)), for full level geopotential - lapserate_corr::Vector{NF} # ? - - # PARAMETERIZATIONS - entrainment_profile::Vector{NF} -end - -""" - G = Geometry(P::Parameters) - -Generator function to create the Geometry struct from parameters in `P`. -""" -function Geometry(P::Parameters,Grid::Type{<:AbstractGrid}) - - (;trunc, dealiasing, nlev) = P # grid type, spectral truncation, # of vertical levels - (;radius, rotation, gravity) = P.planet # radius of planet, angular frequency, gravity - (;R_dry, cₚ) = P # gas constant for dry air, heat capacity - (;σ_tropopause) = P # number of vertical levels used for stratosphere - (;temp_ref, temp_top, lapse_rate) = P # for reference atmosphere - (;ΔT_stratosphere) = P # used for stratospheric temperature increase - - # RESOLUTION PARAMETERS - # resolution parameter nlat_half (= number of lat rings on one hemisphere (Equator incl) - # from spectral resolution and dealiasing parameter (e.g. quadratic grid for T31) - nlat_half = SpeedyTransforms.get_nlat_half(trunc,dealiasing) - nlat = get_nlat(Grid,nlat_half) # 2nlat_half but -1 if grids have odd # of lat rings - nlon_max = get_nlon_max(Grid,nlat_half) # number of longitudes around the equator - nlon = nlon_max # same (used for compatibility) - npoints = get_npoints(Grid,nlat_half) # total number of grid points - - # LATITUDE VECTORS (based on Gaussian, equi-angle or HEALPix latitudes) - latd = get_latd(Grid,nlat_half) # latitude in 90˚...-90˚ - - # LONGITUDE VECTORS (empty for reduced grids) - lond = get_lond(Grid,nlat_half) # array of longitudes 0...360˚ (full grids only) - - # COORDINATES for every grid point in ring order - latds,londs = get_latdlonds(Grid,nlat_half) # in -90˚...90˚N, -180˚...180˚ - - # SINES AND COSINES OF LATITUDE - sinlat = sind.(latd) - coslat = cosd.(latd) - coslat⁻¹ = 1 ./ coslat - coslat² = coslat.^2 - coslat⁻² = 1 ./ coslat² - - # CORIOLIS FREQUENCY (scaled by radius as is vorticity) - f_coriolis = 2rotation*sinlat*radius - - # VERTICAL SIGMA COORDINATE - # σ = p/p0 (fraction of surface pressure) - # sorted such that σ_levels_half[end] is at the planetary boundary - σ_levels_half = vertical_coordinates(P) - σ_levels_full = 0.5*(σ_levels_half[2:end] + σ_levels_half[1:end-1]) - σ_levels_thick = σ_levels_half[2:end] - σ_levels_half[1:end-1] - ln_σ_levels_full = log.(vcat(σ_levels_full,1)) # include surface (σ=1) as last element - - # ADIABATIC TERM, Simmons and Burridge, 1981, eq. 3.12 - # precompute ln(σ_k+1/2) - ln(σ_k-1/2) but swap sign, include 1/Δσₖ - σ_lnp_A = log.(σ_levels_half[1:end-1]./σ_levels_half[2:end]) ./ σ_levels_thick - σ_lnp_A[1] = 0 # the corresponding sum is 1:k-1 so 0 to replace log(0) from above - - # precompute the αₖ = 1 - p_k-1/2/Δpₖ*log(p_k+1/2/p_k-1/2) term in σ coordinates - σ_lnp_B = 1 .- σ_levels_half[1:end-1]./σ_levels_thick .* - log.(σ_levels_half[2:end]./σ_levels_half[1:end-1]) - σ_lnp_B[1] = σ_levels_half[1] <= 0 ? log(2) : σ_lnp_B[1] # set α₁ = log(2), eq. 3.19 - σ_lnp_B .*= -1 # absorb sign from -1/Δσₖ only, eq. 3.12 - - # TROPOPAUSE/STRATOSPHERIC LEVELS - n_stratosphere_levels = sum(σ_levels_full .< σ_tropopause) # of levels above σ_tropopause - - # VERTICAL REFERENCE TEMPERATURE PROFILE - # integrate hydrostatic equation from pₛ to p, use ideal gas law p = ρRT and linear - # temperature decrease with height: T = Tₛ - ΔzΓ with lapse rate Γ - # for stratosphere (σ < σ_tropopause) increase temperature (Jablonowski & Williamson. 2006, eq. 5) - RΓg⁻¹ = R_dry*lapse_rate/(1000*gravity) # convert lapse rate from [K/km] to [K/m] - temp_ref_profile = [temp_ref*σ^RΓg⁻¹ for σ in σ_levels_full] - temp_ref_profile .+= [σ < σ_tropopause ? ΔT_stratosphere*(σ_tropopause-σ)^5 : 0 for σ in σ_levels_full] - - # GEOPOTENTIAL, coefficients to calculate geopotential - Δp_geopot_half, Δp_geopot_full = initialize_geopotential(σ_levels_full,σ_levels_half,R_dry) - - # LAPSE RATE correction - lapserate_corr = lapserate_correction(σ_levels_full,σ_levels_half,Δp_geopot_full) - - # Compute the entrainment coefficients for the convection parameterization. - (;max_entrainment) = P - entrainment_profile = zeros(nlev) - for k = 2:nlev-1 - entrainment_profile[k] = max(0, (σ_levels_full[k] - 0.5)^2) - end - - # profile as fraction of cloud-base mass flux - entrainment_profile /= sum(entrainment_profile) # Normalise - entrainment_profile *= max_entrainment # fraction of max entrainment - - # conversion to number format NF happens here - Geometry{P.NF}( Grid,nlat_half, - nlon_max,nlon,nlat,nlev,npoints,radius, - latd,lond,londs,latds, - sinlat,coslat,coslat⁻¹,coslat²,coslat⁻²,f_coriolis, - n_stratosphere_levels, - σ_levels_half,σ_levels_full,σ_levels_thick,ln_σ_levels_full, - σ_lnp_A,σ_lnp_B, - temp_ref_profile, - Δp_geopot_half,Δp_geopot_full,lapserate_corr,entrainment_profile) -end - -# use Grid in Parameters if not provided -Geometry(P::Parameters) = Geometry(P,P.Grid) - -"""Coefficients of the generalised logistic function to describe the vertical coordinate. -Default coefficients A,K,C,Q,B,M,ν are fitted to the old L31 configuration at ECMWF. -See geometry.jl and function `vertical_coordinate` for more informaiton. - -Following the notation of [https://en.wikipedia.org/wiki/Generalised_logistic_function](https://en.wikipedia.org/wiki/Generalised_logistic_function) (Dec 15 2021). - -Change default parameters for more/fewer levels in the stratosphere vs troposphere vs boundary layer.""" -Base.@kwdef struct GenLogisticCoefs <: Coefficients - A::Float64 = -0.283 # obtained from a fit in /input_date/vertical_coordinate/vertical_resolution.ipynb - K::Float64 = 0.871 - C::Float64 = 0.414 - Q::Float64 = 6.695 - B::Float64 = 10.336 - M::Float64 = 0.602 - ν::Float64 = 5.812 -end - -"""Generalised logistic function based on the coefficients in `coefs`.""" -function generalised_logistic(x,coefs::GenLogisticCoefs) - (; A,K,C,Q,B,M,ν ) = coefs - return @. A + (K-A)/(C+Q*exp(-B*(x-M)))^inv(ν) -end - -""" - σ_levels_half = vertical_coordinates(P::Parameters) - -Vertical sigma coordinates defined by their nlev+1 half levels `σ_levels_half`. Sigma coordinates are -fraction of surface pressure (p/p0) and are sorted from top (stratosphere) to bottom (surface). -The first half level is at 0 the last at 1. Evaluate a generalised logistic function with -coefficients in `P` for the distribution of values in between. Default coefficients follow -the L31 configuration historically used at ECMWF.""" -function vertical_coordinates(P::Parameters) - (;nlev,GLcoefs,σ_levels_half) = P - - if length(σ_levels_half) == 0 # choose σ levels automatically - z = range(0,1,nlev+1) # normalised = level/nlev - σ_levels_half = generalised_logistic(z,GLcoefs) - σ_levels_half .-= σ_levels_half[1] # topmost half-level is at 0 pressure - σ_levels_half ./= σ_levels_half[end] # lowermost half-level is at p=p_surface - else # choose σ levels manually - @assert σ_levels_half[1] >= 0 "First manually specified σ_levels_half has to be >0" - @assert σ_levels_half[end] == 1 "Last manually specified σ_levels_half has to be 1." - @assert nlev == (length(σ_levels_half) - 1) "nlev has to be length of σ_levels_half - 1" - end - - @assert isincreasing(σ_levels_half) "Vertical coordinates are not increasing." - return σ_levels_half -end - -""" - S = SpectralTransform(P::Parameters) - -Generator function for a SpectralTransform struct pulling in parameters from a Parameters struct.""" -function SpeedyTransforms.SpectralTransform(P::Parameters) - (;NF, Grid, trunc, dealiasing, recompute_legendre, legendre_shortcut) = P - return SpectralTransform(NF,Grid,trunc,recompute_legendre;legendre_shortcut,dealiasing) -end \ No newline at end of file diff --git a/src/dynamics/geopotential.jl b/src/dynamics/geopotential.jl index 371dad058..68d2fa9bc 100644 --- a/src/dynamics/geopotential.jl +++ b/src/dynamics/geopotential.jl @@ -1,9 +1,11 @@ """ - Δp_geopot_half, Δp_geopot_full = initialize_geopotential( σ_levels_full::Vector, - σ_levels_half::Vector, - R_dry::Real) +$(TYPEDSIGNATURES) +Precomputes constants for the vertical integration of the geopotential, defined as -Precomputes """ +`Φ_{k+1/2} = Φ_{k+1} + R*T_{k+1}*(ln(p_{k+1}) - ln(p_{k+1/2}))` (half levels) +`Φ_k = Φ_{k+1/2} + R*T_k*(ln(p_{k+1/2}) - ln(p_k))` (full levels) + +Same formula but `k → k-1/2`.""" function initialize_geopotential( σ_levels_full::Vector, σ_levels_half::Vector, R_dry::Real) @@ -29,43 +31,42 @@ function initialize_geopotential( σ_levels_full::Vector, return Δp_geopot_half, Δp_geopot_full end -function lapserate_correction( σ_levels_full::Vector, - σ_levels_half::Vector, - Δp_geopot_full::Vector) +# function lapserate_correction( σ_levels_full::Vector, +# σ_levels_half::Vector, +# Δp_geopot_full::Vector) - nlev = length(σ_levels_full) - @assert nlev+1 == length(σ_levels_half) "σ half levels must have length nlev+1" - @assert nlev == length(Δp_geopot_full) "σ half levels must have length nlev" +# nlev = length(σ_levels_full) +# @assert nlev+1 == length(σ_levels_half) "σ half levels must have length nlev+1" +# @assert nlev == length(Δp_geopot_full) "σ half levels must have length nlev" - lapserate_corr = zeros(nlev) - for k in 2:nlev-1 # only in the free troposphere - # TODO reference - lapserate_corr[k] = 0.5*Δp_geopot_full[k]* - log(σ_levels_half[k+1]/σ_levels_full[k]) / log(σ_levels_full[k+1]/σ_levels_full[k-1]) - end +# lapserate_corr = zeros(nlev) +# for k in 2:nlev-1 # only in the free troposphere +# # TODO reference +# lapserate_corr[k] = 0.5*Δp_geopot_full[k]* +# log(σ_levels_half[k+1]/σ_levels_full[k]) / log(σ_levels_full[k+1]/σ_levels_full[k-1]) +# end - return lapserate_corr -end +# return lapserate_corr +# end """ - geopotential!(diagn,B,G) - +$(TYPEDSIGNATURES) Compute spectral geopotential `geopot` from spectral temperature `temp` and spectral surface geopotential `geopot_surf` (orography*gravity). """ -function geopotential!( diagn::DiagnosticVariables{NF}, - B::Boundaries{NF}, # contains surface geopotential - G::Geometry{NF} # contains precomputed layer-thickness arrays - ) where NF # number format NF +function geopotential!( + diagn::DiagnosticVariables, + O::AbstractOrography, # contains surface geopotential + C::DynamicsConstants, # contains precomputed layer-thickness arrays +) - (;geopot_surf) = B.orography # = orography*gravity - (;Δp_geopot_half, Δp_geopot_full) = G # = R*Δlnp either on half or full levels - (;lapserate_corr) = G - (;nlev) = G # number of vertical levels + (;geopot_surf) = O # = orography*gravity + (;Δp_geopot_half, Δp_geopot_full) = C # = R*Δlnp either on half or full levels + (;nlev) = diagn # number of vertical levels - @boundscheck diagn.nlev == length(Δp_geopot_full) || throw(BoundsError) + @boundscheck nlev == length(Δp_geopot_full) || throw(BoundsError) - # for PrimitiveDryCore virtual temperature = absolute temperature here + # for PrimitiveDry virtual temperature = absolute temperature here # note these are not anomalies here as they are only in grid-point fields # BOTTOM FULL LAYER @@ -103,12 +104,19 @@ function geopotential!( diagn::DiagnosticVariables{NF}, # end end +""" +$(TYPEDSIGNATURES) +Calculate the geopotential based on `temp` in a single column. +This exclues the surface geopotential that would need to be added to the returned vector. +Function not used in the dynamical core but for post-processing and analysis.""" function geopotential!( temp::Vector, - G::Geometry) - - (;Δp_geopot_half, Δp_geopot_full, nlev) = G # = R*Δlnp either on half or full levels + C::DynamicsConstants) + nlev = length(temp) + (;Δp_geopot_half, Δp_geopot_full) = C # = R*Δlnp either on half or full levels geopot = zero(temp) + @boundscheck length(temp) == length(Δp_geopot_full) || throw(BoundsError) + # bottom layer geopot[nlev] = temp[nlev]*Δp_geopot_full[end] @@ -120,21 +128,27 @@ function geopotential!( temp::Vector, return geopot end +""" +$(TYPEDSIGNATURES) +calculates the geopotential in the ShallowWaterModel as g*η, +i.e. gravity times the interface displacement (field `pres`)""" function geopotential!( diagn::DiagnosticVariablesLayer, pres::LowerTriangularMatrix, C::DynamicsConstants) - (;gravity) = C (;geopot) = diagn.dynamics_variables - geopot .= pres*gravity end +# function barrier +function virtual_temperature!( diagn::DiagnosticVariablesLayer, + temp::LowerTriangularMatrix, # only needed for dispatch compat with DryCore + model::PrimitiveWet) + virtual_temperature!(diagn,temp,model.constants) +end + """ - virtual_temperature!( diagn::DiagnosticVariablesLayer, - temp::LowerTriangularMatrix, - M::PrimitiveWetCore) - +$(TYPEDSIGNATURES) Calculates the virtual temperature Tᵥ as Tᵥ = T(1+μq) @@ -143,58 +157,88 @@ With absolute temperature T, specific humidity q and μ = (1-ξ)/ξ, ξ = R_dry/R_vapour. -In grid-point space and then transforms Tᵥ back into spectral space -for the geopotential calculation.""" -function virtual_temperature!( diagn::DiagnosticVariablesLayer, - ::LowerTriangularMatrix, # only needed for dispatch compat with DryCore - model::PrimitiveWetCore) +in grid-point space.""" +function virtual_temperature!( + diagn::DiagnosticVariablesLayer, + temp::LowerTriangularMatrix, # only needed for dispatch compat with DryCore + constants::DynamicsConstants, + ) (;temp_grid, humid_grid, temp_virt_grid) = diagn.grid_variables (;temp_virt) = diagn.dynamics_variables - μ = model.constants.μ_virt_temp - S = model.spectral_transform + μ = constants.μ_virt_temp @inbounds for ij in eachgridpoint(temp_virt_grid, temp_grid, humid_grid) temp_virt_grid[ij] = temp_grid[ij]*(1 + μ*humid_grid[ij]) end + # TODO check that doing a non-linear virtual temperature in grid-point space + # but a linear virtual temperature in spectral space to avoid another transform + # does not cause any problems. Alternative do the transform or have a linear + # virtual temperature in both grid and spectral space # spectral!(temp_virt,temp_virt_grid,S) end -function linear_virtual_temperature!( diagn::DiagnosticVariablesLayer, - progn::PrognosticLayerTimesteps, - model::PrimitiveWetCore, - lf::Int) - - (;temp_virt) = diagn.dynamics_variables - μ = model.constants.μ_virt_temp - Tₖ = model.geometry.temp_ref_profile[diagn.k] - (;temp,humid) = progn.timesteps[lf] - - @. temp_virt = temp + Tₖ*μ*humid -end - - """ -For the PrimitiveDryCore temperautre and virtual temperature are the same (humidity=0). -Just copy over the arrays.""" +$(TYPEDSIGNATURES) +Virtual temperature in grid-point space: For the PrimitiveDry temperature +and virtual temperature are the same (humidity=0). Just copy over the arrays.""" function virtual_temperature!( diagn::DiagnosticVariablesLayer, temp::LowerTriangularMatrix, - ::PrimitiveDryCore) + model::PrimitiveDry) (;temp_grid, temp_virt_grid) = diagn.grid_variables (;temp_virt) = diagn.dynamics_variables copyto!(temp_virt_grid,temp_grid) - # copyto!(temp_virt,temp) end +""" +$(TYPEDSIGNATURES) +Linear virtual temperature for `model::PrimitiveDry`: Just copy over +arrays from `temp` to `temp_virt` at timestep `lf` in spectral space +as humidity is zero in this `model`.""" +function linear_virtual_temperature!( + diagn::DiagnosticVariablesLayer, + progn::PrognosticLayerTimesteps, + model::PrimitiveDry, + lf::Integer, +) + (;temp_virt) = diagn.dynamics_variables + (;temp) = progn.timesteps[lf] + copyto!(temp_virt,temp) +end + +# function barrier +function linear_virtual_temperature!( + diagn::DiagnosticVariablesLayer, + progn::PrognosticLayerTimesteps, + model::PrimitiveWet, + lf::Integer, +) + linear_virtual_temperature!(diagn,progn,model.constants,lf) +end + +""" +$(TYPEDSIGNATURES) +Calculates a linearised virtual temperature Tᵥ as + + Tᵥ = T + Tₖμq + +With absolute temperature T, layer-average temperarture Tₖ (computed in temperature_average!), +specific humidity q and + + μ = (1-ξ)/ξ, ξ = R_dry/R_vapour. + +in spectral space.""" function linear_virtual_temperature!( diagn::DiagnosticVariablesLayer, progn::PrognosticLayerTimesteps, - ::PrimitiveDryCore, + constants::DynamicsConstants, lf::Int) (;temp_virt) = diagn.dynamics_variables - (;temp) = progn.timesteps[lf] + μ = constants.μ_virt_temp + Tₖ = diagn.temp_average[] + (;temp,humid) = progn.timesteps[lf] - copyto!(temp_virt,temp) -end + @. temp_virt = temp + (Tₖ*μ)*humid +end \ No newline at end of file diff --git a/src/dynamics/horizontal_diffusion.jl b/src/dynamics/horizontal_diffusion.jl new file mode 100644 index 000000000..1ea907bc8 --- /dev/null +++ b/src/dynamics/horizontal_diffusion.jl @@ -0,0 +1,272 @@ +""" +Struct for horizontal hyper diffusion of vor, div, temp; implicitly in spectral space +with a `power` of the Laplacian (default=4) and the strength controlled by +`time_scale`. Options exist to scale the diffusion by resolution, and adaptive +depending on the current vorticity maximum to increase diffusion in active +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} + # DIMENSIONS + "spectral resolution" + trunc::Int + + "number of vertical levels" + nlev::Int + + # PARAMETERS + "power of Laplacian" + power::Float64 = 4.0 + + "diffusion time scales [hrs]" + time_scale::Float64 = 2.4 + + "stronger diffusion with resolution? 0: constant with trunc, 1: (inverse) linear with trunc, etc" + resolution_scaling::Float64 = 0.5 + + # incrased diffusion in stratosphere + "different power for tropopause/stratosphere" + power_stratosphere::Float64 = 2.0 + + "linearly scale towards power_stratosphere above this σ" + tapering_σ::Float64 = 0.2 + + # increase diffusion based on high vorticity levels + "adaptive = higher diffusion for layers with higher vorticity levels." + adaptive::Bool = true # swith on/off + + "above this (absolute) vorticity level [1/s], diffusion is increased" + vor_max::Float64 = 1e-4 + + "increase strength above `vor_max` by this factor times `max(abs(vor))/vor_max`" + adaptive_strength::Float64 = 2.0 + + # ARRAYS, precalculated for each spherical harmonics degree + # Barotropic and ShallowWater are fine with a constant time scale + ∇²ⁿ_2D::Vector{NF} = zeros(NF,trunc+2) # initialized with zeros, ones + ∇²ⁿ_2D_implicit::Vector{NF} = ones(NF,trunc+2) # as this corresponds to no diffusion + + # PrimitiveEquation models need something more adaptive + # and for each layer (to allow for varying orders/strength in the vertical) + ∇²ⁿ::Vector{Vector{NF}} = [zeros(NF,trunc+2) for _ in 1:nlev] # explicit part + ∇²ⁿ_implicit::Vector{Vector{NF}} = [ones(NF,trunc+2) for _ in 1:nlev] # implicit part +end + +"""$(TYPEDSIGNATURES) +Generator function based on the resolutin in `spectral_grid`. +Passes on keyword arguments.""" +function HyperDiffusion(spectral_grid::SpectralGrid,kwargs...) + (;NF,trunc,nlev) = spectral_grid # take resolution parameters from spectral_grid + return HyperDiffusion{NF}(;trunc,nlev,kwargs...) +end + +function Base.show(io::IO,HD::HyperDiffusion) + print(io,"$(typeof(HD)):") + keys = (:trunc,:nlev,:power,:time_scale,:resolution_scaling,:power_stratosphere, + :tapering_σ,:adaptive,:vor_max,:adaptive_strength) + for key in keys + val = getfield(HD,key) + s = "\n $key::$(typeof(val)) = $val" + print(io,s) + end +end + +"""$(TYPEDSIGNATURES) +Precomputes the hyper diffusion terms in `scheme` based on the +model time step, and possibly with a changing strength/power in +the vertical. +""" +function initialize!( scheme::HyperDiffusion, + model::ModelSetup) + # always initialize the 2D arrays + initialize!(scheme,model.time_stepping) + + # and the 3D arrays (different diffusion per layer) for primitive eq + for k in 1:scheme.nlev + initialize!(scheme,k,model.geometry,model.time_stepping) + end +end + +"""$(TYPEDSIGNATURES) +Precomputes the 2D hyper diffusion terms in `scheme` based on the +model time step.""" +function initialize!( scheme::HyperDiffusion, + L::TimeStepper) + + (;trunc,time_scale,∇²ⁿ_2D,∇²ⁿ_2D_implicit,power) = scheme + (;Δt, radius) = L + + # time scale times 1/radius because time step Δt is scaled with 1/radius + # time scale*3600 for [hrs] → [s] + time_scale = 1/radius*(3600*time_scale) + + # NORMALISATION + # Diffusion is applied by multiplication of the eigenvalues of the Laplacian -l*(l+1) + # normalise by the largest eigenvalue -lmax*(lmax+1) such that the highest wavenumber lmax + # is dampened to 0 at the given time scale raise to a power of the Laplacian for hyperdiffusion + # (=more scale-selective for smaller wavenumbers) + largest_eigenvalue = -trunc*(trunc+1) + + @inbounds for l in 0:trunc+1 # diffusion for every degree l, but indendent of order m + eigenvalue_norm = -l*(l+1)/largest_eigenvalue # normalised diffusion ∇², power=1 + + # Explicit part (=-ν∇²ⁿ), time scales to damping frequencies [1/s] times norm. eigenvalue + ∇²ⁿ_2D[l+1] = -eigenvalue_norm^power/time_scale + + # and implicit part of the diffusion (= 1/(1-2Δtν∇²ⁿ)) + ∇²ⁿ_2D_implicit[l+1] = 1/(1-2Δt*∇²ⁿ_2D[l+1]) + end +end + +"""$(TYPEDSIGNATURES) +Precomputes the hyper diffusion terms in `scheme` for layer `k` based on the +model time step in `L`, the vertical level sigma level in `G`, and +the current (absolute) vorticity maximum level `vor_max`""" +function initialize!( + scheme::HyperDiffusion, + k::Int, + G::Geometry, + L::TimeStepper, + vor_max::Real = 0, +) + (;trunc,time_scale,resolution_scaling,∇²ⁿ,∇²ⁿ_implicit) = scheme + (;power, power_stratosphere, tapering_σ) = scheme + (;Δt, radius) = L + σ = G.σ_levels_full[k] + + # Reduce diffusion time scale (=increase diffusion) with resolution + # times 1/radius because time step Δt is scaled with 1/radius + # time scale*3600 for [hrs] → [s] + time_scale = 1/radius*(3600*time_scale) * (32/(trunc+1))^resolution_scaling + + # ADAPTIVE/FLOW AWARE + # increase diffusion if maximum vorticity per layer is larger than scheme.vor_max + if scheme.adaptive + # /= as 1/time_scale*∇²ⁿ below + time_scale /= 1 + (scheme.adaptive_strength-1)*max(0,vor_max/scheme.vor_max - 1) + end + + # NORMALISATION + # Diffusion is applied by multiplication of the eigenvalues of the Laplacian -l*(l+1) + # normalise by the largest eigenvalue -lmax*(lmax+1) such that the highest wavenumber lmax + # is dampened to 0 at the given time scale raise to a power of the Laplacian for hyperdiffusion + # (=more scale-selective for smaller wavenumbers) + largest_eigenvalue = -trunc*(trunc+1) + + # VERTICAL TAPERING for the stratosphere + # go from 1 to 0 between σ=0 and tapering_σ + tapering = max(0,(tapering_σ-σ)/tapering_σ) # ∈ [0,1] + p = power + tapering*(power_stratosphere - power) + + @inbounds for l in 0:trunc+1 # diffusion for every degree l, but indendent of order m + eigenvalue_norm = -l*(l+1)/largest_eigenvalue # normalised diffusion ∇², power=1 + + # Explicit part (=-ν∇²ⁿ), time scales to damping frequencies [1/s] times norm. eigenvalue + ∇²ⁿ[k][l+1] = -eigenvalue_norm^p/time_scale + + # and implicit part of the diffusion (= 1/(1-2Δtν∇²ⁿ)) + ∇²ⁿ_implicit[k][l+1] = 1/(1-2Δt*∇²ⁿ[k][l+1]) + end +end + +"""$(TYPEDSIGNATURES) +Pre-function to other `initialize!(::HyperDiffusion)` initialisors that +calculates the (absolute) vorticity maximum for the layer of `diagn`.""" +function initialize!( + scheme::HyperDiffusion, + diagn::DiagnosticVariablesLayer, + G::Geometry, + L::TimeStepper, +) + scheme.adaptive || return nothing + vor_min, vor_max = extrema(diagn.grid_variables.vor_grid) + vor_abs_max = max(abs(vor_min), abs(vor_max))/G.radius + initialize!(scheme,diagn.k,G,L,vor_abs_max) +end + +"""$(TYPEDSIGNATURES) +Apply horizontal diffusion to a 2D field `A` in spectral space by updating its tendency `tendency` +with an implicitly calculated diffusion term. The implicit diffusion of the next time step is split +into an explicit part `∇²ⁿ_expl` and an implicit part `∇²ⁿ_impl`, such that both can be calculated +in a single forward step by using `A` as well as its tendency `tendency`.""" +function horizontal_diffusion!( tendency::LowerTriangularMatrix{Complex{NF}}, # tendency of a + A::LowerTriangularMatrix{Complex{NF}}, # spectral horizontal field + ∇²ⁿ_expl::AbstractVector{NF}, # explicit spectral damping + ∇²ⁿ_impl::AbstractVector{NF} # implicit spectral damping + ) where {NF<:AbstractFloat} + lmax,mmax = size(tendency) # 1-based + @boundscheck size(tendency) == size(A) || throw(BoundsError) + @boundscheck lmax <= length(∇²ⁿ_expl) == length(∇²ⁿ_impl) || throw(BoundsError) + + lm = 0 + @inbounds for m in 1:mmax # loops over all columns/order m + for l in m:lmax-1 # but skips the lmax+2 degree (1-based) + lm += 1 # single index lm corresponding to harmonic l,m + tendency[lm] = (tendency[lm] + ∇²ⁿ_expl[l]*A[lm])*∇²ⁿ_impl[l] + end + lm += 1 # skip last row for scalar quantities + end +end + +"""$(TYPEDSIGNATURES) +Apply horizontal diffusion to vorticity in the Barotropic models.""" +function horizontal_diffusion!( diagn::DiagnosticVariablesLayer, + progn::PrognosticLayerTimesteps, + model::Barotropic, + lf::Int=1) # leapfrog index used (2 is unstable) + + HD = model.horizontal_diffusion + ∇²ⁿ = HD.∇²ⁿ_2D + ∇²ⁿ_implicit = HD.∇²ⁿ_2D_implicit + + # Barotropic model diffuses vorticity (only variable) + (;vor) = progn.timesteps[lf] + (;vor_tend) = diagn.tendencies + horizontal_diffusion!(vor_tend,vor,∇²ⁿ,∇²ⁿ_implicit) +end + +"""$(TYPEDSIGNATURES) +Apply horizontal diffusion to vorticity and diffusion in the ShallowWater models.""" +function horizontal_diffusion!( progn::PrognosticLayerTimesteps, + diagn::DiagnosticVariablesLayer, + model::ShallowWater, + lf::Int=1) # leapfrog index used (2 is unstable) + + HD = model.horizontal_diffusion + ∇²ⁿ = HD.∇²ⁿ_2D + ∇²ⁿ_implicit = HD.∇²ⁿ_2D_implicit + + # ShallowWater model diffuses vorticity and divergence + (;vor,div) = progn.timesteps[lf] + (;vor_tend,div_tend) = diagn.tendencies + horizontal_diffusion!(vor_tend,vor,∇²ⁿ,∇²ⁿ_implicit) + horizontal_diffusion!(div_tend,div,∇²ⁿ,∇²ⁿ_implicit) +end + +"""$(TYPEDSIGNATURES) +Apply horizontal diffusion applied to vorticity, diffusion and temperature +in the PrimitiveEquation models. Uses the constant diffusion for temperature +but possibly adaptive diffusion for vorticity and divergence.""" +function horizontal_diffusion!( progn::PrognosticLayerTimesteps, + diagn::DiagnosticVariablesLayer, + model::PrimitiveEquation, + lf::Int=1) # leapfrog index used (2 is unstable) + + HD = model.horizontal_diffusion + initialize!(HD,diagn,model.geometry,model.time_stepping) + k = diagn.k # current layer k + ∇²ⁿ = HD.∇²ⁿ[k] # now pick operators at k + ∇²ⁿ_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 + horizontal_diffusion!(vor_tend,vor,∇²ⁿ,∇²ⁿ_implicit) + horizontal_diffusion!(div_tend,div,∇²ⁿ,∇²ⁿ_implicit) + + # but use the weaker normal diffusion for temperature + ∇²ⁿ = HD.∇²ⁿ_2D + ∇²ⁿ_implicit = HD.∇²ⁿ_2D_implicit + horizontal_diffusion!(temp_tend,temp,∇²ⁿ,∇²ⁿ_implicit) +end \ No newline at end of file diff --git a/src/dynamics/implicit.jl b/src/dynamics/implicit.jl index 094e2eeec..1dc3fe7c0 100644 --- a/src/dynamics/implicit.jl +++ b/src/dynamics/implicit.jl @@ -1,45 +1,62 @@ -# BAROTROPIC MODEL -""" - nothing = initialize_implicit!(::Real,::Barotropic) - -Just passes, as implicit terms are not used in the barotropic model.""" -initialize_implicit!(::Barotropic,::DiagnosticVariables,::Real) = nothing - -# !!! structs are defined in define_implicit.jl !!! +# BAROTROPIC MODEL (no implicit needed) +struct NoImplicit{NF} <: AbstractImplicit{NF} end +NoImplicit(SG::SpectralGrid) = NoImplicit{SG.NF}() +initialize!(I::NoImplicit,dt::Real,::DiagnosticVariables,::ModelSetup) = nothing # SHALLOW WATER MODEL """ - I = Implicit(P::Parameters{<:ShallowWater}) +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} + + # DIMENSIONS + trunc::Int + + "coefficient for semi-implicit computations to filter gravity waves" + α::Float64 = 1 + + # PRECOMPUTED ARRAYS, to be initiliased with initialize! + H₀::Base.RefValue{NF} = Ref(zero(NF)) # layer_thicknes + ξH₀::Base.RefValue{NF} = Ref(zero(NF)) # = 2αΔt*layer_thickness, store in RefValue for mutability + g∇²::Vector{NF} = zeros(NF,trunc+2) # = gravity*eigenvalues + ξg∇²::Vector{NF} = zeros(NF,trunc+2) # = 2αΔt*gravity*eigenvalues + S⁻¹::Vector{NF} = zeros(NF,trunc+2) # = 1 / (1-ξH₀*ξg∇²), implicit operator +end -Zero generator function for an `ImplicitShallowWater` struct, which holds precomputed arrays -for the implicit correction in the shallow water model. Actual precomputation happens in -initialize_implicit!.""" -function Implicit(P::Parameters{<:ShallowWater}) # shallow water model only - - (;NF,trunc) = P +""" +$(TYPEDSIGNATURES) +Generator using the resolution from `spectral_grid`.""" +function ImplicitShallowWater(spectral_grid::SpectralGrid,kwargs...) + (;NF,trunc) = spectral_grid + return ImplicitShallowWater{NF}(;trunc,kwargs...) +end - # 0-initialize, actual initialization depends on time step, done in initialize_implicit! - ξH₀ = [zero(NF)] # time step ξ, layer thickness at rest H₀ (in vec for mutability) - g∇² = zeros(NF,trunc+2) # gravity times Laplace operator - ξg∇² = zeros(NF,trunc+2) # time step ξ times gravity times Laplace operator - S⁻¹ = zeros(NF,trunc+2) # combined operator to be inverted +function Base.show(io::IO,I::ImplicitShallowWater) + print(io,"$(typeof(I)):") + keys = (:trunc,:α) + for key in keys + val = getfield(I,key) + s = "\n $key::$(typeof(val)) = $val" + print(io,s) + end +end - return ImplicitShallowWater(ξH₀,g∇²,ξg∇²,S⁻¹) +# function barrier to unpack the constants struct for shallow water +function initialize!(I::ImplicitShallowWater,dt::Real,::DiagnosticVariables,model::ShallowWater) + initialize!(I,dt,model.constants) end """ - initialize_implicit!(dt::Real,M::BarotropicModel) - -Update the implicit terms in `M` for the shallow water model as they depend on the time step `dt`.""" -function initialize_implicit!( model::ShallowWater, # update Implicit struct in model - ::DiagnosticVariables, - dt::Real) # time step +$(TYPEDSIGNATURES) +Update the implicit terms in `implicit` for the shallow water model as they depend on the time step `dt`.""" +function initialize!( implicit::ImplicitShallowWater, + dt::Real, # time step used [s] + constants::DynamicsConstants) - (;implicit_α) = model.parameters # = [0,0.5,1], time step fraction for implicit - (;eigenvalues) = model.spectral_transform # = -l*(l+1), degree l of harmonics - (;ξH₀,g∇²,ξg∇²,S⁻¹) = model.implicit # pull precomputed arrays to be updated - (;layer_thickness) = model.constants # shallow water layer thickness [m] - (;gravity) = model.constants # gravitational acceleration [m/s²] + (;α,H₀,ξH₀,g∇²,ξg∇²,S⁻¹) = implicit # precomputed arrays to be updated + (;gravity,layer_thickness) = constants # shallow water layer thickness [m] + # gravitational acceleration [m/s²] # implicit time step between i-1 and i+1 # α = 0 means the gravity wave terms are evaluated at i-1 (forward) @@ -49,32 +66,29 @@ function initialize_implicit!( model::ShallowWater, # update Implicit struct # α = 0.5 slows gravity waves and prevents them from amplifying # α > 0.5 will dampen the gravity waves within days to a few timesteps (α=1) - ξ = implicit_α*dt # new implicit timestep ξ = α*dt = 2αΔt (for leapfrog) from input dt - ξH₀[1] = ξ*layer_thickness # update ξ*H₀ with new ξ, in vec for mutability + ξ = α*dt # new implicit timestep ξ = α*dt = 2αΔt (for leapfrog) from input dt + H₀[] = layer_thickness # update H₀ the undisturbed layer thickness without mountains + ξH₀[] = ξ*layer_thickness # update ξ*H₀ with new ξ, in RefValue for mutability # loop over degree l of the harmonics (implicit terms are independent of order m) - @inbounds for l in eachindex(g∇²,ξg∇²,S⁻¹,eigenvalues) - g∇²[l] = gravity*eigenvalues[l] # doesn't actually change with dt + @inbounds for l in eachindex(g∇²,ξg∇²,S⁻¹) + eigenvalue = -l*(l-1) # =∇², with without 1/radius², 1-based -l*l(l+1) → -l*(l-1) + g∇²[l] = gravity*eigenvalue # doesn't actually change with dt ξg∇²[l] = ξ*g∇²[l] # update ξg∇² with new ξ - S⁻¹[l] = inv(1 - ξH₀[1]*ξg∇²[l]) # update 1/(1-ξ²gH₀∇²) with new ξ + S⁻¹[l] = inv(1 - ξH₀[]*ξg∇²[l]) # update 1/(1-ξ²gH₀∇²) with new ξ end end """ - implicit_correction!( diagn::DiagnosticVariablesLayer, - progn::PrognosticLayerTimesteps, - surface::SurfaceVariables, - pres::PrognosticSurfaceTimesteps, - M::ShallowWaterModel) - -Apply correction to the tendencies in `diag` to prevent the gravity waves from amplifying. -The correction is implicitly evaluated using the parameter `implicit_α` to switch between +$(TYPEDSIGNATURES) +Apply correction to the tendencies in `diagn` to prevent the gravity waves from amplifying. +The correction is implicitly evaluated using the parameter `implicit.α` to switch between forward, centered implicit or backward evaluation of the gravity wave terms.""" function implicit_correction!( diagn::DiagnosticVariablesLayer{NF}, progn::PrognosticLayerTimesteps{NF}, diagn_surface::SurfaceVariables{NF}, progn_surface::PrognosticSurfaceTimesteps{NF}, - model::ShallowWater) where NF + implicit::ImplicitShallowWater) where NF (;div_tend) = diagn.tendencies # divergence tendency div_old = progn.timesteps[1].div # divergence at t @@ -83,16 +97,11 @@ function implicit_correction!( diagn::DiagnosticVariablesLayer{NF}, pres_new = progn_surface.timesteps[2].pres # pressure/η at t+dt (;pres_tend) = diagn_surface # tendency of pressure/η - (;g∇²,ξg∇²,S⁻¹) = model.implicit - ξH₀ = model.implicit.ξH₀[1] # unpack as it's stored in a vec for mutation - H₀ = model.constants.layer_thickness + (;g∇²,ξg∇²,S⁻¹) = implicit + H₀ = implicit.H₀[] # unpack as it's stored in a RefValue for mutation + ξH₀ = implicit.ξH₀[] # unpack as it's stored in a RefValue for mutation - @boundscheck size(div_old) == size(div_new) || throw(BoundsError) - @boundscheck size(pres_old) == size(pres_new) || throw(BoundsError) - @boundscheck size(div_tend) == size(pres_new) || throw(BoundsError) - @boundscheck size(pres_tend) == size(div_new) || throw(BoundsError) lmax,mmax = size(div_tend) .- (2,1) - @boundscheck length(S⁻¹) == lmax+2 || throw(BoundsError) @boundscheck length(ξg∇²) == lmax+2 || throw(BoundsError) @boundscheck length(g∇²) == lmax+2 || throw(BoundsError) @@ -117,49 +126,111 @@ function implicit_correction!( diagn::DiagnosticVariablesLayer{NF}, end end -# PRIMITIVE EQUATION MODEL """ - I = Implicit(P::Parameters{<:PrimitiveEquation}) +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} + + # DIMENSIONS + "spectral resolution" + trunc::Int + + "number of vertical levels" + nlev::Int + + # PARAMETERS + "time-step coefficient: 0=explicit, 0.5=centred implicit, 1=backward implicit" + α::Float64 = 1 + + "recalculate the implicit terms occasionally based on the current temperature profile?" + adaptive::Bool = true -Zero generator function for an `ImplicitPrimitiveEq` struct, which holds precomputed arrays for -the implicit correction in the primitive equation model. Actual precomputation happens in initialize_implicit!.""" -function Implicit(P::Parameters{<:PrimitiveEquation}) # primitive equation only + "recalculate operators based on new temperature profile every `recalculate` time steps" + recalculate::Int = adaptive ? 100 : typemax(Int) + + # PRECOMPUTED ARRAYS, to be initiliased with initialize! + "vertical temperature profile" + temp_profile::Vector{NF} = zeros(NF,nlev) + + "time step 2α*Δt packed in RefValue for mutability" + ξ::Base.RefValue{NF} = Ref{NF}(0) + + "divergence: operator for the geopotential calculation" + R::Matrix{NF} = zeros(NF,nlev,nlev) + + "divergence: the -RdTₖ∇² term excl the eigenvalues from ∇² for divergence" + U::Vector{NF} = zeros(NF,nlev) - (;NF,trunc,nlev) = P - - # initialize with zeros only, actual initialization depends on time step, done in initialize_implicit! - ξ = Ref{NF}(0) # time step 2α*Δt packed in RefValue for mutability - R = zeros(NF,nlev,nlev) # divergence: operator for the geopotential calculation - U = zeros(NF,nlev) # divergence: the -RdTₖ∇² term excl the eigenvalues from ∇² for divergence - L = zeros(NF,nlev,nlev) # temperature: operator for the TₖD + κTₖDlnps/Dt term - W = zeros(NF,nlev) # pressure: vertical averaging of the -D̄ term in the log surface pres equation + "temperature: operator for the TₖD + κTₖDlnps/Dt term" + L::Matrix{NF} = zeros(NF,nlev,nlev) + + "pressure: vertical averaging of the -D̄ term in the log surface pres equation" + W::Vector{NF} = zeros(NF,nlev) - L0 = zeros(NF,nlev) # components to construct L, 1/ 2Δσ - L1 = zeros(NF,nlev,nlev) # vert advection term in the temperature equation (below+above) - L2 = zeros(NF,nlev) # factor in front of the div_sum_above term - L3 = zeros(NF,nlev,nlev) # _sum_above operator itself - L4 = zeros(NF,nlev) # factor in front of div term in Dlnps/Dt - - S = zeros(NF,nlev,nlev) # for every l the matrix to be inverted  - S⁻¹ = zeros(NF,trunc+1,nlev,nlev) # combined inverted operator: S = 1 - ξ²(RL + UW) - return ImplicitPrimitiveEq(ξ,R,U,L,W,L0,L1,L2,L3,L4,S,S⁻¹) + "components to construct L, 1/ 2Δσ" + L0::Vector{NF} = zeros(NF,nlev) + + "vert advection term in the temperature equation (below+above)" + L1::Matrix{NF} = zeros(NF,nlev,nlev) + + "factor in front of the div_sum_above term" + L2::Vector{NF} = zeros(NF,nlev) + + "_sum_above operator itself" + L3::Matrix{NF} = zeros(NF,nlev,nlev) + + "factor in front of div term in Dlnps/Dt" + L4::Vector{NF} = zeros(NF,nlev) + + "for every l the matrix to be inverted" + S::Matrix{NF} = zeros(NF,nlev,nlev) + + "combined inverted operator: S = 1 - ξ²(RL + UW)" + S⁻¹::Array{NF,3} = zeros(NF,trunc+1,nlev,nlev) end -function initialize_implicit!( model::PrimitiveEquation, - diagn::DiagnosticVariables, - dt::Real) # the scaled time step radius*dt - - (;S,S⁻¹,L,R,U,W,L0,L1,L2,L3,L4)= model.implicit - (;nlev, σ_levels_full, σ_levels_thick) = model.geometry - (;Δp_geopot_half, Δp_geopot_full, σ_lnp_A, σ_lnp_B) = model.geometry - (;R_dry, κ) = model.constants - (;eigenvalues, lmax) = model.spectral_transform - α = model.parameters.implicit_α - - # use an occasionally updated vertical temperature profile - (;temp_profile) = diagn - for t in temp_profile # return immediately if temp_profile contains - if !isfinite(t) return nothing end # NaRs, model blew up in that case +"""$(TYPEDSIGNATURES) +Generator using the resolution from SpectralGrid.""" +function ImplicitPrimitiveEq(spectral_grid::SpectralGrid,kwargs...) + (;NF,trunc,nlev) = spectral_grid + return ImplicitPrimitiveEq{NF}(;trunc,nlev,kwargs...) +end + +function Base.show(io::IO,I::ImplicitPrimitiveEq) + print(io,"$(typeof(I)):") + keys = (:trunc,:nlev,:α,:adaptive,:recalculate) + for key in keys + val = getfield(I,key) + s = "\n $key::$(typeof(val)) = $val" + print(io,s) + end +end + +# function barrier to unpack the constants struct for primitive eq models +function initialize!(I::ImplicitPrimitiveEq,dt::Real,diagn::DiagnosticVariables,model::PrimitiveEquation) + initialize!(I,dt,diagn,model.geometry,model.constants) +end + +"""$(TYPEDSIGNATURES) +Initialize the implicit terms for the PrimitiveEquation models.""" +function initialize!( implicit::ImplicitPrimitiveEq, + dt::Real, # the scaled time step radius*dt + diagn::DiagnosticVariables, + geometry::Geometry, + constants::DynamicsConstants) + + (;trunc, nlev, α,temp_profile,S,S⁻¹,L,R,U,W,L0,L1,L2,L3,L4) = implicit + (;σ_levels_full, σ_levels_thick) = geometry + (;R_dry, κ, Δp_geopot_half, Δp_geopot_full, σ_lnp_A, σ_lnp_B) = constants + + if implicit.adaptive # use current vertical temperature profile + for k in 1:nlev + temp_profile[k] = diagn.layers[k].temp_average[] # return immediately if temp_profile contains + if !isfinite(temp_profile[k]) return nothing end # NaRs, model blew up in that case + end + else # or use reference profile + temp_profile .= constants.temp_ref_profile end # set up R, U, L, W operators from @@ -175,8 +246,8 @@ function initialize_implicit!( model::PrimitiveEquation, # R, U, L, W are linear operators that are therefore defined here and inverted # to obtain δD first, and then δT and δlnps through substitution - ξ = α*dt # dt = 2Δt for leapfrog, but = Δt, Δ/2 in first_timesteps! - model.implicit.ξ[] = ξ # also store in Implicit struct + ξ = α*dt # dt = 2Δt for leapfrog, but = Δt, Δ/2 in first_timesteps! + implicit.ξ[] = ξ # also store in Implicit struct # DIVERGENCE OPERATORS (called g in Hoskins and Simmons 1975, eq 11 and Appendix 1) @inbounds for k in 1:nlev # vertical geopotential integration as matrix operator @@ -221,8 +292,9 @@ function initialize_implicit!( model::PrimitiveEquation, # δD = SG, with G = G_D + ξRG_T + ξUG_lnps and the operator S # S = 1 - ξ²(RL + UW) that has to be inverted to obtain δD from the Gs I = LinearAlgebra.I(nlev) - @inbounds for l in 1:lmax+1 - S .= I .- ξ^2*eigenvalues[l]*(R*L .+ U*W') + @inbounds for l in 1:trunc+1 + eigenvalue = -l*(l-1) # 1-based, -l*(l+1) → -l*(l-1) + S .= I .- ξ^2*eigenvalue*(R*L .+ U*W') # inv(S) but saving memory: luS = LinearAlgebra.lu!(S) # in-place LU decomposition (overwriting S) @@ -233,46 +305,35 @@ function initialize_implicit!( model::PrimitiveEquation, end end -function initialize_implicit!( model::PrimitiveEquation, - diagn::DiagnosticVariables, - progn::PrognosticVariables, - dt::Real, - i::Integer, - lf::Integer) +"""$(TYPEDSIGNATURES) +Reinitialize implicit occasionally based on time step `i` and `implicit.recalculate`.""" +function initialize!( + implicit::ImplicitPrimitiveEq, + i::Integer, + dt::Real, # the scaled time step radius*dt + diagn::DiagnosticVariables, + geometry::Geometry, + constants::DynamicsConstants +) # only reinitialize occasionally, otherwise exit immediately - i % model.parameters.recalculate_implicit == 0 || return nothing - temperature_profile!(diagn,progn,model,lf) - initialize_implicit!(model,diagn,dt) + i % implicit.recalculate == 0 || return nothing + initialize!(implicit,dt,diagn,geometry,constants) end -# just pass for Barotropic and ShallowWater -temperature_profile!(::DiagnosticVariables,::PrognosticVariables,::ModelSetup,lf::Integer=1) = nothing - -function temperature_profile!( diagn::DiagnosticVariables, - progn::PrognosticVariables, - model::PrimitiveEquation, - lf::Integer=1) - (;temp_profile) = diagn - (;norm_sphere) = model.spectral_transform - - @inbounds for k in 1:diagn.nlev - temp_profile[k] = real(progn.layers[k].timesteps[lf].temp[1])/norm_sphere - end -end - -function implicit_correction!( diagn::DiagnosticVariables{NF}, - progn::PrognosticVariables, - model::PrimitiveEquation, - ) where NF +"""$(TYPEDSIGNATURES) +Apply the implicit corrections to dampen gravity waves in the primitive equation models.""" +function implicit_correction!( + diagn::DiagnosticVariables, + implicit::ImplicitPrimitiveEq, + progn::PrognosticVariables, +) # escape immediately if explicit - model.parameters.implicit_α == 0 && return nothing + implicit.α == 0 && return nothing - (;nlev) = model.geometry - (;eigenvalues, lmax, mmax) = model.spectral_transform - (;Δp_geopot_half, Δp_geopot_full) = model.geometry # = R*Δlnp on half or full levels - (;S⁻¹,R,U,L,W) = model.implicit - ξ = model.implicit.ξ[] + # (;Δp_geopot_half, Δp_geopot_full) = model.geometry # = R*Δlnp on half or full levels + (;nlev, trunc, S⁻¹, R, U, L, W) = implicit + ξ = implicit.ξ[] # MOVE THE IMPLICIT TERMS OF THE TEMPERATURE EQUATION FROM TIME STEP i TO i-1 # geopotential and linear pressure gradient (divergence equation) are already evaluated at i-1 @@ -320,11 +381,12 @@ function implicit_correction!( diagn::DiagnosticVariables{NF}, # 2. the G = G_D + ξRG_T + ξUG_lnps terms using geopot from above lm = 0 - @inbounds for m in 1:mmax+1 # loops over all columns/order m - for l in m:lmax+1 # but skips the lmax+2 degree (1-based) + @inbounds for m in 1:trunc+1 # loops over all columns/order m + for l in m:trunc+1 # but skips the lmax+2 degree (1-based) lm += 1 # single index lm corresponding to harmonic l,m # ∇² not part of U so *eigenvalues here - G[lm] = div_tend[lm] + ξ*eigenvalues[l]*(U[k]*pres_tend[lm] + geopot[lm]) + eigenvalue = -l*(l-1) # 1-based, -l*(l+1) → -l*(l-1) + G[lm] = div_tend[lm] + ξ*eigenvalue*(U[k]*pres_tend[lm] + geopot[lm]) end lm += 1 # skip last row, LowerTriangularMatrices are of size lmax+2 x mmax+1 end @@ -342,8 +404,8 @@ function implicit_correction!( diagn::DiagnosticVariables{NF}, G = diagn.layers[r].dynamics_variables.a # reuse work arrays lm = 0 - for m in 1:mmax+1 # loops over all columns/order m - for l in m:lmax+1 # but skips the lmax+2 degree (1-based) + for m in 1:trunc+1 # loops over all columns/order m + for l in m:trunc+1 # but skips the lmax+2 degree (1-based) lm += 1 # single index lm corresponding to harmonic l,m div_tend[lm] += S⁻¹[l,k,r]*G[lm] end diff --git a/src/dynamics/initial_conditions.jl b/src/dynamics/initial_conditions.jl index 7ecdeeb75..c0c08f701 100644 --- a/src/dynamics/initial_conditions.jl +++ b/src/dynamics/initial_conditions.jl @@ -1,96 +1,125 @@ -""" - prognostic_variables = initial_conditions(M::ModelSetup) +# default initial conditions by model +initial_conditions_default(::Type{<:Barotropic}) = StartWithVorticity() +initial_conditions_default(::Type{<:ShallowWater}) = ZonalJet() +initial_conditions_default(::Type{<:PrimitiveEquation}) = ZonalWind() +""" +$(TYPEDSIGNATURES) Allocate the prognostic variables and then set to initial conditions.""" function initial_conditions(model::ModelSetup) - - progn = allocate_prognostic_variables(model) # allocate variables in any case - IC = model.parameters.initial_conditions # initial conditions struct - initial_conditions!(IC,progn,model) # dispatch to initial conditions + progn = allocate(PrognosticVariables,model.spectral_grid) # allocate variables in any case + IC = model.initial_conditions # initial conditions struct + initial_conditions!(progn,IC,model) # dispatch to initial conditions return progn end -function allocate_prognostic_variables(model::ModelSetup) - - (;NF) = model.parameters - (;nlev) = model.geometry - (;lmax, mmax) = model.spectral_transform - - return zeros(PrognosticVariables{NF},model,lmax,mmax,nlev) +""" +$(TYPEDSIGNATURES)""" +function allocate(::Type{PrognosticVariables},spectral_grid::SpectralGrid{Model}) where Model + (;NF,trunc,nlev) = spectral_grid + return zeros(PrognosticVariables{NF},Model,trunc,nlev) end -struct StartFromRest <: InitialConditions end +@with_kw struct StartFromRest <: InitialConditions + pressure_on_orography::Bool = false +end -function initial_conditions!( ::StartFromRest, - progn::PrognosticVariables, +function initial_conditions!( progn::PrognosticVariables, + initial_conditions::StartFromRest, model::ModelSetup) - return nothing + return nothing # everything remains zero end -function initial_conditions!( ::StartFromRest, - progn::PrognosticVariables, +function initial_conditions!( progn::PrognosticVariables, + initial_conditions::StartFromRest, model::PrimitiveEquation) homogeneous_temperature!(progn,model) - model.parameters.pressure_on_orography && pressure_on_orography!(progn,model) + initial_conditions.pressure_on_orography && pressure_on_orography!(progn,model) # TODO initialise humidity end -struct StartWithVorticity <: InitialConditions end +"""Start with random vorticity as initial conditions +$(TYPEDFIELDS)""" +@with_kw struct StartWithVorticity <: InitialConditions -function initial_conditions!( ::StartWithVorticity, - progn::PrognosticVariables, - model::ModelSetup) + "Power law the vorticity should be spectrally distributed by" + power_law::Float64 = -3 - (;radius) = model.geometry + "(approximate) amplitude in [1/s], used as standard deviation of spherical harmonic coefficients" + amplitude::Float64 = 1e-5 +end - mmax = min(15,progn.mmax+1) # perturb only larger modes - lmax = min(15,progn.lmax+1) +""" +$(TYPEDSIGNATURES) +Start with random vorticity as initial conditions""" +function initial_conditions!( progn::PrognosticVariables{NF}, + initial_conditions::StartWithVorticity, + model::ModelSetup) where NF - ξ = randn(Complex{model.parameters.NF},mmax,lmax) + lmax = progn.trunc+1 + power = initial_conditions.power_law + 1 # +1 as power is summed of orders m + ξ = randn(Complex{NF},lmax,lmax)*convert(NF,initial_conditions.amplitude) for progn_layer in progn.layers - - # zonal wind - progn_layer.timesteps[1].vor[4,1] = 80/radius - progn_layer.timesteps[1].vor[6,1] = -160/radius - progn_layer.timesteps[1].vor[8,1] = 80/radius - - # perturbation - for m in 2:min(15,progn.mmax+1) - for l in m:min(15,progn.lmax+1) - progn_layer.timesteps[1].vor[l,m] = 10/radius*ξ[l,m] + for m in 1:lmax + for l in m:lmax + progn_layer.timesteps[1].vor[l,m] = ξ[l,m]*l^power end end + # don't perturb l=m=0 mode to have zero mean + progn_layer.timesteps[1].vor[1] = 0 end end """ - Z = ZonalJet(;kwargs...) <: InitialConditions - +$(TYPEDSIGNATURES) Create a struct that contains all parameters for the Galewsky et al, 2004 zonal jet -intitial conditions for the shallow water model. Default values as in Galewsky.""" -Base.@kwdef struct ZonalJet <: InitialConditions - # jet - latitude = 45 # degrees north [˚N] - width = (1/4-1/7)*180 # ≈ 19.29˚ as in Galewsky et al, 2004 - umax = 80 # [m/s] +intitial conditions for the shallow water model. Default values as in Galewsky. +$(TYPEDFIELDS)""" +@with_kw struct ZonalJet <: InitialConditions + "jet latitude [˚N]" + latitude::Float64 = 45 - # perturbation - perturb_lat = latitude # [˚N], position in jet by default - perturb_lon = 270 # [˚E] - perturb_xwidth = 1/3*360/2π # ≈ 19.1˚E zonal extent [˚E] - perturb_ywidth = 1/15*360/2π # ≈ 3.8˚N meridional extent [˚N] - perturb_height = 120 # amplitude [m] + "jet width [˚], default ≈ 19.29˚" + width::Float64 = (1/4-1/7)*180 + + "jet maximum velocity [m/s]" + umax::Float64 = 80 + + "perturbation latitude [˚N], position in jet by default" + perturb_lat::Float64 = latitude + + "perturbation longitude [˚E]" + perturb_lon::Float64 = 270 + + "perturbation zonal extent [˚], default ≈ 19.1˚" + perturb_xwidth::Float64 = 1/3*360/2π + + "perturbation meridinoal extent [˚], default ≈ 3.8˚" + perturb_ywidth::Float64 = 1/15*360/2π + + "perturbation amplitude [m]" + perturb_height::Float64 = 120 end -"""Initial conditions from Galewsky, 2004, Tellus""" -function initial_conditions!( coefs::ZonalJet, - progn::PrognosticVariables, +function Base.show(io::IO,IC::InitialConditions) + print(io,"$(typeof(IC)) <: InitialConditions:") + for key in propertynames(IC) + val = getfield(IC,key) + print(io,"\n $key::$(typeof(val)) = $val") + end +end + +""" +$(TYPEDSIGNATURES) +Initial conditions from Galewsky, 2004, Tellus""" +function initial_conditions!( progn::PrognosticVariables, + initial_conditions::ZonalJet, model::ShallowWater) - (;latitude, width, umax) = coefs # for jet - (;perturb_lat, perturb_lon, perturb_xwidth, # for perturbation - perturb_ywidth, perturb_height) = coefs + (;latitude, width, umax) = initial_conditions # for jet + (;perturb_lat, perturb_lon, perturb_xwidth, # for perturbation + perturb_ywidth, perturb_height) = initial_conditions θ₀ = (latitude-width)/360*2π # southern boundary of jet [radians] θ₁ = (latitude+width)/360*2π # northern boundary of jet @@ -101,7 +130,8 @@ function initial_conditions!( coefs::ZonalJet, β = perturb_ywidth*2π/360 # meridional extent of interface perturbation [radians] λ = perturb_lon*2π/360 # perturbation longitude [radians] - (;radius, rotation, gravity) = model.parameters.planet + (;rotation, gravity) = model.planet + (;radius) = model.spectral_grid # always create on F64 grid then convert to spectral and interpolate there Grid = FullGaussianGrid @@ -147,7 +177,7 @@ function initial_conditions!( coefs::ZonalJet, # interpolate in spectral space to desired resolution (;lmax,mmax) = model.spectral_transform - (;NF) = model.parameters + (;NF) = model.spectral_grid u = spectral_truncation(complex(NF),u,lmax+1,mmax) # get vorticity initial conditions from curl of u,v @@ -162,39 +192,55 @@ function initial_conditions!( coefs::ZonalJet, end """ - Z = ZonalWind(;kwargs...) <: InitialConditions - +$(TYPEDSIGNATURES) Create a struct that contains all parameters for the Jablonowski and Williamson, 2006 -intitial conditions for the primitive equation model. Default values as in Jablonowski.""" -Base.@kwdef struct ZonalWind <: InitialConditions +intitial conditions for the primitive equation model. Default values as in Jablonowski. +$(TYPEDFIELDS)""" +@with_kw struct ZonalWind <: InitialConditions + "conversion from σ to Jablonowski's ηᵥ-coordinates" + η₀::Float64 = 0.252 - # vertical - η₀ = 0.252 # conversion from σ to Jablonowski's ηᵥ-coordinates - u₀ = 35 # max amplitude of zonal wind [m/s] + "max amplitude of zonal wind [m/s]" + u₀::Float64 = 35 + + # PERTURBATION + "perturbation centred at [˚N]" + perturb_lat::Float64 = 40 + + "perturbation centred at [˚E]" + perturb_lon::Float64 = 20 + + "perturbation strength [m/s]" + perturb_uₚ::Float64 = 1 - # perturbation - perturb_lat = 40 # Gaussian profile perturbation centred at [˚N] - perturb_lon = 20 # and [˚E] - perturb_uₚ = 1 # strength of perturbation [m/s] - perturb_radius = 1/10 # radius of Gaussian perturbation in units of Earth's radius [1] - - # temperature - ΔT = 0 # temperature difference used for stratospheric lapse rate [K] - # Jablonowski uses ΔT = 4.8e5 [K] - Tmin = 200 # minimum temperature [K] of profile + "radius of Gaussian perturbation in units of Earth's radius [1]" + perturb_radius::Float64 = 1/10 + + # TERMPERATURE + "temperature difference used for stratospheric lapse rate [K], Jablonowski uses ΔT = 4.8e5 [K]" + ΔT::Float64 = 0 + + "minimum temperature [K] of profile" + Tmin::Float64 = 200 + + # PRESSURE + "initialize pressure given the `atmosphere.lapse_rate` on orography?" + pressure_on_orography::Bool = false end -"""Initial conditions from Jablonowski and Williamson, 2006, QJR Meteorol. Soc""" -function initial_conditions!( coefs::ZonalWind, - progn::PrognosticVariables{NF}, +""" +$(TYPEDSIGNATURES) +Initial conditions from Jablonowski and Williamson, 2006, QJR Meteorol. Soc""" +function initial_conditions!( progn::PrognosticVariables{NF}, + initial_conditions::ZonalWind, model::PrimitiveEquation) where NF - (;u₀, η₀, ΔT, Tmin) = coefs - (;perturb_lat, perturb_lon, perturb_uₚ, perturb_radius) = coefs - (;temp_ref, R_dry, lapse_rate, pres_ref) = model.parameters - (;radius, rotation, gravity) = model.parameters.planet - (;σ_tropopause, pressure_on_orography) = model.parameters - (;σ_levels_full, Grid, nlat_half, nlev) = model.geometry + (;u₀, η₀, ΔT, Tmin, pressure_on_orography) = initial_conditions + (;perturb_lat, perturb_lon, perturb_uₚ, perturb_radius) = initial_conditions + (;temp_ref, R_dry, lapse_rate, pres_ref, σ_tropopause) = model.atmosphere + (;radius, Grid, nlat_half, nlev) = model.spectral_grid + (;rotation, gravity) = model.planet + (;σ_levels_full) = model.geometry (;norm_sphere) = model.spectral_transform φ, λ = model.geometry.latds, model.geometry.londs @@ -294,49 +340,60 @@ function initial_conditions!( coefs::ZonalWind, pressure_on_orography && pressure_on_orography!(progn,model) end -struct StartFromFile <: InitialConditions end +""" +Restart from a previous SpeedyWeather.jl simulation via the restart file restart.jld2 +Applies interpolation in the horizontal but not in the vertical. restart.jld2 is +identified by +$(TYPEDFIELDS)""" +@with_kw struct StartFromFile <: InitialConditions + "path for restart file" + path::String = pwd() + + "`run_id` of restart file in `run_????/restart.jld2`" + id::Union{String,Int} = 1 +end -function initial_conditions!( ::StartFromFile, - progn_new::PrognosticVariables, +""" +$(TYPEDSIGNATURES) +Restart from a previous SpeedyWeather.jl simulation via the restart file restart.jld2 +Applies interpolation in the horizontal but not in the vertical.""" +function initial_conditions!( progn_new::PrognosticVariables, + initial_conditions::StartFromFile, model::ModelSetup) - (; restart_path, restart_id ) = model.parameters + (; path, id ) = initial_conditions - restart_file = jldopen(joinpath(restart_path,string("run-",run_id_string(restart_id)),"restart.jld2")) + restart_file = jldopen(joinpath(path,string("run_",run_id_to_string(id)),"restart.jld2")) progn_old = restart_file["prognostic_variables"] - version = restart_file["version"] # currently unused, TODO check for compat with version - time = restart_file["time"] # currently unused - + # version = restart_file["version"] # currently unused + model.clock.time = restart_file["time"] # synchronize clocks return copy!(progn_new, progn_old) end function homogeneous_temperature!( progn::PrognosticVariables, model::PrimitiveEquation) - - P = model.parameters - B = model.boundaries - G = model.geometry - S = model.spectral_transform - - (; geopot_surf ) = B.orography # spectral surface geopotential [m²/s²] (orography*gravity) + (; geopot_surf ) = model.orography # spectral surface geopotential [m²/s²] (orography*gravity) # temp_ref: Reference absolute T [K] at surface z = 0, constant lapse rate # temp_top: Reference absolute T in the stratosphere [K], lapse rate = 0 # lapse_rate: Reference temperature lapse rate -dT/dz [K/km] # gravity: Gravitational acceleration [m/s^2] # R_dry: Specific gas constant for dry air [J/kg/K] - (; temp_ref, temp_top, lapse_rate, R_dry ) = P - (; gravity ) = P.planet - (; n_stratosphere_levels, nlev ) = G # number of vertical levels used for stratosphere - (; norm_sphere ) = S # normalization of the l=m=0 spherical harmonic - - Γg⁻¹ = lapse_rate/gravity/1000 # Lapse rate scaled by gravity [K/m / (m²/s²)] - - # SURFACE TEMPERATURE - temp_surf = progn.layers[end].timesteps[1].temp # spectral temperature at k=nlev=surface - temp_surf[1] = norm_sphere*temp_ref # set global mean surface temperature + (;temp_ref, temp_top, lapse_rate, R_dry, σ_tropopause) = model.atmosphere + (;gravity) = model.planet + (;nlev, σ_levels_full) = model.geometry + (; norm_sphere ) = model.spectral_transform # normalization of the l=m=0 spherical harmonic + n_stratosphere_levels = findfirst(σ->σ>=σ_tropopause,σ_levels_full) + + # Lapse rate scaled by gravity [K/m / (m²/s²)] + Γg⁻¹ = lapse_rate/gravity/1000 # /1000 for lapse rate [K/km] → [K/m] + + # SURFACE TEMPERATURE (store in k = nlev, but it's actually surface, i.e. k=nlev+1/2) + # overwrite with lowermost layer further down + temp_surf = progn.layers[end].timesteps[1].temp # spectral temperature at k=nlev+1/2 + temp_surf[1] = norm_sphere*temp_ref # set global mean surface temperature for lm in eachharmonic(geopot_surf,temp_surf) - temp_surf[lm] -= Γg⁻¹*geopot_surf[lm] # lower temperature for higher mountains + temp_surf[lm] -= Γg⁻¹*geopot_surf[lm] # lower temperature for higher mountains end # TROPOPAUSE/STRATOSPHERE set the l=m=0 spectral coefficient (=mean value) only @@ -347,9 +404,10 @@ function homogeneous_temperature!( progn::PrognosticVariables, end # TROPOSPHERE use lapserate and vertical coordinate σ for profile - for k in n_stratosphere_levels+1:nlev + for k in n_stratosphere_levels+1:nlev # k=nlev overwrites the surface temperature + # with lowermost layer temperature temp = progn.layers[k].timesteps[1].temp - σₖᴿ = G.σ_levels_full[k]^(R_dry*Γg⁻¹) # TODO reference + σₖᴿ = σ_levels_full[k]^(R_dry*Γg⁻¹) # from hydrostatic equation for lm in eachharmonic(temp,temp_surf) temp[lm] = temp_surf[lm]*σₖᴿ @@ -357,26 +415,20 @@ function homogeneous_temperature!( progn::PrognosticVariables, end end +""" +$(TYPEDSIGNATURES) +Initialize surface pressure on orography by integrating the +hydrostatic equation with the reference temperature lapse rate.""" function pressure_on_orography!(progn::PrognosticVariables, model::PrimitiveEquation) - - P = model.parameters - B = model.boundaries - G = model.geometry - S = model.spectral_transform - - (; Grid,nlat_half ) = G - (; lmax,mmax ) = S - # temp_ref: Reference absolute T [K] at surface z = 0, constant lapse rate - # temp_top: Reference absolute T in the stratosphere [K], lapse rate = 0 # lapse_rate: Reference temperature lapse rate -dT/dz [K/km] # gravity: Gravitational acceleration [m/s^2] # R: Specific gas constant for dry air [J/kg/K] # pres_ref: Reference surface pressure [hPa] - (; temp_ref, temp_top, lapse_rate, pres_ref, R_dry ) = P - (; gravity ) = P.planet - (; orography ) = B.orography # orography on the grid + (; temp_ref, lapse_rate, pres_ref, R_dry ) = model.atmosphere + (; gravity ) = model.planet + (; orography ) = model.orography # orography on the grid Γ = lapse_rate/1000 # Lapse rate [K/km] -> [K/m] lnp₀ = log(pres_ref*100) # logarithm of reference surface pressure, *100 for [hPa] to [Pa] @@ -391,8 +443,7 @@ function pressure_on_orography!(progn::PrognosticVariables, lnp = progn.pres.timesteps[1] spectral!(lnp,lnp_grid,S) - spectral_truncation!(lnp,lmax) # set lmax+1 row to zero - + spectral_truncation!(lnp) # set lmax+1 row to zero return lnp_grid # return grid for use in initialize_humidity! end diff --git a/src/dynamics/models.jl b/src/dynamics/models.jl index 931972ada..23494c7ce 100644 --- a/src/dynamics/models.jl +++ b/src/dynamics/models.jl @@ -1,106 +1,268 @@ """ - M = BarotropicModel(::Parameters, - ::DynamicsConstants, - ::Geometry, - ::SpectralTransform, - ::HorizontalDiffusion) +$(TYPEDSIGNATURES) +Simulation is a container struct to be used with `run!(::Simulation)`. +It contains +$(TYPEDFIELDS)""" +struct Simulation{Model<:ModelSetup} + "define the current state of the model" + prognostic_variables::PrognosticVariables + "contain the tendencies and auxiliary arrays to compute them" + diagnostic_variables::DiagnosticVariables + + "all parameters, constant at runtime" + model::Model +end + +""" +$(SIGNATURES) The BarotropicModel struct holds all other structs that contain precalculated constants, -whether scalars or arrays that do not change throughout model integration. In contrast to -`ShallowWaterModel` or `PrimitiveEquation` it does not contain a `Boundaries` struct -as not needed.""" -struct BarotropicModel{NF<:AbstractFloat, D<:AbstractDevice} <: Barotropic - parameters::Parameters - constants::DynamicsConstants{NF} - geometry::Geometry{NF} - spectral_transform::SpectralTransform{NF} - horizontal_diffusion::HorizontalDiffusion{NF} - device_setup::DeviceSetup{D} +whether scalars or arrays that do not change throughout model integration. +$(TYPEDFIELDS)""" +@with_kw struct BarotropicModel{NF<:AbstractFloat, D<:AbstractDevice} <: Barotropic + "dictates resolution for many other components" + spectral_grid::SpectralGrid = SpectralGrid() + + # DYNAMICS + "contains physical and orbital characteristics" + planet::AbstractPlanet = Earth() + atmosphere::AbstractAtmosphere = EarthAtmosphere() + forcing::AbstractForcing{NF} = NoForcing(spectral_grid) + initial_conditions::InitialConditions = StartWithVorticity() + + # NUMERICS + time_stepping::TimeStepper{NF} = Leapfrog(spectral_grid) + spectral_transform::SpectralTransform{NF} = SpectralTransform(spectral_grid) + horizontal_diffusion::HorizontalDiffusion{NF} = HyperDiffusion(spectral_grid) + implicit::AbstractImplicit{NF} = NoImplicit(spectral_grid) + + # INTERNALS + clock::Clock = Clock() + geometry::Geometry{NF} = Geometry(spectral_grid) + constants::DynamicsConstants{NF} = DynamicsConstants(spectral_grid,planet,atmosphere,geometry) + device_setup::DeviceSetup{D} = DeviceSetup(CPUDevice()) + + # OUTPUT + output::AbstractOutputWriter = OutputWriter(spectral_grid) + feedback::AbstractFeedback = Feedback() end has(::Type{<:Barotropic}, var_name::Symbol) = var_name in (:vor,) default_concrete_model(::Type{Barotropic}) = BarotropicModel """ - M = ShallowWaterModel( ::Parameters, - ::DynamicsConstants, - ::Geometry, - ::SpectralTransform, - ::Boundaries, - ::HorizontalDiffusion) +$(TYPEDSIGNATURES) +Calls all `initialize!` functions for components of `model`, +except for `model.output` and `model.feedback` which are always called +at in `time_stepping!`.""" +function initialize!(model::Barotropic) + (;spectral_grid,forcing,horizontal_diffusion) = model + initialize!(forcing,model) + initialize!(horizontal_diffusion,model) + + prognostic_variables = initial_conditions(model) + diagnostic_variables = DiagnosticVariables(spectral_grid) + return Simulation(prognostic_variables,diagnostic_variables,model) +end +""" +$(SIGNATURES) The ShallowWaterModel struct holds all other structs that contain precalculated constants, -whether scalars or arrays that do not change throughout model integration.""" -struct ShallowWaterModel{NF<:AbstractFloat, D<:AbstractDevice} <: ShallowWater - parameters::Parameters - constants::DynamicsConstants{NF} - geometry::Geometry{NF} - spectral_transform::SpectralTransform{NF} - boundaries::Boundaries{NF} - horizontal_diffusion::HorizontalDiffusion{NF} - implicit::ImplicitShallowWater{NF} - device_setup::DeviceSetup{D} +whether scalars or arrays that do not change throughout model integration. +$(TYPEDFIELDS)""" +@with_kw struct ShallowWaterModel{NF<:AbstractFloat, D<:AbstractDevice} <: ShallowWater + "dictates resolution for many other components" + spectral_grid::SpectralGrid = SpectralGrid() + + # DYNAMICS + "contains physical and orbital characteristics" + planet::AbstractPlanet = Earth() + atmosphere::AbstractAtmosphere = EarthAtmosphere() + forcing::AbstractForcing{NF} = NoForcing(spectral_grid) + initial_conditions::InitialConditions = ZonalJet() + orography::AbstractOrography{NF} = EarthOrography(spectral_grid) + + # NUMERICS + time_stepping::TimeStepper{NF} = Leapfrog(spectral_grid) + spectral_transform::SpectralTransform{NF} = SpectralTransform(spectral_grid) + horizontal_diffusion::HorizontalDiffusion{NF} = HyperDiffusion(spectral_grid) + implicit::AbstractImplicit{NF} = ImplicitShallowWater(spectral_grid) + + # INTERNALS + clock::Clock = Clock() + geometry::Geometry{NF} = Geometry(spectral_grid) + constants::DynamicsConstants{NF} = DynamicsConstants(spectral_grid,planet,atmosphere,geometry) + device_setup::DeviceSetup{D} = DeviceSetup(CPUDevice()) + + # OUTPUT + output::AbstractOutputWriter = OutputWriter(spectral_grid) + feedback::AbstractFeedback = Feedback() end has(::Type{<:ShallowWater}, var_name::Symbol) = var_name in (:vor, :div, :pres) default_concrete_model(::Type{ShallowWater}) = ShallowWaterModel """ - M = PrimitiveDryCoreModel( ::Parameters, - ::DynamicsConstants, - ::Geometry, - ::SpectralTransform, - ::Boundaries, - ::HorizontalDiffusion - ::Implicit) - -The PrimitiveDryCoreModel struct holds all other structs that contain precalculated constants, -whether scalars or arrays that do not change throughout model integration.""" -struct PrimitiveDryCoreModel{NF<:AbstractFloat,D<:AbstractDevice} <: PrimitiveDryCore - parameters::Parameters - constants::DynamicsConstants{NF} - parameterization_constants::ParameterizationConstants{NF} - geometry::Geometry{NF} - spectral_transform::SpectralTransform{NF} - boundaries::Boundaries{NF} - horizontal_diffusion::HorizontalDiffusion{NF} - implicit::ImplicitPrimitiveEq{NF} - device_setup::DeviceSetup{D} +$(TYPEDSIGNATURES) +Calls all `initialize!` functions for components of `model`, +except for `model.output` and `model.feedback` which are always called +at in `time_stepping!` and `model.implicit` which is done in `first_timesteps!`.""" +function initialize!(model::ShallowWater) + (;spectral_grid,forcing,horizontal_diffusion, + orography,planet,spectral_transform,geometry) = model + + initialize!(forcing,model) + initialize!(horizontal_diffusion,model) + initialize!(orography,planet,spectral_transform,geometry) + + prognostic_variables = initial_conditions(model) + diagnostic_variables = DiagnosticVariables(spectral_grid) + return Simulation(prognostic_variables,diagnostic_variables,model) +end + +""" +$(SIGNATURES) +The PrimitiveDryModel struct holds all other structs that contain precalculated constants, +whether scalars or arrays that do not change throughout model integration. +$(TYPEDFIELDS)""" +@with_kw struct PrimitiveDryModel{NF<:AbstractFloat, D<:AbstractDevice} <: PrimitiveDry + "dictates resolution for many other components" + spectral_grid::SpectralGrid = SpectralGrid() + + # DYNAMICS + "contains physical and orbital characteristics" + planet::AbstractPlanet = Earth() + atmosphere::AbstractAtmosphere = EarthAtmosphere() + initial_conditions::InitialConditions = ZonalWind() + orography::AbstractOrography{NF} = EarthOrography(spectral_grid) + + # PHYSICS/PARAMETERIZATIONS + physics::Bool = true + boundary_layer_drag::BoundaryLayerDrag{NF} = LinearDrag(spectral_grid) + temperature_relaxation::TemperatureRelaxation{NF} = HeldSuarez(spectral_grid) + static_energy_diffusion::VerticalDiffusion{NF} = StaticEnergyDiffusion(spectral_grid) + # vertical_diffusion::VerticalDiffusion{NF} = VerticalLaplacian(spectral_grid) + + # NUMERICS + time_stepping::TimeStepper{NF} = Leapfrog(spectral_grid) + spectral_transform::SpectralTransform{NF} = SpectralTransform(spectral_grid) + horizontal_diffusion::HorizontalDiffusion{NF} = HyperDiffusion(spectral_grid) + implicit::AbstractImplicit{NF} = ImplicitPrimitiveEq(spectral_grid) + + # INTERNALS + clock::Clock = Clock() + geometry::Geometry{NF} = Geometry(spectral_grid) + constants::DynamicsConstants{NF} = DynamicsConstants(spectral_grid,planet,atmosphere,geometry) + device_setup::DeviceSetup{D} = DeviceSetup(CPUDevice()) + + # OUTPUT + output::AbstractOutputWriter = OutputWriter(spectral_grid) + feedback::AbstractFeedback = Feedback() end +has(::Type{<:PrimitiveDry}, var_name::Symbol) = var_name in (:vor, :div, :temp, :pres) +default_concrete_model(::Type{PrimitiveDry}) = PrimitiveDryModel +default_concrete_model(::Type{PrimitiveEquation}) = PrimitiveDryModel + """ - M = PrimitiveWetCoreModel( ::Parameters, - ::DynamicsConstants, - ::Geometry, - ::SpectralTransform, - ::Boundaries, - ::HorizontalDiffusion - ::Implicit) - -The PrimitiveWetCoreModel struct holds all other structs that contain precalculated constants, -whether scalars or arrays that do not change throughout model integration.""" -struct PrimitiveWetCoreModel{NF<:AbstractFloat,D<:AbstractDevice} <: PrimitiveWetCore - parameters::Parameters - constants::DynamicsConstants{NF} - parameterization_constants::ParameterizationConstants{NF} - geometry::Geometry{NF} - spectral_transform::SpectralTransform{NF} - boundaries::Boundaries{NF} - horizontal_diffusion::HorizontalDiffusion{NF} - implicit::ImplicitPrimitiveEq{NF} - device_setup::DeviceSetup{D} +$(TYPEDSIGNATURES) +Calls all `initialize!` functions for components of `model`, +except for `model.output` and `model.feedback` which are always called +at in `time_stepping!` and `model.implicit` which is done in `first_timesteps!`.""" +function initialize!(model::PrimitiveDry) + (;spectral_grid,horizontal_diffusion, + orography,planet,spectral_transform,geometry) = model + + initialize!(horizontal_diffusion,model) + initialize!(orography,planet,spectral_transform,geometry) + + # parameterizations + initialize!(model.boundary_layer_drag,model) + initialize!(model.temperature_relaxation,model) + initialize!(model.static_energy_diffusion,model) + # initialize!(model.vertical_diffusion,model) + + prognostic_variables = initial_conditions(model) + diagnostic_variables = DiagnosticVariables(spectral_grid) + return Simulation(prognostic_variables,diagnostic_variables,model) end -has(::Type{<:PrimitiveDryCore}, var_name::Symbol) = var_name in (:vor, :div, :temp, :pres) -has(::Type{<:PrimitiveWetCore}, var_name::Symbol) = var_name in (:vor, :div, :temp, :humid, :pres) -default_concrete_model(::Type{PrimitiveEquation}) = PrimitiveDryCoreModel -default_concrete_model(Model::Type{<:ModelSetup}) = Model +""" +$(SIGNATURES) +The PrimitiveDryModel struct holds all other structs that contain precalculated constants, +whether scalars or arrays that do not change throughout model integration. +$(TYPEDFIELDS)""" +@with_kw struct PrimitiveWetModel{NF<:AbstractFloat, D<:AbstractDevice} <: PrimitiveWet + "dictates resolution for many other components" + spectral_grid::SpectralGrid = SpectralGrid() + # DYNAMICS + "contains physical and orbital characteristics" + planet::AbstractPlanet = Earth() + atmosphere::AbstractAtmosphere = EarthAtmosphere() + initial_conditions::InitialConditions = ZonalWind() + orography::AbstractOrography{NF} = EarthOrography(spectral_grid) + + # PHYSICS/PARAMETERIZATIONS + physics::Bool = true + thermodynamics::Thermodynamics{NF} = Thermodynamics(spectral_grid,atmosphere) + boundary_layer_drag::BoundaryLayerDrag{NF} = LinearDrag(spectral_grid) + temperature_relaxation::TemperatureRelaxation{NF} = HeldSuarez(spectral_grid) + static_energy_diffusion::VerticalDiffusion{NF} = StaticEnergyDiffusion(spectral_grid) + # vertical_diffusion::VerticalDiffusion{NF} = VerticalLaplacian(spectral_grid) + + # NUMERICS + time_stepping::TimeStepper{NF} = Leapfrog(spectral_grid) + spectral_transform::SpectralTransform{NF} = SpectralTransform(spectral_grid) + horizontal_diffusion::HorizontalDiffusion{NF} = HyperDiffusion(spectral_grid) + implicit::AbstractImplicit{NF} = ImplicitPrimitiveEq(spectral_grid) + + # INTERNALS + clock::Clock = Clock() + geometry::Geometry{NF} = Geometry(spectral_grid) + constants::DynamicsConstants{NF} = DynamicsConstants(spectral_grid,planet,atmosphere,geometry) + device_setup::DeviceSetup{D} = DeviceSetup(CPUDevice()) + + # OUTPUT + output::AbstractOutputWriter = OutputWriter(spectral_grid) + feedback::AbstractFeedback = Feedback() +end + +has(::Type{<:PrimitiveWet}, var_name::Symbol) = var_name in (:vor, :div, :temp, :pres, :humid) +default_concrete_model(::Type{PrimitiveWet}) = PrimitiveWetModel +default_concrete_model(::Type{PrimitiveEquation}) = PrimitiveDryModel + """ - has(M::ModelSetup, var_name::Symbol) +$(TYPEDSIGNATURES) +Calls all `initialize!` functions for components of `model`, +except for `model.output` and `model.feedback` which are always called +at in `time_stepping!` and `model.implicit` which is done in `first_timesteps!`.""" +function initialize!(model::PrimitiveWet) + (;spectral_grid,horizontal_diffusion, + orography,planet,spectral_transform,geometry) = model + initialize!(horizontal_diffusion,model) + initialize!(orography,planet,spectral_transform,geometry) + + # parameterizations + initialize!(model.boundary_layer_drag,model) + initialize!(model.temperature_relaxation,model) + initialize!(model.static_energy_diffusion,model) + # initialize!(model.vertical_diffusion,model) + + prognostic_variables = initial_conditions(model) + diagnostic_variables = DiagnosticVariables(spectral_grid) + return Simulation(prognostic_variables,diagnostic_variables,model) +end + +"""$(TYPEDSIGNATURES) 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}, var_name::Symbol) = var_name in (:vor, :div, :temp, :humid, :pres) -has(M::ModelSetup, var_name) = has(typeof(M), var_name) \ No newline at end of file +The default fallback is that all variables are included. """ +has(M::Type{<:ModelSetup}, var_name::Symbol) = var_name in (:vor, :div, :temp, :humid, :pres) +has(M::ModelSetup, var_name) = has(typeof(M), var_name) + +"""$(TYPEDSIGNATURES) +Creates a concrete model depending on the abstract model type of `spectral_grid`.""" +function Model(;spectral_grid::SpectralGrid{WhichModel},kwargs...) where WhichModel + return default_concrete_model(WhichModel)(;spectral_grid,kwargs...) +end \ No newline at end of file diff --git a/src/dynamics/orography.jl b/src/dynamics/orography.jl new file mode 100644 index 000000000..93e9ae9ce --- /dev/null +++ b/src/dynamics/orography.jl @@ -0,0 +1,197 @@ +"""Orography with zero height in `orography` and zero surface geopotential `geopot_surf`. +$(TYPEDFIELDS)""" +struct NoOrography{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} <: AbstractOrography{NF,Grid} + "height [m] on grid-point space." + orography::Grid + + "surface geopotential, height*gravity [m²/s²]" + geopot_surf::LowerTriangularMatrix{Complex{NF}} +end + +""" +$(TYPEDSIGNATURES) +Generator function pulling the resolution information from `spectral_grid`.""" +function NoOrography(spectral_grid::SpectralGrid) + (;NF, Grid, nlat_half, trunc) = spectral_grid + orography = zeros(Grid{NF},nlat_half) + geopot_surf = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + return NoOrography{NF,Grid{NF}}(orography,geopot_surf) +end + +function Base.show(io::IO,orog::NoOrography) + print(io,"$(typeof(orog))") +end + +# no further initialization needed +initialize!(::NoOrography,::AbstractPlanet,::SpectralTransform,::Geometry) = nothing + + +"""Zonal ridge orography after Jablonowski and Williamson, 2006. +$(TYPEDFIELDS)""" +@with_kw struct ZonalRidge{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} <: AbstractOrography{NF,Grid} + + "conversion from σ to Jablonowski's ηᵥ-coordinates" + η₀::Float64 = 0.252 + + "max amplitude of zonal wind [m/s] that scales orography height" + u₀::Float64 = 35 + + # FIELDS (to be initialized in initialize!) + "height [m] on grid-point space." + orography::Grid + + "surface geopotential, height*gravity [m²/s²]" + geopot_surf::LowerTriangularMatrix{Complex{NF}} +end + +function Base.show(io::IO,orog::ZonalRidge) + print(io,"$(typeof(orog)):") + keys = (:η₀,:u₀) + for key in keys + val = getfield(orog,key) + print(io,"\n $key::$(typeof(val)) = $val") + end +end + +""" +$(TYPEDSIGNATURES) +Generator function pulling the resolution information from `spectral_grid`.""" +function ZonalRidge(spectral_grid::SpectralGrid;kwargs...) + (;NF, Grid, nlat_half, trunc) = spectral_grid + orography = zeros(Grid{NF},nlat_half) + geopot_surf = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + return ZonalRidge{NF,Grid{NF}}(;orography,geopot_surf,kwargs...) +end + +""" +$(TYPEDSIGNATURES) +Initialize the arrays `orography`,`geopot_surf` in `orog` following +Jablonowski and Williamson, 2006. +""" +function initialize!( orog::ZonalRidge, + P::AbstractPlanet, + S::SpectralTransform, + G::Geometry) + + (;gravity, rotation) = P + (;radius) = G + φ = G.latds # latitude for each grid point [˚N] + + (;orography, geopot_surf, η₀, u₀) = orog + + ηᵥ = (1-η₀)*π/2 # ηᵥ-coordinate of the surface [1] + A = u₀*cos(ηᵥ)^(3/2) # amplitude [m/s] + RΩ = radius*rotation # [m/s] + g⁻¹ = inv(gravity) # inverse gravity [s²/m] + + for ij in eachindex(φ,orography) + sinφ = sind(φ[ij]) + cosφ = cosd(φ[ij]) + + # Jablonowski & Williamson, 2006, eq. (7) + orography[ij] = g⁻¹*A*(A*(-2*sinφ^6*(cosφ^2 + 1/3) + 10/63) + (8/5*cosφ^3*(sinφ^2 + 2/3) - π/4)*RΩ) + end + + spectral!(geopot_surf,orography,S) # to grid-point space + geopot_surf .*= gravity # turn orography into surface geopotential + spectral_truncation!(geopot_surf) # set the lmax+1 harmonics to zero +end + + +"""Earth's orography read from file, with smoothing. +$(TYPEDFIELDS)""" +@with_kw struct EarthOrography{NF<:AbstractFloat,Grid<:AbstractGrid{NF}} <: AbstractOrography{NF,Grid} + + # OPTIONS + "path to the folder containing the orography file, pkg path default" + path::String = "SpeedyWeather.jl/input_data" + + "filename of orography" + file::String = "orography_F512.nc" + + "Grid the orography file comes on" + file_Grid::Type{<:AbstractGrid} = FullGaussianGrid + + "scale orography by a factor" + scale::Float64 = 1 + + "smooth the orography field?" + smoothing::Bool = true + + "power of Laplacian for smoothing" + smoothing_power::Float64 = 1.0 + + "highest degree l is multiplied by" + smoothing_strength::Float64 = 0.1 + + "resolution of orography in spectral trunc" + smoothing_truncation::Int = 85 + + + # FIELDS (to be initialized in initialize!) + "height [m] on grid-point space." + orography::Grid + + "surface geopotential, height*gravity [m²/s²]" + geopot_surf::LowerTriangularMatrix{Complex{NF}} +end + +""" +$(TYPEDSIGNATURES) +Generator function pulling the resolution information from `spectral_grid`.""" +function EarthOrography(spectral_grid::SpectralGrid;kwargs...) + (;NF, Grid, nlat_half, trunc) = spectral_grid + orography = zeros(Grid{NF},nlat_half) + geopot_surf = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + return EarthOrography{NF,Grid{NF}}(;orography,geopot_surf,kwargs...) +end + +function Base.show(io::IO,orog::EarthOrography) + print(io,"$(typeof(orog)):") + keys = (:path,:file,:scale,:smoothing,:smoothing_power, + :smoothing_strength,:smoothing_truncation) + for key in keys + val = getfield(orog,key) + print(io,"\n $key::$(typeof(val)) = $val") + end +end + +""" +$(TYPEDSIGNATURES) +Initialize the arrays `orography`,`geopot_surf` in `orog` by reading the +orography field from file. +""" +function initialize!( orog::EarthOrography, + P::AbstractPlanet, + S::SpectralTransform, + G::Geometry) + + (;orography, geopot_surf) = orog + (;gravity) = P + + # LOAD NETCDF FILE + if orog.path == "SpeedyWeather.jl/input_data" + path = joinpath(@__DIR__,"../../input_data",orog.file) + else + path = joinpath(orog.path,orog.file) + end + ncfile = NetCDF.open(path) + + orography_highres = ncfile.vars["orog"][:,:] # height [m] + + # Interpolate/coarsen to desired resolution + # TODO also read lat,lon from file and flip array in case it's not as expected + recompute_legendre = true # don't allocate large arrays as spectral transform is not reused + orography_spec = spectral(orography_highres;Grid=orog.file_Grid,recompute_legendre) + + copyto!(geopot_surf,orography_spec) # truncates to the size of geopot_surf, no *gravity yet + if orog.smoothing # smooth orography in spectral space? + SpeedyTransforms.spectral_smoothing!(geopot_surf,orog.smoothing_strength, + power=orog.smoothing_power, + truncation=orog.smoothing_truncation) + end + + gridded!(orography,geopot_surf,S) # to grid-point space + geopot_surf .*= gravity # turn orography into surface geopotential + spectral_truncation!(geopot_surf) # set the lmax+1 harmonics to zero +end \ No newline at end of file diff --git a/src/dynamics/planets.jl b/src/dynamics/planets.jl index 270503150..124d99a3f 100644 --- a/src/dynamics/planets.jl +++ b/src/dynamics/planets.jl @@ -1,13 +1,41 @@ -Base.@kwdef struct Earth <: Planet - radius::Float64 = 6.371e6 # radius of Earth [m] - rotation::Float64 = 7.29e-5 # angular frequency of Earth's rotation [rad/s] - gravity::Float64 = 9.81 # gravitational acceleration [m/s^2] - - daily_cycle::Bool = true # daily cycle? - length_of_day::Float64 = 24 # time [hrs] of a day - - seasonal_cycle::Bool = true # Seasonal cycle? - length_of_year::Float64 = 365.25 # time [days] of a year - equinox::DateTime = DateTime(2000,3,20) # Spring equinox (year irrelevant) - axial_tilt::Float64 = 23.4 # angle [˚] rotation axis tilt wrt to orbit +""" +$(TYPEDSIGNATURES) +Create a struct `Earth<:AbstractPlanet`, with the following physical/orbital +characteristics. Note that `radius` is not part of it as this should be chosen +in `SpectralGrid`. Keyword arguments are +$(TYPEDFIELDS) +""" +@with_kw struct Earth <: AbstractPlanet + + "angular frequency of Earth's rotation [rad/s]" + rotation::Float64 = 7.29e-5 + + "gravitational acceleration [m/s^2]" + gravity::Float64 = 9.81 + + "switch on/off daily cycle" + daily_cycle::Bool = true + + "[hrs] in a day" + length_of_day::Float64 = 24 + + "switch on/off seasonal cycle" + seasonal_cycle::Bool = true + + "[days] in a year" + length_of_year::Float64 = 365.25 + + "time of spring equinox (year irrelevant)" + equinox::DateTime = DateTime(2000,3,20) + + "angle [˚] rotation axis tilt wrt to orbit" + axial_tilt::Float64 = 23.4 +end + +function Base.show(io::IO,planet::AbstractPlanet) + print(io,"$(typeof(planet)):") + for key in propertynames(planet) + val = getfield(planet,key) + print(io,"\n $key::$(typeof(val)) = $val") + end end \ No newline at end of file diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl index b9e3ef5fb..b23e3efc3 100644 --- a/src/dynamics/prognostic_variables.jl +++ b/src/dynamics/prognostic_variables.jl @@ -29,94 +29,96 @@ struct PrognosticVariables{NF<:AbstractFloat,M<:ModelSetup} surface::PrognosticSurfaceTimesteps{NF} # dimensions - lmax::Int # two spectral dimensions: max meridional wavenumber - mmax::Int # max zonal wavenumber + trunc::Int # max degree of spherical harmonics n_steps::Int # N_STEPS time steps that are stored nlev::Int # number of vertical levels + + # scaling + scale::Base.RefValue{NF} + + # # scaling TODO think about including more flexible scaling factors here + # scaling_vor::Base.RefValue{NF} + # scaling_div::Base.RefValue{NF} + # scaling_temp::Base.RefValue{NF} + # scaling_humid::Base.RefValue{NF} + # scaling_pres::Base.RefValue{NF} end # ZERO GENERATOR FUNCTIONS # general version -function Base.zeros(::Type{PrognosticVariablesLayer{NF}},lmax::Integer,mmax::Integer) where NF +function Base.zeros(::Type{PrognosticVariablesLayer{NF}},trunc::Integer) where NF # use one more l for size compatibility with vector quantities - vor = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) - div = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) - temp = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) - humid = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) + # size trunc+2 x trunc+1 corresponds to lmax+1 x mmax + vor = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + div = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + temp = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + humid = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) return PrognosticVariablesLayer(vor,div,temp,humid) end # reduce size of unneeded variables if ModelSetup is provided -function Base.zeros(::Type{PrognosticVariablesLayer{NF}},model::ModelSetup,lmax::Integer,mmax::Integer) where NF +function Base.zeros(::Type{PrognosticVariablesLayer{NF}},Model::Type{<:ModelSetup},trunc::Integer) where NF # use one more l for size compatibility with vector quantities LTM = LowerTriangularMatrix{Complex{NF}} - vor = has(model, :vor) ? zeros(LTM,lmax+2,mmax+1) : LTM(undef, 0, 0) - div = has(model, :div) ? zeros(LTM,lmax+2,mmax+1) : LTM(undef, 0, 0) - temp = has(model, :temp) ? zeros(LTM,lmax+2,mmax+1) : LTM(undef, 0, 0) - humid = has(model, :humid) ? zeros(LTM,lmax+2,mmax+1) : LTM(undef, 0, 0) + vor = has(Model, :vor) ? zeros(LTM,trunc+2,trunc+1) : LTM(undef, 0, 0) + div = has(Model, :div) ? zeros(LTM,trunc+2,trunc+1) : LTM(undef, 0, 0) + temp = has(Model, :temp) ? zeros(LTM,trunc+2,trunc+1) : LTM(undef, 0, 0) + humid = has(Model, :humid) ? zeros(LTM,trunc+2,trunc+1) : LTM(undef, 0, 0) return PrognosticVariablesLayer(vor,div,temp,humid) end -function Base.zeros(::Type{PrognosticVariablesSurface{NF}},lmax::Integer,mmax::Integer) where NF +function Base.zeros(::Type{PrognosticVariablesSurface{NF}},trunc::Integer) where NF # use one more l for size compatibility with vector quantities - pres = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) + pres = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) return PrognosticVariablesSurface(pres) end -function Base.zeros(::Type{PrognosticVariablesSurface{NF}},model::ModelSetup,lmax::Integer,mmax::Integer) where NF +function Base.zeros(::Type{PrognosticVariablesSurface{NF}},Model::Type{<:ModelSetup},trunc::Integer) where NF # use one more l for size compatibility with vector quantities LTM = LowerTriangularMatrix{Complex{NF}} - pres = has(model, :pres) ? zeros(LTM,lmax+2,mmax+1) : LTM(undef, 0, 0) + pres = has(Model, :pres) ? zeros(LTM,trunc+2,trunc+1) : LTM(undef, 0, 0) return PrognosticVariablesSurface(pres) end # create time steps as N_STEPS-element vector of PrognosticVariablesLayer -function Base.zeros(::Type{PrognosticLayerTimesteps{NF}},lmax::Integer,mmax::Integer) where NF - return PrognosticLayerTimesteps([zeros(PrognosticVariablesLayer{NF},lmax,mmax) for _ in 1:N_STEPS]) +function Base.zeros(::Type{PrognosticLayerTimesteps{NF}},trunc::Integer) where NF + return PrognosticLayerTimesteps([zeros(PrognosticVariablesLayer{NF},trunc) for _ in 1:N_STEPS]) end -function Base.zeros(::Type{PrognosticSurfaceTimesteps{NF}},lmax::Integer,mmax::Integer) where NF - return PrognosticSurfaceTimesteps([zeros(PrognosticVariablesSurface{NF},lmax,mmax) for _ in 1:N_STEPS]) +function Base.zeros(::Type{PrognosticSurfaceTimesteps{NF}},trunc::Integer) where NF + return PrognosticSurfaceTimesteps([zeros(PrognosticVariablesSurface{NF},trunc) for _ in 1:N_STEPS]) end # also pass on model if available -function Base.zeros(::Type{PrognosticLayerTimesteps{NF}}, - model::ModelSetup, - lmax::Integer, - mmax::Integer) where NF - return PrognosticLayerTimesteps([zeros(PrognosticVariablesLayer{NF},model,lmax,mmax) for _ in 1:N_STEPS]) +function Base.zeros(::Type{PrognosticLayerTimesteps{NF}},Model::Type{<:ModelSetup},trunc::Integer) where NF + return PrognosticLayerTimesteps([zeros(PrognosticVariablesLayer{NF},Model,trunc) for _ in 1:N_STEPS]) end -function Base.zeros(::Type{PrognosticSurfaceTimesteps{NF}}, - model::ModelSetup, - lmax::Integer, - mmax::Integer) where NF - return PrognosticSurfaceTimesteps([zeros(PrognosticVariablesSurface{NF},model,lmax,mmax) for _ in 1:N_STEPS]) +function Base.zeros(::Type{PrognosticSurfaceTimesteps{NF}},Model::Type{<:ModelSetup},trunc::Integer) where NF + return PrognosticSurfaceTimesteps([zeros(PrognosticVariablesSurface{NF},Model,trunc) for _ in 1:N_STEPS]) end # general function to initiate all prognostic variables with zeros -function Base.zeros(::Type{PrognosticVariables{NF}}, - lmax::Integer, - mmax::Integer, - nlev::Integer) where NF - - layers = [zeros(PrognosticLayerTimesteps{NF},lmax,mmax) for _ in 1:nlev] # vector of nlev layers - surface = zeros(PrognosticSurfaceTimesteps{NF},lmax,mmax) - return PrognosticVariables{NF,ModelSetup}(layers,surface,lmax,mmax,N_STEPS,nlev) +function Base.zeros(::Type{PrognosticVariables{NF}},trunc::Integer,nlev::Integer) where NF + layers = [zeros(PrognosticLayerTimesteps{NF},trunc) for _ in 1:nlev] # vector of nlev layers + surface = zeros(PrognosticSurfaceTimesteps{NF},trunc) + + # initialize with scale=1, wrapped in RefValue for mutability + scale = Ref(one(NF)) + return PrognosticVariables{NF,ModelSetup}(layers,surface,trunc,N_STEPS,nlev,scale) end # pass on model to reduce size -function Base.zeros(::Type{PrognosticVariables{NF}}, - model::ModelSetup, - lmax::Integer, - mmax::Integer, - nlev::Integer) where NF - - layers = [zeros(PrognosticLayerTimesteps{NF},lmax,mmax) for _ in 1:nlev] # vector of nlev layers +function Base.zeros(::Type{PrognosticVariables{NF}},Model::Type{<:ModelSetup},trunc::Integer,nlev::Integer) where NF + + layers = [zeros(PrognosticLayerTimesteps{NF},Model,trunc) for _ in 1:nlev] # vector of nlev layers PST = PrognosticSurfaceTimesteps{NF} - surface = has(model, :pres) ? zeros(PST,lmax,mmax) : zeros(PST,-2,-1) - return PrognosticVariables{NF,typeof(model)}(layers,surface,lmax,mmax,N_STEPS,nlev) + surface = has(Model, :pres) ? zeros(PST,trunc) : zeros(PST,-1) + + # initialize with scale=1, wrapped in RefValue for mutability + scale = Ref(one(NF)) + return PrognosticVariables{NF,Model}(layers,surface,trunc,N_STEPS,nlev,scale) end has(progn::PrognosticVariables{NF,M}, var_name::Symbol) where {NF,M} = has(M, var_name) @@ -350,3 +352,26 @@ get_divergence(progn::PrognosticVariables; kwargs...) = get_var(progn, :div; kwa 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.surface.timesteps[lf].pres + +function Base.show(io::IO, P::PrognosticVariables) + + ζ = P.layers[end].timesteps[1].vor # create a view on vorticity + ζ_grid = Matrix(gridded(ζ)) # to grid space + ζ_grid = ζ_grid[:,end:-1:1] # flip latitudes + + nlon,nlat = size(ζ_grid) + + plot_kwargs = pairs(( xlabel="˚E", + xfact=360/(nlon-1), + ylabel="˚N", + yfact=180/(nlat-1), + yoffset=-90, + title="Surface relative vorticity", + colormap=:viridis, + compact=true, + colorbar=true, + width=60, + height=30)) + + print(io,UnicodePlots.heatmap(ζ_grid';plot_kwargs...)) +end \ No newline at end of file diff --git a/src/dynamics/scaling.jl b/src/dynamics/scaling.jl index b800c1a11..0652c7383 100644 --- a/src/dynamics/scaling.jl +++ b/src/dynamics/scaling.jl @@ -11,8 +11,7 @@ scale_coslat⁻¹!(A::AbstractMatrix,G::Geometry) = A.*G.coslat⁻¹' scale_coslat⁻²!(A::AbstractMatrix,G::Geometry) = A.*G.coslat⁻²' """ - _scale_lat!(A::AbstractGrid,v::AbstractVector) - +$(TYPEDSIGNATURES) Generic latitude scaling applied to `A` in-place with latitude-like vector `v`.""" function _scale_lat!(A::AbstractGrid{NF},v::AbstractVector) where {NF<:AbstractFloat} @boundscheck get_nlat(A) == length(v) || throw(BoundsError) @@ -28,63 +27,51 @@ function _scale_lat!(A::AbstractGrid{NF},v::AbstractVector) where {NF<:AbstractF end """ - scale!( progn::PrognosticVariables{NF}, - var::Symbol, - s::Number) where NF - -Scale the variable `var` inside `progn` with scalar `s`. +$(TYPEDSIGNATURES) +Scale the variable `var` inside `progn` with scalar `scale`. """ function scale!(progn::PrognosticVariables{NF}, var::Symbol, - s::Number) where NF - - if var == :pres # surface pressure is not stored in layers + scale::Real) where NF + if var == :pres for pres in progn.pres.timesteps - pres .*= s # pres*s but in-place + pres .*= scale end else for layer in progn.layers for step in layer.timesteps variable = getfield(step,var) - variable .*= s # var*s but in-place + variable .*= scale end end end end """ - scale!( progn::PrognosticVariables, - model::ModelSetup) - +$(TYPEDSIGNATURES) Scales the prognostic variables vorticity and divergence with the Earth's radius which is used in the dynamical core.""" function scale!(progn::PrognosticVariables, - model::ModelSetup) - - (; radius ) = model.geometry - scale!(progn,:vor,radius) - scale!(progn,:div,radius) + scale::Real) + scale!(progn,:vor,scale) + scale!(progn,:div,scale) + progn.scale[] = scale # store scaling information end """ - unscale!( progn::PrognosticVariables, - model::ModelSetup) - -Undo the radius-scaling of vorticity and divergence from scale!(progn,model).""" -function unscale!( progn::PrognosticVariables, - model::ModelSetup) - - (; radius ) = model.geometry - scale!(progn,:vor,inv(radius)) - scale!(progn,:div,inv(radius)) +$(TYPEDSIGNATURES) +Undo the radius-scaling of vorticity and divergence from scale!(progn,scale::Real).""" +function unscale!(progn::PrognosticVariables) + scale = progn.scale[] + scale!(progn,:vor,inv(scale)) + scale!(progn,:div,inv(scale)) + progn.scale[] = 1 # set scale back to 1=unscaled end """ - unscale!( variable::AbstractArray, - model::ModelSetup) - +$(TYPEDSIGNATURES) Undo the radius-scaling for any variable. Method used for netcdf output.""" function unscale!( variable::AbstractArray, - model::ModelSetup) - variable ./= model.geometry.radius + scale::Real) + variable ./= scale end \ No newline at end of file diff --git a/src/dynamics/spectral_grid.jl b/src/dynamics/spectral_grid.jl new file mode 100644 index 000000000..eb48ee4c2 --- /dev/null +++ b/src/dynamics/spectral_grid.jl @@ -0,0 +1,179 @@ +const DEFAULT_NF = Float32 +const DEFAULT_MODEL = PrimitiveDry +const DEFAULT_GRID = OctahedralGaussianGrid + +""" +Defines the horizontal spectral resolution and corresponding grid and the +vertical coordinate for SpeedyWeather.jl. Options are +$(TYPEDFIELDS) + +`nlat_half` and `npoints` should not be chosen but are derived from `trunc`, +`Grid` and `dealiasing`.""" +@with_kw struct SpectralGrid{Model<:ModelSetup} + "number format used throughout the model" + NF::Type{<:AbstractFloat} = DEFAULT_NF + + # HORIZONTAL + "horizontal resolution as the maximum degree of spherical harmonics" + trunc::Int = 31 + + "horizontal grid used for calculations in grid-point space" + Grid::Type{<:AbstractGrid} = DEFAULT_GRID + + "how to match spectral with grid resolution: dealiasing factor, 1=linear, 2=quadratic, 3=cubic grid" + dealiasing::Float64 = 2 + + "radius of the sphere [m]" + radius::Float64 = 6.371e6 + + # SIZE OF GRID from trunc, Grid, dealiasing: + "number of latitude rings on one hemisphere (Equator incl)" + nlat_half::Int = SpeedyTransforms.get_nlat_half(trunc,dealiasing) + + "total number of grid points in the horizontal" + npoints::Int = RingGrids.get_npoints(Grid,nlat_half) + + # VERTICAL + "number of vertical levels" + nlev::Int = default_nlev(Model) + + "coordinates used to discretize the vertical" + vertical_coordinates::VerticalCoordinates = default_vertical_coordinates(Model)(;nlev) + + # make sure nlev and vertical_coordinates.nlev match + function SpectralGrid{Model}(NF,trunc,Grid,dealiasing,radius,nlat_half,npoints,nlev,vertical_coordinates) where Model + if nlev == vertical_coordinates.nlev + return new(NF,trunc,Grid,dealiasing,radius,nlat_half,npoints, + nlev,vertical_coordinates) + else # use nlev from vert_coords: + return new(NF,trunc,Grid,dealiasing,radius,nlat_half,npoints, + vertical_coordinates.nlev,vertical_coordinates) + end + end +end + +# generator functions +SpectralGrid(NF::Type{<:AbstractFloat};kwargs...) = SpectralGrid(;NF,kwargs...) +SpectralGrid(Grid::Type{<:AbstractGrid};kwargs...) = SpectralGrid(;Grid,kwargs...) +SpectralGrid(NF::Type{<:AbstractFloat},Grid::Type{<:AbstractGrid};kwargs...) = SpectralGrid(;NF,Grid,kwargs...) +SpectralGrid(Model::Type{<:ModelSetup};kwargs...) = SpectralGrid{Model}(;kwargs...) +SpectralGrid(;kwargs...) = SpectralGrid{DEFAULT_MODEL}(;kwargs...) + +function Base.show(io::IO,SG::SpectralGrid) + (;NF,trunc,Grid,dealiasing,radius,nlat_half,npoints,nlev,vertical_coordinates) = SG + truncation = if dealiasing < 2 "linear" elseif dealiasing < 3 "quadratic" else "cubic" end + res = sqrt(4π*radius^2/npoints)/1000 # in [km] + println(io,"$(typeof(SG)):") + println(io," Spectral: T$trunc LowerTriangularMatrix{Complex{$NF}}, radius = $radius m") + println(io," Grid: $npoints-element, $(get_nlat(Grid,nlat_half))-ring $Grid{$NF} ($truncation)") + println(io," Resolution: $(@sprintf("%.3g",res))km (average)") + print(io," Vertical: $nlev-level $(typeof(vertical_coordinates))") +end + +""" +$(TYPEDSIGNATURES) +Construct Geometry struct containing parameters and arrays describing an iso-latitude grid <:AbstractGrid +and the vertical levels. Pass on `SpectralGrid` to calculate the following fields +$(TYPEDFIELDS) +""" +@with_kw struct Geometry{NF<:AbstractFloat} <: AbstractGeometry{NF} # NF: Number format + + "SpectralGrid that defines spectral and grid resolution" + spectral_grid::SpectralGrid + + "grid of the dynamical core" + Grid::Type{<:AbstractGrid} = spectral_grid.Grid + + "resolution parameter nlat_half of Grid, # of latitudes on one hemisphere (incl Equator)" + nlat_half::Int = spectral_grid.nlat_half + + + # GRID-POINT SPACE + "maximum number of longitudes (at/around Equator)" + nlon_max::Int = get_nlon_max(Grid,nlat_half) + + "=nlon_max, same (used for compatibility), TODO: still needed?" + nlon::Int = nlon_max + + "number of latitude rings" + nlat::Int = get_nlat(Grid,nlat_half) + + "number of vertical levels" + nlev::Int = spectral_grid.nlev + + "total number of grid points" + npoints::Int = spectral_grid.npoints + + "Planet's radius [m]" + radius::NF = spectral_grid.radius + + + # ARRAYS OF LANGITUDES/LONGITUDES + "array of latitudes in degrees (90˚...-90˚)" + latd::Vector{Float64} = get_latd(Grid,nlat_half) + + "array of longitudes in degrees (0...360˚), empty for non-full grids" + lond::Vector{Float64} = get_lond(Grid,nlat_half) + + "longitude (-180˚...180˚) for each grid point in ring order" + londs::Vector{NF} = get_latdlonds(Grid,nlat_half)[2] + + "latitude (-90˚...˚90) for each grid point in ring order" + latds::Vector{NF} = get_latdlonds(Grid,nlat_half)[1] + + "sin of latitudes" + sinlat::Vector{NF} = sind.(latd) + + "cos of latitudes" + coslat::Vector{NF} = cosd.(latd) + + "= 1/cos(lat)" + coslat⁻¹::Vector{NF} = 1 ./ coslat + + "= cos²(lat)" + coslat²::Vector{NF} = coslat.^2 + + "# = 1/cos²(lat)" + coslat⁻²::Vector{NF} = 1 ./ coslat² + + + # VERTICAL SIGMA COORDINATE σ = p/p0 (fraction of surface pressure) + "σ at half levels, σ_k+1/2" + σ_levels_half::Vector{NF} = spectral_grid.vertical_coordinates.σ_half + + "σ at full levels, σₖ" + σ_levels_full::Vector{NF} = 0.5*(σ_levels_half[2:end] + σ_levels_half[1:end-1]) + + "σ level thicknesses, σₖ₊₁ - σₖ" + σ_levels_thick::Vector{NF} = σ_levels_half[2:end] - σ_levels_half[1:end-1] + + "log of σ at full levels, include surface (σ=1) as last element" + ln_σ_levels_full::Vector{NF} = log.(vcat(σ_levels_full,1)) +end + +""" +$(TYPEDSIGNATURES) +Generator function for `Geometry` struct based on `spectral_grid`.""" +function Geometry(spectral_grid::SpectralGrid) + return Geometry{spectral_grid.NF}(;spectral_grid) +end + +# for barotropic/shallowwater always set the σ level to be defined between 0 and 1 +function Geometry(spectral_grid::SpectralGrid{Model}) where {Model<:Union{Barotropic,ShallowWater}} + return Geometry{spectral_grid.NF}(;spectral_grid,σ_levels_half=[0,1]) +end + +function Base.show(io::IO,G::Geometry) + print(io,"$(typeof(G)) for $(G.spectral_grid)") +end + +""" +$(TYPEDSIGNATURES) +Generator function for a SpectralTransform struct pulling in parameters from a SpectralGrid struct.""" +function SpeedyTransforms.SpectralTransform(spectral_grid::SpectralGrid; + recompute_legendre::Bool = false, + kwargs...) + (;NF, Grid, trunc, dealiasing) = spectral_grid + return SpectralTransform(NF,Grid,trunc;recompute_legendre,dealiasing,kwargs...) +end + diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl index c12192bc4..fdcdfa535 100644 --- a/src/dynamics/tendencies.jl +++ b/src/dynamics/tendencies.jl @@ -1,54 +1,52 @@ """ - dynamics_tendencies!(diagn,model) - +$(TYPEDSIGNATURES) Calculate all tendencies for the BarotropicModel.""" function dynamics_tendencies!( diagn::DiagnosticVariablesLayer, - model::BarotropicModel) - - # only (absolute) vorticity advection for the barotropic model - vorticity_flux_divcurl!(diagn,model,curl=false) # = -∇⋅(u(ζ+f),v(ζ+f)) + time::DateTime, + model::Barotropic) + forcing!(diagn,model.forcing,time) # = (Fᵤ, Fᵥ) forcing for u,v + vorticity_flux!(diagn,model) # = ∇×(v(ζ+f) + Fᵤ,v(ζ+f) + Fᵥ) end """ - dynamics_tendencies!(diagn,surface,pres,time,model) - +$(TYPEDSIGNATURES) Calculate all tendencies for the ShallowWaterModel.""" function dynamics_tendencies!( diagn::DiagnosticVariablesLayer, surface::SurfaceVariables, pres::LowerTriangularMatrix, # spectral pressure/η for geopotential time::DateTime, # time to evaluate the tendencies at - model::ShallowWaterModel) # struct containing all constants + model::ShallowWater) # struct containing all constants - S,C = model.spectral_transform, model.constants + S,C,G,O,F = model.spectral_transform, model.constants, model.geometry, model.orography, model.forcing # for compatibility with other ModelSetups pressure pres = interface displacement η here - vorticity_flux_divcurl!(diagn,model,curl=true) # = -∇⋅(u(ζ+f),v(ζ+f)), tendency for vorticity - # and ∇×(u(ζ+f),v(ζ+f)), tendency for divergence + forcing!(diagn,F,time) # = (Fᵤ, Fᵥ, Fₙ) forcing for u,v,η + vorticity_flux!(diagn,model) # = ∇×(v(ζ+f) + Fᵤ,v(ζ+f) + Fᵥ), tendency for vorticity + # = ∇⋅(v(ζ+f) + Fᵤ,v(ζ+f) + Fᵥ), tendency for divergence + geopotential!(diagn,pres,C) # geopotential Φ = gη in the shallow water model bernoulli_potential!(diagn,S) # = -∇²(E+Φ), tendency for divergence - volume_flux_divergence!(diagn,surface,model) # = -∇⋅(uh,vh), tendency pressure - - # interface forcing - (; interface_relaxation ) = model.parameters - interface_relaxation && interface_relaxation!(pres,surface,time,model) + volume_flux_divergence!(diagn,surface,O,C,G,S) # = -∇⋅(uh,vh), tendency pressure end """ - dynamics_tendencies!(diagn,surface,pres,time,model) - -Calculate all tendencies for the primitive equation model (wet or dry).""" +$(TYPEDSIGNATURES) +Calculate all tendencies for the PrimitiveEquation model (wet or dry).""" function dynamics_tendencies!( diagn::DiagnosticVariables, progn::PrognosticVariables, model::PrimitiveEquation, lf::Int=2) # leapfrog index for tendencies - B, G, S = model.boundaries, model.geometry, model.spectral_transform + O = model.orography + G = model.geometry + S = model.spectral_transform + C = model.constants (; surface ) = diagn # for semi-implicit corrections (α >= 0.5) linear gravity-wave related tendencies are # evaluated at previous timestep i-1 (i.e. lf=1 leapfrog time step) # nonlinear terms and parameterizations are always evaluated at lf - lf_implicit = model.parameters.implicit_α == 0 ? lf : 1 + lf_implicit = model.implicit.α == 0 ? lf : 1 pressure_gradient!(diagn,progn,lf,S) # calculate ∇ln(pₛ) @@ -57,46 +55,32 @@ function dynamics_tendencies!( diagn::DiagnosticVariables, # calculate Tᵥ = T + Tₖμq in spectral as a approxmation to Tᵥ = T(1+μq) used for geopotential linear_virtual_temperature!(diagn_layer,progn_layer,model,lf_implicit) - temperature_anomaly!(diagn_layer,diagn) # temperature relative to profile + temperature_anomaly!(diagn_layer) # temperature relative to profile end - geopotential!(diagn,B,G) # from ∂Φ/∂ln(pₛ) = -RTᵥ, used in bernoulli_potential! + geopotential!(diagn,O,C) # from ∂Φ/∂ln(pₛ) = -RTᵥ, used in bernoulli_potential! vertical_integration!(diagn,progn,lf_implicit,G) # get ū,v̄,D̄ on grid; and and D̄ in spectral - surface_pressure_tendency!(surface,model) # ∂ln(pₛ)/∂t = -(ū,v̄)⋅∇ln(pₛ) - D̄ - - # SINGLE THREADED VERSION - # for layer in diagn.layers - # vertical_velocity!(layer,surface,model) # calculate σ̇ for the vertical mass flux M = pₛσ̇ - # # add the RTₖlnpₛ term to geopotential - # linear_pressure_gradient!(layer,progn,model,lf_implicit) - # end - - # vertical_advection!(diagn,model) # use σ̇ for the vertical advection of u,v,T,q + surface_pressure_tendency!(surface,S) # ∂ln(pₛ)/∂t = -(ū,v̄)⋅∇ln(pₛ) - D̄ - # for layer in diagn.layers - # vordiv_tendencies!(layer,surface,model) # vorticity advection, pressure gradient term - # temperature_tendency!(layer,surface,model) # hor. advection + adiabatic term - # humidity_tendency!(layer,model) # horizontal advection of humidity (nothing for wetcore) - # bernoulli_potential!(layer,S) # add -∇²(E+ϕ+RTₖlnpₛ) term to div tendency - # end - - # MULTI-THREADED VERSION @floop for layer in diagn.layers - vertical_velocity!(layer,surface,model) # calculate σ̇ for the vertical mass flux M = pₛσ̇ + vertical_velocity!(layer,surface,G) # calculate σ̇ for the vertical mass flux M = pₛσ̇ # add the RTₖlnpₛ term to geopotential - linear_pressure_gradient!(layer,diagn,progn,model,lf_implicit) + linear_pressure_gradient!(layer,progn.surface,lf_implicit,C) end # wait all because vertical_velocity! needs to # finish before vertical_advection! @floop for layer in diagn.layers vertical_advection!(layer,diagn,model) # use σ̇ for the vertical advection of u,v,T,q vordiv_tendencies!(layer,surface,model) # vorticity advection, pressure gradient term - temperature_tendency!(layer,surface,model) # hor. advection + adiabatic term + temperature_tendency!(layer,model) # hor. advection + adiabatic term humidity_tendency!(layer,model) # horizontal advection of humidity (nothing for wetcore) bernoulli_potential!(layer,S) # add -∇²(E+ϕ+RTₖlnpₛ) term to div tendency end end +""" +$(TYPEDSIGNATURES) +Set the tendencies in `diagn` to zero.""" function zero_tendencies!(diagn::DiagnosticVariables) for layer in diagn.layers fill!(layer.tendencies.u_tend_grid,0) @@ -105,4 +89,6 @@ function zero_tendencies!(diagn::DiagnosticVariables) fill!(layer.tendencies.humid_tend_grid,0) end fill!(diagn.surface.pres_tend_grid,0) + fill!(diagn.surface.pres_tend,0) + return nothing end diff --git a/src/dynamics/tendencies_dynamics.jl b/src/dynamics/tendencies_dynamics.jl index 877af5f60..e36430b6b 100644 --- a/src/dynamics/tendencies_dynamics.jl +++ b/src/dynamics/tendencies_dynamics.jl @@ -6,7 +6,7 @@ function pressure_gradient!(diagn::DiagnosticVariables, (;pres) = progn.surface.timesteps[lf] # log of surface pressure ∇lnp_x_spec = diagn.layers[1].dynamics_variables.a # reuse work arrays for gradients ∇lnp_y_spec = diagn.layers[1].dynamics_variables.b # in spectral space - (; ∇lnp_x, ∇lnp_y ) = diagn.surface # but store in grid space + (;∇lnp_x, ∇lnp_y) = diagn.surface # but store in grid space ∇!(∇lnp_x_spec,∇lnp_y_spec,pres,S) # CALCULATE ∇ln(pₛ) gridded!(∇lnp_x,∇lnp_x_spec,S,unscale_coslat=true) # transform to grid: zonal gradient @@ -24,10 +24,9 @@ function pressure_flux!(diagn::DiagnosticVariablesLayer, end """Convert absolute and virtual temperature to anomalies wrt to the reference profile""" -function temperature_anomaly!( diagn_layer::DiagnosticVariablesLayer, - diagn::DiagnosticVariables) +function temperature_anomaly!(diagn_layer::DiagnosticVariablesLayer) - Tₖ = diagn.temp_profile[diagn_layer.k] # mean temperature at this level k + Tₖ = diagn_layer.temp_average[] # mean temperature on this layer (; temp_grid, temp_virt_grid ) = diagn_layer.grid_variables @. temp_grid -= Tₖ # absolute temperature -> anomaly @@ -110,10 +109,10 @@ of the logarithm of surface pressure ln(p_s) and D̄ the vertically averaged div 2. Multiply ū,v̄ with ∇ln(p_s) in grid-point space, convert to spectral. 3. D̄ is subtracted in spectral space. 4. Set tendency of the l=m=0 mode to 0 for better mass conservation.""" -function surface_pressure_tendency!(surf::SurfaceVariables{NF}, - model::PrimitiveEquation - ) where {NF<:AbstractFloat} - +function surface_pressure_tendency!( + surf::SurfaceVariables, + S::SpectralTransform, +) (; pres_tend, pres_tend_grid, ∇lnp_x, ∇lnp_y ) = surf # vertical averages need to be computed first! @@ -123,35 +122,36 @@ function surface_pressure_tendency!(surf::SurfaceVariables{NF}, # in grid-point space the the (ū,v̄)⋅∇lnpₛ term (swap sign in spectral) @. pres_tend_grid = ū*∇lnp_x + v̄*∇lnp_y - spectral!(pres_tend,pres_tend_grid,model.spectral_transform) + spectral!(pres_tend,pres_tend_grid,S) # for semi-implicit D̄ is calc at time step i-1 in vertical_integration! @. pres_tend = -pres_tend - D̄ # the -D̄ term in spectral and swap sign - pres_tend[1] = zero(NF) # for mass conservation + pres_tend[1] = 0 # for mass conservation spectral_truncation!(pres_tend) # remove lmax+1 row, only vectors use it return nothing end -function vertical_velocity!(diagn::DiagnosticVariablesLayer, - surf::SurfaceVariables, - model::PrimitiveEquation) - - (; k ) = diagn # vertical level - Δσₖ = model.geometry.σ_levels_thick[k] # σ level thickness at k - σk_half = model.geometry.σ_levels_half[k+1] # σ at k+1/2 +function vertical_velocity!( + diagn::DiagnosticVariablesLayer, + surf::SurfaceVariables, + G::Geometry, +) + (;k) = diagn # vertical level + Δσₖ = G.σ_levels_thick[k] # σ level thickness at k + σk_half = G.σ_levels_half[k+1] # σ at k+1/2 σ̇ = diagn.dynamics_variables.σ_tend # vertical mass flux M = pₛσ̇ at k+1/2 # sum of Δσ-weighted div, uv∇lnp from 1:k-1 - (; div_sum_above, uv∇lnp, uv∇lnp_sum_above ) = diagn.dynamics_variables - (; div_grid ) = diagn.grid_variables + (;div_sum_above, uv∇lnp, uv∇lnp_sum_above) = diagn.dynamics_variables + (;div_grid) = diagn.grid_variables ūv̄∇lnp = surf.pres_tend_grid # calc'd in surface_pressure_tendency! (excl -D̄) D̄ = surf.div_mean_grid # vertical avrgd div to be added to ūv̄∇lnp # mass flux σ̇ is zero at k=1/2 (not explicitly stored) and k=nlev+1/2 (stored in layer k) # set to zero for bottom layer then, and exit immediately - k == model.geometry.nlev && (fill!(σ̇,0); return nothing) + k == G.nlev && (fill!(σ̇,0); return nothing) # Hoskins and Simmons, 1975 just before eq. (6) σ̇ .= σk_half*(D̄ .+ ūv̄∇lnp) .- @@ -159,76 +159,13 @@ function vertical_velocity!(diagn::DiagnosticVariablesLayer, (uv∇lnp_sum_above .+ Δσₖ*uv∇lnp) # and level k σₖ-weighted uv∇lnp here end -# MULTI LAYER VERSION -function vertical_advection!( diagn::DiagnosticVariables, - model::PrimitiveEquation) - - wet_core = model isa PrimitiveWetCore - (; σ_levels_thick, nlev ) = model.geometry - @boundscheck nlev == diagn.nlev || throw(BoundsError) - - # set the k=1 level to zero in the beginning - u_tend_top = diagn.layers[1].tendencies.u_tend_grid - v_tend_top = diagn.layers[1].tendencies.v_tend_grid - temp_tend_top = diagn.layers[1].tendencies.temp_tend_grid - humid_tend_top = diagn.layers[1].tendencies.humid_tend_grid - fill!(u_tend_top,0) - fill!(v_tend_top,0) - fill!(temp_tend_top,0) - wet_core && fill!(humid_tend_top,0) - - # ALL LAYERS (but use indexing tricks to avoid out of bounds access for top/bottom) - @inbounds for k in 1:nlev - # for k==1 "above" term is 0, for k==nlev "below" term is zero - # avoid out-of-bounds indexing with k_above, k_below as follows - k_below = min(k+1,nlev) # just saturate, because M_nlev+1/2 = 0 (which zeros that term) - - # mass fluxes, M_1/2 = M_nlev+1/2 = 0, but k=1/2 isn't explicitly stored - σ_tend = diagn.layers[k].dynamics_variables.σ_tend - - # layer thickness Δσ on level k - Δσₖ = σ_levels_thick[k] - - u_tend_k = diagn.layers[k].tendencies.u_tend_grid - u_tend_below = diagn.layers[k_below].tendencies.u_tend_grid - u = diagn.layers[k].grid_variables.u_grid - u_below = diagn.layers[k_below].grid_variables.u_grid - - _vertical_advection!(u_tend_below,u_tend_k,σ_tend,u_below,u,Δσₖ) - - v_tend_k = diagn.layers[k].tendencies.v_tend_grid - v_tend_below = diagn.layers[k_below].tendencies.v_tend_grid - v = diagn.layers[k].grid_variables.v_grid - v_below = diagn.layers[k_below].grid_variables.v_grid - - _vertical_advection!(v_tend_below,v_tend_k,σ_tend,v_below,v,Δσₖ) - - T_tend_k = diagn.layers[k].tendencies.temp_tend_grid - T_tend_below = diagn.layers[k_below].tendencies.temp_tend_grid - T = diagn.layers[k].grid_variables.temp_grid - T_below = diagn.layers[k_below].grid_variables.temp_grid - - _vertical_advection!(T_tend_below,T_tend_k,σ_tend,T_below,T,Δσₖ) - - if wet_core - q_tend_k = diagn.layers[k].tendencies.humid_tend_grid - q_tend_below = diagn.layers[k_below].tendencies.humid_tend_grid - q = diagn.layers[k].grid_variables.humid_grid - q_below = diagn.layers[k_below].grid_variables.humid_grid - - _vertical_advection!(q_tend_below,q_tend_k,σ_tend,q_below,q,Δσₖ) - end - end -end - -# SINGLE LAYER VERSION function vertical_advection!( layer::DiagnosticVariablesLayer, diagn::DiagnosticVariables, model::PrimitiveEquation) (; k ) = layer # which layer are we on? - wet_core = model isa PrimitiveWetCore + wet_core = model isa PrimitiveWet (; σ_levels_thick, nlev ) = model.geometry # for k==1 "above" term is 0, for k==nlev "below" term is zero @@ -274,19 +211,6 @@ function vertical_advection!( layer::DiagnosticVariablesLayer, end end -# SINGLE THREADED VERSION uses the layer below to store intermediate result -function _vertical_advection!( ξ_tend_below::Grid, # tendency of quantity ξ at k+1 - ξ_tend_k::Grid, # tendency of quantity ξ at k - σ_tend::Grid, # vertical velocity at k+1/2 - ξ_below::Grid, # quantity ξ at k+1 - ξ::Grid, # quantity ξ at k - Δσₖ::NF # layer thickness on σ levels - ) where {NF<:AbstractFloat,Grid<:AbstractGrid{NF}} - Δσₖ2⁻¹ = -1/2Δσₖ # precompute - @. ξ_tend_below = σ_tend * (ξ_below - ξ) # store without Δσ-scaling in layer below - @. ξ_tend_k = Δσₖ2⁻¹ * (ξ_tend_k + ξ_tend_below) # combine with layer above and scale -end - # MULTI THREADED VERSION only writes into layer k function _vertical_advection!( ξ_tend::Grid, # tendency of quantity ξ at k σ_tend_above::Grid, # vertical velocity at k-1/2 @@ -312,12 +236,44 @@ function _vertical_advection!( ξ_tend::Grid, # tendency of quantity # @. ξ_tend += Δσₖ2⁻¹ * (σ_tend_below*(ξ_below - ξ) + σ_tend_above*(ξ - ξ_above)) end -function vordiv_tendencies!(diagn::DiagnosticVariablesLayer, - surf::SurfaceVariables, - model::PrimitiveEquation) - - (;f_coriolis, coslat⁻¹) = model.geometry - (;R_dry) = model.constants +""" +$(TYPEDSIGNATURES) +Function barrier to unpack `model`.""" +function vordiv_tendencies!( + diagn::DiagnosticVariablesLayer, + surf::SurfaceVariables, + model::PrimitiveEquation, +) + vordiv_tendencies!(diagn,surf,model.constants,model.geometry,model.spectral_transform) +end + +"""$(TYPEDSIGNATURES) +Tendencies for vorticity and divergence. Excluding Bernoulli potential with geopotential +and linear pressure gradient inside the Laplace operator, which are added later in +spectral space. + + u_tend += v*(f+ζ) - RTᵥ'*∇lnp_x + v_tend += -u*(f+ζ) - RTᵥ'*∇lnp_y + +`+=` because the tendencies already contain the parameterizations and vertical advection. +`f` is coriolis, `ζ` relative vorticity, `R` the gas constant `Tᵥ'` the virtual temperature +anomaly, `∇lnp` the gradient of surface pressure and `_x` and `_y` its zonal/meridional +components. The tendencies are then curled/dived to get the tendencies for vorticity/divergence in +spectral space + + ∂ζ/∂t = ∇×(u_tend,v_tend) + ∂D/∂t = ∇⋅(u_tend,v_tend) + ... + +`+ ...` because there's more terms added later for divergence.""" +function vordiv_tendencies!( + diagn::DiagnosticVariablesLayer, + surf::SurfaceVariables, + C::DynamicsConstants, + G::Geometry, + S::SpectralTransform, +) + (;R_dry, f_coriolis) = C + (;coslat⁻¹) = G (;u_tend_grid, v_tend_grid) = diagn.tendencies # already contains vertical advection u = diagn.grid_variables.u_grid # velocity @@ -345,7 +301,6 @@ function vordiv_tendencies!(diagn::DiagnosticVariablesLayer, (; vor_tend, div_tend ) = diagn.tendencies u_tend = diagn.dynamics_variables.a v_tend = diagn.dynamics_variables.b - S = model.spectral_transform spectral!(u_tend,u_tend_grid,S) spectral!(v_tend,v_tend_grid,S) @@ -356,29 +311,45 @@ function vordiv_tendencies!(diagn::DiagnosticVariablesLayer, # only vectors make use of the lmax+1 row, set to zero for scalars spectral_truncation!(vor_tend) spectral_truncation!(div_tend) + return nothing end """ -Compute the temperature tendency +$(TYPEDSIGNATURES) +Function barrier to unpack `model`.""" +function temperature_tendency!( + diagn::DiagnosticVariablesLayer, + model::PrimitiveEquation, +) + temperature_tendency!(diagn,model.constants,model.geometry,model.spectral_transform) +end + """ -function temperature_tendency!( diagn::DiagnosticVariablesLayer, - surf::SurfaceVariables, - model::PrimitiveEquation) +$(TYPEDSIGNATURES) +Compute the temperature tendency - (; temp_tend, temp_tend_grid ) = diagn.tendencies - (; div_grid, temp_grid ) = diagn.grid_variables - (; uv∇lnp, uv∇lnp_sum_above, div_sum_above ) = diagn.dynamics_variables - (; κ ) = model.constants # thermodynamic kappa - (; k ) = diagn # model level - (; nlev ) = model.geometry + ∂T/∂t += -∇⋅((u,v)*T') + T'D + κTᵥ*Dlnp/Dt + +`+=` because the tendencies already contain parameterizations and vertical advection. +`T'` is the anomaly with respect to the reference/average temperature. Tᵥ is the virtual +temperature used in the adiabatic term κTᵥ*Dlnp/Dt.""" +function temperature_tendency!( + diagn::DiagnosticVariablesLayer, + C::DynamicsConstants, + G::Geometry, + S::SpectralTransform, +) + (;temp_tend, temp_tend_grid) = diagn.tendencies + (;div_grid, temp_grid) = diagn.grid_variables + (;uv∇lnp, uv∇lnp_sum_above, div_sum_above) = diagn.dynamics_variables + (;κ) = C # thermodynamic kappa Tᵥ = diagn.grid_variables.temp_virt_grid # anomaly wrt to Tₖ - Tₖs = model.geometry.temp_ref_profile # reference temperatures - Tₖ = Tₖs[k] # reference temperature at k + Tₖ = diagn.temp_average[] # average temperatures # coefficients from Simmons and Burridge 1981 - σ_lnp_A = model.geometry.σ_lnp_A[k] # eq. 3.12, -1/Δσₖ*ln(σ_k+1/2/σ_k-1/2) - σ_lnp_B = model.geometry.σ_lnp_B[k] # eq. 3.12 -αₖ + σ_lnp_A = C.σ_lnp_A[diagn.k] # eq. 3.12, -1/Δσₖ*ln(σ_k+1/2/σ_k-1/2) + σ_lnp_B = C.σ_lnp_B[diagn.k] # eq. 3.12 -αₖ # semi-implicit: terms here are explicit+implicit evaluated at time step i # implicit_correction! then calculated the implicit terms from Vi-1 minus Vi @@ -391,63 +362,46 @@ function temperature_tendency!( diagn::DiagnosticVariablesLayer, σ_lnp_A * (div_sum_above+uv∇lnp_sum_above) + # eq. 3.12 1st term σ_lnp_B * (div_grid+uv∇lnp) + # eq. 3.12 2nd term uv∇lnp) # eq. 3.13 - - # SEMI-IMPLICIT ALTERNATIVE - # evaluate only the explicit terms at time step i and the implicit terms - # in implicit_correction! at i-1, however, this is more expensive then above - - # ūv̄∇lnp = surf.pres_tend_grid - # # for explicit vertical advection get reference temperature differences - # k_above = max(1,k-1) # layer index above - # k_below = min(k+1,nlev) # layer index below - # ΔT_above = Tₖ - Tₖs[k_above] # temperature difference to layer above - # ΔT_below = Tₖs[k_below] - Tₖ # and to layer below - - # # for explicit vertical advection terms - # σₖ = model.geometry.σ_levels_full[k] # should be Σ_r=1^k Δσᵣ for model top at >0hPa - # σₖ_above = model.geometry.σ_levels_full[k_above] - # Δσₖ2⁻¹ = -1/2model.geometry.σ_levels_thick[k] - - # # Hoskins and Simmons 1975, Appendix 1 but the adiabatic term therein as above - # @. temp_tend_grid += temp_grid*div_grid + - # Δσₖ2⁻¹*ΔT_below*( σₖ*ūv̄∇lnp - (uv∇lnp_sum_above + σₖ*uv∇lnp)) + - # Δσₖ2⁻¹*ΔT_above*(σₖ_above*ūv̄∇lnp - uv∇lnp_sum_above) + - # κ*Tₖ*(σ_lnp_A*uv∇lnp_sum_above + σ_lnp_B*uv∇lnp) + - # κ*Tᵥ*(σ_lnp_A * (div_sum_above+uv∇lnp_sum_above) + σ_lnp_B * (div_grid+uv∇lnp)) + - # κ*(Tᵥ+Tₖ)*uv∇lnp - - spectral!(temp_tend,temp_tend_grid,model.spectral_transform) + + spectral!(temp_tend,temp_tend_grid,S) # now add the -∇⋅((u,v)*T') term - flux_divergence!(temp_tend,temp_grid,diagn,model,add=true,flipsign=true) + flux_divergence!(temp_tend,temp_grid,diagn,G,S,add=true,flipsign=true) # only vectors make use of the lmax+1 row, set to zero for scalars - spectral_truncation!(temp_tend) + spectral_truncation!(temp_tend) + return nothing end function humidity_tendency!(diagn::DiagnosticVariablesLayer, - model::PrimitiveWetCore) + model::PrimitiveWet) + G = model.geometry + S = model.spectral_transform (; humid_tend, humid_tend_grid ) = diagn.tendencies (; humid_grid ) = diagn.grid_variables - horizontal_advection!(humid_tend,humid_tend_grid,humid_grid,diagn,model) - + # add horizontal advection to parameterization + vertical advection tendencies + horizontal_advection!(humid_tend,humid_tend_grid,humid_grid,diagn,G,S,add=true) + # only vectors make use of the lmax+1 row, set to zero for scalars spectral_truncation!(humid_tend) end # no humidity tendency for dry core -humidity_tendency!(::DiagnosticVariablesLayer,::PrimitiveDryCore) = nothing - -function horizontal_advection!( A_tend::LowerTriangularMatrix{Complex{NF}}, # Ouput: tendency to write into - A_tend_grid::AbstractGrid{NF}, # Input: tendency incl prev terms - A_grid::AbstractGrid{NF}, # Input: grid field to be advected - diagn::DiagnosticVariablesLayer{NF}, - model::ModelSetup; - add::Bool=true) where NF # add/overwrite A_tend_grid? - - (; div_grid ) = diagn.grid_variables +humidity_tendency!(::DiagnosticVariablesLayer,::PrimitiveDry) = nothing + +function horizontal_advection!( + A_tend::LowerTriangularMatrix{Complex{NF}}, # Ouput: tendency to write into + A_tend_grid::AbstractGrid{NF}, # Input: tendency incl prev terms + A_grid::AbstractGrid{NF}, # Input: grid field to be advected + diagn::DiagnosticVariablesLayer{NF}, + G::Geometry, + S::SpectralTransform; + add::Bool=true # add/overwrite A_tend_grid? +) where NF + + (;div_grid) = diagn.grid_variables @inline kernel(a,b,c) = add ? a+b*c : b*c @@ -457,22 +411,32 @@ function horizontal_advection!( A_tend::LowerTriangularMatrix{Complex{NF}}, # Ou A_tend_grid[ij] = kernel(A_tend_grid[ij],A_grid[ij],div_grid[ij]) end - spectral!(A_tend,A_tend_grid,model.spectral_transform) # for +A*div in spectral space + spectral!(A_tend,A_tend_grid,S) # for +A*div in spectral space # now add the -∇⋅((u,v)*A) term - flux_divergence!(A_tend,A_grid,diagn,model,add=true,flipsign=true) + flux_divergence!(A_tend,A_grid,diagn,G,S,add=true,flipsign=true) end -"""Computes -∇⋅((u,v)*A)""" +""" +$(TYPEDSIGNATURES) +Computes ∇⋅((u,v)*A) with the option to add/overwrite `A_tend` and to +`flip_sign` of the flux divergence by doing so. + +- `A_tend = ∇⋅((u,v)*A)` for `add=false`, `flip_sign=false` +- `A_tend = -∇⋅((u,v)*A)` for `add=false`, `flip_sign=true` +- `A_tend += ∇⋅((u,v)*A)` for `add=true`, `flip_sign=false` +- `A_tend -= ∇⋅((u,v)*A)` for `add=true`, `flip_sign=true` +""" function flux_divergence!( A_tend::LowerTriangularMatrix{Complex{NF}}, # Ouput: tendency to write into A_grid::AbstractGrid{NF}, # Input: grid field to be advected diagn::DiagnosticVariablesLayer{NF}, - model::ModelSetup; + G::Geometry{NF}, + S::SpectralTransform{NF}; add::Bool=true, # add result to A_tend or overwrite for false flipsign::Bool=true) where NF # compute -∇⋅((u,v)*A) (true) or ∇⋅((u,v)*A)? - (; u_grid, v_grid ) = diagn.grid_variables - (; coslat⁻¹ ) = model.geometry + (;u_grid, v_grid) = diagn.grid_variables + (;coslat⁻¹) = G # reuse general work arrays a,b,a_grid,b_grid uA = diagn.dynamics_variables.a # = u*A in spectral @@ -491,62 +455,115 @@ function flux_divergence!( A_tend::LowerTriangularMatrix{Complex{NF}}, # Ouput: end end - spectral!(uA,uA_grid,model.spectral_transform) - spectral!(vA,vA_grid,model.spectral_transform) + spectral!(uA,uA_grid,S) + spectral!(vA,vA_grid,S) - divergence!(A_tend,uA,vA,model.spectral_transform;add,flipsign) - return uA,vA # return for curl calculation (ShallowWater) + divergence!(A_tend,uA,vA,S;add,flipsign) + return nothing end """ -function vorticity_flux_divcurl!( diagn::DiagnosticVariablesLayer, - model::ModelSetup; - curl::Bool=true) # calculate curl of vor flux? - -1) Compute the vorticity advection as the (negative) divergence of the vorticity fluxes -∇⋅(uv*(ζ+f)). -First, compute the uv*(ζ+f), then transform to spectral space and take the divergence and flip the sign. -2) Compute the curl of the vorticity fluxes ∇×(uω,vω) and store in divergence tendency.""" -function vorticity_flux_divcurl!( diagn::DiagnosticVariablesLayer, - model::ModelSetup; - curl::Bool=true) # calculate curl of vor flux? - - G = model.geometry - S = model.spectral_transform +$(TYPEDSIGNATURES) +Compute the vorticity advection as the curl/div of the vorticity fluxes - (; u_grid, v_grid, vor_grid ) = diagn.grid_variables - (; vor_tend, div_tend ) = diagn.tendencies +`∂ζ/∂t = ∇×(u_tend,v_tend)` +`∂D/∂t = ∇⋅(u_tend,v_tend)` - # add the planetary vorticity f to relative vorticity ζ = absolute vorticity ω - absolute_vorticity!(vor_grid,G) +with - # now do -∇⋅(uω,vω) and store in vor_tend - uω,vω = flux_divergence!(vor_tend,vor_grid,diagn,model,add=false,flipsign=true) +`u_tend = Fᵤ + v*(ζ+f)` +`v_tend = Fᵥ - u*(ζ+f)` - # = ∇×(uω,vω) = ∇×(uv*(ζ+f)), write directly into tendency - # curl not needed for BarotropicModel - curl && curl!(div_tend,uω,vω,S,add=false,flipsign=false) -end +with `Fᵤ,Fᵥ` from `u_tend_grid`/`v_tend_grid` that are assumed to be alread +set in `forcing!`. Set `div=false` for the BarotropicModel which doesn't +require the divergence tendency.""" +function vorticity_flux_curldiv!( diagn::DiagnosticVariablesLayer, + C::DynamicsConstants, + G::Geometry, + S::SpectralTransform; + div::Bool=true) # also calculate div of vor flux? + + (;f_coriolis) = C + (;coslat⁻¹) = G -function absolute_vorticity!( vor::AbstractGrid, - G::Geometry) + (;u_tend_grid, v_tend_grid) = diagn.tendencies # already contains forcing + u = diagn.grid_variables.u_grid # velocity + v = diagn.grid_variables.v_grid # velocity + vor = diagn.grid_variables.vor_grid # relative vorticity - (; f_coriolis ) = G - @boundscheck length(f_coriolis) == get_nlat(vor) || throw(BoundsError) + # precompute ring indices and boundscheck + rings = eachring(u_tend_grid,v_tend_grid,u,v,vor) - rings = eachring(vor) @inbounds for (j,ring) in enumerate(rings) + coslat⁻¹j = coslat⁻¹[j] f = f_coriolis[j] for ij in ring - vor[ij] += f + ω = vor[ij] + f # absolute vorticity + u_tend_grid[ij] = (u_tend_grid[ij] + v[ij]*ω)*coslat⁻¹j + v_tend_grid[ij] = (v_tend_grid[ij] - u[ij]*ω)*coslat⁻¹j end end + + # divergence and curl of that u,v_tend vector for vor,div tendencies + (; vor_tend, div_tend ) = diagn.tendencies + u_tend = diagn.dynamics_variables.a + v_tend = diagn.dynamics_variables.b + + spectral!(u_tend,u_tend_grid,S) + spectral!(v_tend,v_tend_grid,S) + + curl!(vor_tend,u_tend,v_tend,S) # ∂ζ/∂t = ∇×(u_tend,v_tend) + div && divergence!(div_tend,u_tend,v_tend,S) # ∂D/∂t = ∇⋅(u_tend,v_tend) + + # only vectors make use of the lmax+1 row, set to zero for scalars + spectral_truncation!(vor_tend) + div && spectral_truncation!(div_tend) + return nothing end """ - bernoulli_potential!( diagn::DiagnosticVariablesLayer, - G::Geometry, - S::SpectralTransform) +$(TYPEDSIGNATURES) +Vorticity flux tendency in the shallow water equations + +`∂ζ/∂t = ∇×(u_tend,v_tend)` +`∂D/∂t = ∇⋅(u_tend,v_tend)` + +with + +`u_tend = Fᵤ + v*(ζ+f)` +`v_tend = Fᵥ - u*(ζ+f)` + +with Fᵤ,Fᵥ the forcing from `forcing!` already in `u_tend_grid`/`v_tend_grid` and +vorticity ζ, coriolis f.""" +function vorticity_flux!(diagn::DiagnosticVariablesLayer,model::ShallowWater) + C = model.constants + G = model.geometry + S = model.spectral_transform + vorticity_flux_curldiv!(diagn,C,G,S,div=true) +end +""" +$(TYPEDSIGNATURES) +Vorticity flux tendency in the barotropic vorticity equation + +`∂ζ/∂t = ∇×(u_tend,v_tend)` + +with + +`u_tend = Fᵤ + v*(ζ+f)` +`v_tend = Fᵥ - u*(ζ+f)` + +with Fᵤ,Fᵥ the forcing from `forcing!` already in `u_tend_grid`/`v_tend_grid` and +vorticity ζ, coriolis f.""" +function vorticity_flux!(diagn::DiagnosticVariablesLayer,model::Barotropic) + C = model.constants + G = model.geometry + S = model.spectral_transform + vorticity_flux_curldiv!(diagn,C,G,S,div=false) +end + +""" +$(TYPEDSIGNATURES) Computes the Laplace operator ∇² of the Bernoulli potential `B` in spectral space. 1. computes the kinetic energy KE = ½(u²+v²) on the grid 2. transforms KE to spectral space @@ -572,16 +589,26 @@ function bernoulli_potential!( diagn::DiagnosticVariablesLayer{NF}, ∇²!(div_tend,bernoulli,S,add=true,flipsign=true) # add -∇²(½(u² + v²) + ϕ) end -function linear_pressure_gradient!( diagn_layer::DiagnosticVariablesLayer, - diagn::DiagnosticVariables, - progn::PrognosticVariables, - model::PrimitiveEquation, - lf::Int) # leapfrog index to evaluate tendencies on - - (; R_dry ) = model.constants - Tₖ = diagn.temp_profile[diagn_layer.k] # mean temperature at layer k - (;pres) = progn.surface.timesteps[lf] - (; geopot ) = diagn_layer.dynamics_variables +""" +$(TYPEDSIGNATURES) +Add the linear contribution of the pressure gradient to the geopotential. +The pressure gradient in the divergence equation takes the form + +`-∇⋅(Rd*Tᵥ*∇lnpₛ) = -∇⋅(Rd*Tᵥ'*∇lnpₛ) - ∇²(Rd*Tₖ*lnpₛ)` + +So that the second term inside the Laplace operator can be added to the geopotential. +Rd is the gas constant, Tᵥ the virtual temperature and Tᵥ' its anomaly wrt to the +average or reference temperature Tₖ, lnpₛ is the logarithm of surface pressure.""" +function linear_pressure_gradient!( + diagn::DiagnosticVariablesLayer, + surface::PrognosticSurfaceTimesteps, + lf::Int, # leapfrog index to evaluate tendencies on + C::DynamicsConstants, +) + (; R_dry ) = C + Tₖ = diagn.temp_average[] # mean temperature at layer k + (;pres) = surface.timesteps[lf] + (;geopot) = diagn.dynamics_variables # -R_dry*Tₖ*∇²lnpₛ, linear part of the ∇⋅RTᵥ∇lnpₛ pressure gradient term # Tₖ being the reference temperature profile, the anomaly term T' = Tᵥ - Tₖ is calculated @@ -591,54 +618,38 @@ function linear_pressure_gradient!( diagn_layer::DiagnosticVariablesLayer, end """ - volume_flux_divergence!(diagn::DiagnosticVariablesLayer, - surface::SurfaceVariables, - model::ShallowWater) - +$(TYPEDSIGNATURES) Computes the (negative) divergence of the volume fluxes `uh,vh` for the continuity equation, -∇⋅(uh,vh).""" function volume_flux_divergence!( diagn::DiagnosticVariablesLayer, surface::SurfaceVariables, - model::ShallowWater) + orog::AbstractOrography, + constants::DynamicsConstants, + G::Geometry, + S::SpectralTransform) (; pres_grid, pres_tend ) = surface - (; orography ) = model.boundaries.orography - H₀ = model.constants.layer_thickness + (; orography ) = orog + H₀ = constants.layer_thickness # compute dynamic layer thickness h on the grid # pres_grid is η, the interface displacement, update to # layer thickness h = η + H, H is the layer thickness at rest - # H = H₀ - orography, H₀ is the layer thickness without mountains + # H = H₀ - orography, H₀ is the layer thickness at rest without mountains pres_grid .+= H₀ .- orography # now do -∇⋅(uh,vh) and store in pres_tend - flux_divergence!(pres_tend,pres_grid,diagn,model,add=false,flipsign=true) + flux_divergence!(pres_tend,pres_grid,diagn,G,S,add=true,flipsign=true) end -function interface_relaxation!( η::LowerTriangularMatrix{Complex{NF}}, - surface::SurfaceVariables{NF}, - time::DateTime, # time of relaxation - M::ShallowWaterModel, # contains η⁰, which η is relaxed to - ) where NF - - (; pres_tend ) = surface - (; seasonal_cycle, equinox, axial_tilt ) = M.parameters.planet - A = M.parameters.interface_relax_amplitude - - s = 45/23.5 # heuristic conversion to Legendre polynomials - θ = seasonal_cycle ? s*axial_tilt*sin(Dates.days(time - equinox)/365.25*2π) : 0 - η2 = convert(NF,A*(2sind(θ))) # l=1,m=0 harmonic - η3 = convert(NF,A*(0.2-1.5cosd(θ))) # l=2,m=0 harmonic - - τ⁻¹ = inv(M.constants.interface_relax_time) - pres_tend[2] += τ⁻¹*(η2-η[2]) - pres_tend[3] += τ⁻¹*(η3-η[3]) -end - -function SpeedyTransforms.gridded!( diagn::DiagnosticVariables, # all diagnostic variables - progn::PrognosticVariables, # all prognostic variables - lf::Int, # leapfrog index - model::ModelSetup, - ) +""" +$(TYPEDSIGNATURES) +Propagate the spectral state of `progn` to `diagn` using time step/leapfrog index `lf`. +Function barrier that calls gridded! for the respective `model`.""" +function SpeedyTransforms.gridded!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Int, + model::ModelSetup) # all variables on layers for (progn_layer,diagn_layer) in zip(progn.layers,diagn.layers) @@ -655,12 +666,7 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariables, # all diagno end """ - gridded!( diagn::DiagnosticVariables{NF}, # all diagnostic variables - progn::PrognosticVariables{NF}, # all prognostic variables - M::BarotropicModel, # everything that's constant - lf::Int=1 # leapfrog index - ) where NF - +$(TYPEDSIGNATURES) Propagate the spectral state of the prognostic variables `progn` to the diagnostic variables in `diagn` for the barotropic vorticity model. Updates grid vorticity, spectral stream function and spectral and grid velocities u,v.""" @@ -690,16 +696,11 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, end """ - gridded!( diagn::DiagnosticVariables{NF}, # all diagnostic variables - progn::PrognosticVariables{NF}, # all prognostic variables - lf::Int=1 # leapfrog index - M::ShallowWaterModel, # everything that's constant - ) where NF - +$(TYPEDSIGNATURES) Propagate the spectral state of the prognostic variables `progn` to the diagnostic variables in `diagn` for the shallow water model. Updates grid vorticity, grid divergence, grid interface displacement (`pres_grid`) and the velocities -U,V (scaled by cos(lat)).""" +u,v.""" function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, progn::PrognosticLayerTimesteps, lf::Int, # leapfrog index @@ -729,11 +730,16 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, return nothing end +""" +$(TYPEDSIGNATURES) +Propagate the spectral state of the prognostic variables `progn` to the +diagnostic variables in `diagn` for primitive equation models. Updates grid vorticity, +grid divergence, grid temperature, pressure (`pres_grid`) and the velocities +u,v.""" function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, progn::PrognosticLayerTimesteps, lf::Int, # leapfrog index - model::PrimitiveEquation, # everything that's constant - ) + model::PrimitiveEquation) # everything that's constant (; vor_grid, div_grid, u_grid, v_grid ) = diagn.grid_variables (; temp_grid, humid_grid ) = diagn.grid_variables @@ -741,7 +747,7 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, V = diagn.dynamics_variables.b # U = u*coslat, V=v*coslat S = model.spectral_transform - wet_core = model isa PrimitiveWetCore + wet_core = model isa PrimitiveWet vor_lf = progn.timesteps[lf].vor # pick leapfrog index without memory allocation div_lf = progn.timesteps[lf].div @@ -759,6 +765,7 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, wet_core && gridded!(humid_grid,humid_lf,S) # specific humidity (wet core only) # include humidity effect into temp for everything stability-related + temperature_average!(diagn,temp_lf,S) # TODO: do at frequency of reinitialize implicit? virtual_temperature!(diagn,temp_lf,model) # temp = virt temp for dry core # transform from U,V in spectral to u,v on grid (U,V = u,v*coslat) @@ -766,4 +773,18 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, gridded!(v_grid,V,S,unscale_coslat=true) return nothing -end \ No newline at end of file +end + +""" +$(TYPEDSIGNATURES) +Calculates the average temperature of a layer from the l=m=0 harmonic +and stores the result in `diagn.temp_average`""" +function temperature_average!( + diagn::DiagnosticVariablesLayer, + temp::LowerTriangularMatrix, + S::SpectralTransform, +) + + # average from l=m=0 harmonic divided by norm of the sphere + diagn.temp_average[] = real(temp[1])/S.norm_sphere +end \ No newline at end of file diff --git a/src/dynamics/time_integration.jl b/src/dynamics/time_integration.jl index dcc94be0e..48a29bb7a 100644 --- a/src/dynamics/time_integration.jl +++ b/src/dynamics/time_integration.jl @@ -1,27 +1,97 @@ """ - leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}}, # prognostic variable at t - A_new::LowerTriangularMatrix{Complex{NF}}, # prognostic variable at t+dt - tendency::LowerTriangularMatrix{Complex{NF}}, # tendency (dynamics+physics) of A - dt::Real, # time step (=2Δt, but for init steps =Δt,Δt/2) - lf::Int=2, # leapfrog index to dis/enable William's filter - C::DynamicsConstants{NF}, # struct with constants used at runtime - ) where {NF<:AbstractFloat} # number format NF +Clock struct keeps track of the model time, how many days to integrate for +and how many time steps this takes +$(TYPEDFIELDS).""" +@with_kw mutable struct Clock + "current model time" + time::DateTime = DateTime(2000,1,1) + + "number of days to integrate for, set in run!(::Simulation)" + n_days::Float64 = 0 + + "number of time steps to integrate for, set initialize!(::Clock,::TimeStepper)" + n_timesteps::Int = 0 +end + +function initialize!(clock::Clock,time_stepping::TimeStepper) + clock.n_timesteps = ceil(Int,24*clock.n_days/time_stepping.Δt_hrs) + return clock +end + +function Clock(time_stepping::TimeStepper;kwargs...) + clock = Clock(;kwargs...) + initialize!(clock,time_stepping) +end + +""" +Leapfrog time stepping defined by the following fields +$(TYPEDFIELDS) +""" +@with_kw struct Leapfrog{NF} <: TimeStepper{NF} + + # DIMENSIONS + "spectral resolution (max degree of spherical harmonics)" + trunc::Int + + # OPTIONS + "time step in minutes for T31, scale linearly to `trunc`" + Δt_at_T31::Float64 = 30 + + "radius of sphere [m], used for scaling" + radius::NF = 6.371e6 + + # NUMERICS + "Robert (1966) time filter coefficeint to suppress comput. mode" + robert_filter::NF = 0.05 + + "William's time filter (Amezcua 2011) coefficient for 3rd order acc" + william_filter::NF = 0.53 + + + # DERIVED FROM OPTIONS + "time step Δt [s] at specified resolution" + Δt_sec::Int = round(Int,60*Δt_at_T31*(32/(trunc+1))) + + "time step Δt [s/m] at specified resolution, scaled by 1/radius" + Δt::NF = Δt_sec/radius + + "convert time step Δt from minutes to hours" + Δt_hrs::Float64 = Δt_sec/3600 +end -Performs one leapfrog time step with (`lf=2`) or without (`lf=1`) Robert+William's filter -(see William (2009), Montly Weather Review, Eq. 7-9). """ +$(TYPEDSIGNATURES) +Generator function for a Leapfrog struct using `spectral_grid` +for the resolution information.""" +function Leapfrog(spectral_grid::SpectralGrid;kwargs...) + (;NF,trunc,radius) = spectral_grid + return Leapfrog{NF}(;trunc,radius,kwargs...) +end + +function Base.show(io::IO,L::Leapfrog) + print(io,"$(typeof(L)):") + for key = propertynames(L) + val = getfield(L,key) + print(io,"\n $key::$(typeof(val)) = $val") + end +end + +""" +$(TYPEDSIGNATURES) +Performs one leapfrog time step with (`lf=2`) or without (`lf=1`) Robert+William's filter +(see William (2009), Montly Weather Review, Eq. 7-9).""" function leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}}, # prognostic variable at t A_new::LowerTriangularMatrix{Complex{NF}}, # prognostic variable at t+dt tendency::LowerTriangularMatrix{Complex{NF}}, # tendency (dynamics+physics) of A dt::Real, # time step (=2Δt, but for init steps =Δt,Δt/2) lf::Int, # leapfrog index to dis/enable William's filter - C::DynamicsConstants{NF}, # struct with constants used at runtime + L::Leapfrog{NF}, # struct with constants ) where {NF<:AbstractFloat} # number format NF @boundscheck lf == 1 || lf == 2 || throw(BoundsError()) # index lf picks leapfrog dim - A_lf = lf == 1 ? A_old : A_new # view on either t or t+dt to dis/enable William's filter - (; robert_filter, williams_filter ) = C # coefficients for the Robert and William's filter + A_lf = lf == 1 ? A_old : A_new # view on either t or t+dt to dis/enable William's filter + (;robert_filter, william_filter) = L # coefficients for the Robert and William's filter two = convert(NF,2) # 2 in number format NF dt_NF = convert(NF,dt) # time step dt in number format NF @@ -30,8 +100,8 @@ function leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}}, # prognostic # see William (2009), Eq. 7-9 # for lf == 1 (initial time step) no filter applied (w1=w2=0) # for lf == 2 (later steps) Robert+William's filter is applied - w1 = lf == 1 ? zero(NF) : robert_filter*williams_filter/two # = ν*α/2 in William (2009, Eq. 8) - w2 = lf == 1 ? zero(NF) : robert_filter*(1-williams_filter)/two # = ν(1-α)/2 in William (2009, Eq. 9) + w1 = lf == 1 ? zero(NF) : robert_filter*william_filter/two # = ν*α/2 in William (2009, Eq. 8) + w2 = lf == 1 ? zero(NF) : robert_filter*(1-william_filter)/two # = ν(1-α)/2 in William (2009, Eq. 9) @inbounds for lm in eachharmonic(A_old,A_new,A_lf,tendency) a_old = A_old[lm] # double filtered value from previous time step (t-Δt) @@ -45,8 +115,8 @@ end # variables that are leapfrogged in the respective models that are on layers (so excl surface pressure) leapfrog_layer_vars(::Barotropic) = (:vor,) leapfrog_layer_vars(::ShallowWater) = (:vor, :div) -leapfrog_layer_vars(::PrimitiveDryCore) = (:vor, :div, :temp) -leapfrog_layer_vars(::PrimitiveWetCore) = (:vor, :div, :temp, :humid) +leapfrog_layer_vars(::PrimitiveDry) = (:vor, :div, :temp) +leapfrog_layer_vars(::PrimitiveWet) = (:vor, :div, :temp, :humid) function leapfrog!( progn::PrognosticLayerTimesteps, diagn::DiagnosticVariablesLayer, @@ -58,141 +128,110 @@ function leapfrog!( progn::PrognosticLayerTimesteps, var_old = getproperty(progn.timesteps[1],var) var_new = getproperty(progn.timesteps[2],var) var_tend = getproperty(diagn.tendencies,Symbol(var,:_tend)) - leapfrog!(var_old,var_new,var_tend,dt,lf,model.constants) + leapfrog!(var_old,var_new,var_tend,dt,lf,model.time_stepping) end end """ - first_timesteps!( progn::PrognosticVariables, # all prognostic variables - diagn::DiagnosticVariables, # all pre-allocated diagnostic variables - time::DateTime, # time at timestep - M::ModelSetup, # everything that is constant at runtime - feedback::AbstractFeedback # feedback struct - ) - +$(TYPEDSIGNATURES) Performs the first two initial time steps (Euler forward, unfiltered leapfrog) to populate the prognostic variables with two time steps (t=0,Δt) that can then be used in the normal leap frogging.""" function first_timesteps!( progn::PrognosticVariables, # all prognostic variables diagn::DiagnosticVariables, # all pre-allocated diagnostic variables - time::DateTime, # time at timestep + clock::Clock, # current model time at timestep model::ModelSetup, # everything that is constant at runtime - feedback::AbstractFeedback, # feedback struct - outputter::AbstractOutput - ) + output::AbstractOutputWriter) - (; n_timesteps, Δt, Δt_sec ) = model.constants - n_timesteps == 0 && return time # exit immediately for no time steps + (;implicit) = model + (;Δt, Δt_sec) = model.time_stepping + clock.n_timesteps == 0 && return time # exit immediately for no time steps # FIRST TIME STEP (EULER FORWARD with dt=Δt/2) i = 1 # time step index lf1 = 1 # without Robert+William's filter lf2 = 1 # evaluates all tendencies at t=0, # the first leapfrog index (=>Euler forward) - temperature_profile!(diagn,progn,model,lf2) # used for implicit solver, update occasionally - initialize_implicit!(model,diagn,Δt/2) # update precomputed implicit terms with time step Δt/2 - timestep!(progn,diagn,time,Δt/2,i,model,lf1,lf2) - time += Dates.Second(Δt_sec÷2) # update by half the leapfrog time step Δt used here - progress!(feedback,progn) + initialize!(implicit,Δt/2,diagn,model) # update precomputed implicit terms with time step Δt/2 + timestep!(progn,diagn,clock.time,Δt/2,i,model,lf1,lf2) + clock.time += Dates.Second(Δt_sec÷2) # update by half the leapfrog time step Δt used here # SECOND TIME STEP (UNFILTERED LEAPFROG with dt=Δt, leapfrogging from t=0 over t=Δt/2 to t=Δt) - initialize_implicit!(model,diagn,Δt) # update precomputed implicit terms with time step Δt + initialize!(implicit,Δt,diagn,model) # update precomputed implicit terms with time step Δt lf1 = 1 # without Robert+William's filter lf2 = 2 # evaluate all tendencies at t=dt/2, # the 2nd leapfrog index (=>Leapfrog) - timestep!(progn,diagn,time,Δt,i,model,lf1,lf2) - time += Dates.Second(Δt_sec÷2) # now 2nd leapfrog step is at t=Δt - progress!(feedback,progn) - write_netcdf_output!(outputter,time,diagn,model) + timestep!(progn,diagn,clock.time,Δt,i,model,lf1,lf2) + clock.time += Dates.Second(Δt_sec÷2) # now 2nd leapfrog step is at t=Δt + write_output!(output,clock.time,diagn) return time end """ - timestep!( progn::PrognosticVariables, # all prognostic variables - diagn::DiagnosticVariables, # all pre-allocated diagnostic variables - time::DateTime, # time at timestep - dt::Real, # time step (mostly =2Δt, but for init steps =Δt,Δt/2) - lf1::Int=2, # leapfrog index 1 (dis/enables Robert+William's filter) - lf2::Int=2, # leapfrog index 2 (time step used for tendencies) - M::BarotropicModel, # everything that's constant at runtime - ) - -Calculate a single time step for the barotropic vorticity equation model of SpeedyWeather.jl """ +$(TYPEDSIGNATURES) +Calculate a single time step for the `model <: Barotropic`.""" function timestep!( progn::PrognosticVariables, # all prognostic variables diagn::DiagnosticVariables, # all pre-allocated diagnostic variables time::DateTime, # time at time step dt::Real, # time step (mostly =2Δt, but for init steps =Δt,Δt/2) i::Integer, # time step index - M::BarotropicModel, # everything that's constant at runtime + model::Barotropic, # everything that's constant at runtime lf1::Int=2, # leapfrog index 1 (dis/enables Robert+William's filter) - lf2::Int=2, # leapfrog index 2 (time step used for tendencies) - ) + lf2::Int=2) # leapfrog index 2 (time step used for tendencies) - # LOOP OVER LAYERS FOR DIFFUSION, LEAPFROGGING AND PROPAGATE STATE TO GRID + zero_tendencies!(diagn) + + # LOOP OVER LAYERS FOR TENDENCIES, DIFFUSION, LEAPFROGGING AND PROPAGATE STATE TO GRID for (progn_layer,diagn_layer) in zip(progn.layers,diagn.layers) - dynamics_tendencies!(diagn_layer,M) # tendency of vorticity - horizontal_diffusion!(progn_layer,diagn_layer,M) # diffusion for vorticity - leapfrog!(progn_layer,diagn_layer,dt,lf1,M) # leapfrog vorticity forward - gridded!(diagn_layer,progn_layer,lf2,M) # propagate spectral state to grid + dynamics_tendencies!(diagn_layer,time,model) + horizontal_diffusion!(diagn_layer,progn_layer,model) + leapfrog!(progn_layer,diagn_layer,dt,lf1,model) + gridded!(diagn_layer,progn_layer,lf2,model) end end """ - timestep!( progn::PrognosticVariables{NF}, # all prognostic variables - diagn::DiagnosticVariables{NF}, # all pre-allocated diagnostic variables - time::DateTime, # time at timestep - dt::Real, # time step (mostly =2Δt, but for init steps =Δt,Δt/2) - M::ShallowWaterModel, # everything that's constant at runtime - lf1::Int=2, # leapfrog index 1 (dis/enables Robert+William's filter) - lf2::Int=2 # leapfrog index 2 (time step used for tendencies) - ) where {NF<:AbstractFloat} - -Calculate a single time step for the shallow water model of SpeedyWeather.jl """ +$(TYPEDSIGNATURES) +Calculate a single time step for the `model <: ShallowWater`.""" function timestep!( progn::PrognosticVariables{NF}, # all prognostic variables diagn::DiagnosticVariables{NF}, # all pre-allocated diagnostic variables time::DateTime, # time at timestep dt::Real, # time step (mostly =2Δt, but for init steps =Δt,Δt/2) i::Integer, # time step index - M::ShallowWaterModel, # everything that's constant at runtime + model::ShallowWater, # everything that's constant at runtime lf1::Int=2, # leapfrog index 1 (dis/enables Robert+William's filter) lf2::Int=2 # leapfrog index 2 (time step used for tendencies) ) where {NF<:AbstractFloat} - # IMPLICIT, DIFFUSION, LEAPFROGGING AND PROPAGATE STATE TO GRID - progn_layer = progn.layers[1] # only calculate tendencies for the first layer + progn_layer = progn.layers[1] # only calculate tendencies for the first layer diagn_layer = diagn.layers[1] diagn_surface = diagn.surface progn_surface = progn.surface (;pres) = progn.surface.timesteps[lf2] + (;implicit, time_stepping, spectral_transform) = model + + zero_tendencies!(diagn) # GET TENDENCIES, CORRECT THEM FOR SEMI-IMPLICIT INTEGRATION - dynamics_tendencies!(diagn_layer,diagn_surface,pres,time,M) - implicit_correction!(diagn_layer,progn_layer,diagn_surface,progn_surface,M) + dynamics_tendencies!(diagn_layer,diagn_surface,pres,time,model) + implicit_correction!(diagn_layer,progn_layer,diagn_surface,progn_surface,implicit) # APPLY DIFFUSION, STEP FORWARD IN TIME, AND TRANSFORM NEW TIME STEP TO GRID - horizontal_diffusion!(progn_layer,diagn_layer,M) # diffusion for vorticity and divergence - leapfrog!(progn_layer,diagn_layer,dt,lf1,M) # leapfrog vorticity forward - gridded!(diagn_layer,progn_layer,lf2,M) # propagate spectral state to grid + horizontal_diffusion!(progn_layer,diagn_layer,model) + leapfrog!(progn_layer,diagn_layer,dt,lf1,model) + gridded!(diagn_layer,progn_layer,lf2,model) # SURFACE LAYER (pressure), no diffusion though (;pres_grid,pres_tend) = diagn.surface pres_old = progn.surface.timesteps[1].pres pres_new = progn.surface.timesteps[2].pres - leapfrog!(pres_old,pres_new,pres_tend,dt,lf1,M.constants) - gridded!(pres_grid,pres,M.spectral_transform) + leapfrog!(pres_old,pres_new,pres_tend,dt,lf1,time_stepping) + gridded!(pres_grid,pres,spectral_transform) end """ - timestep!( progn::PrognosticVariables{NF}, # all prognostic variables - diagn::DiagnosticVariables{NF}, # all pre-allocated diagnostic variables - time::DateTime, # time at timestep - dt::Real, # time step (mostly =2Δt, but for init steps =Δt,Δt/2) - M::PrimitiveEquation, # everything that's constant at runtime - lf1::Int=2, # leapfrog index 1 (dis/enables Robert+William's filter) - lf2::Int=2 # leapfrog index 2 (time step used for tendencies) - ) where {NF<:AbstractFloat} - -Calculate a single time step for the primitive equation model of SpeedyWeather.jl """ +$(TYPEDSIGNATURES) +Calculate a single time step for the `model<:PrimitiveEquation`""" function timestep!( progn::PrognosticVariables{NF}, # all prognostic variables diagn::DiagnosticVariables{NF}, # all pre-allocated diagnostic variables time::DateTime, # time at timestep @@ -204,15 +243,14 @@ function timestep!( progn::PrognosticVariables{NF}, # all prognostic variables ) where {NF<:AbstractFloat} # switch on/off all physics - (;physics) = model.parameters - physics && parameterization_tendencies!(diagn,time,model) - physics || zero_tendencies!(diagn) # set tendencies to zero otherwise + model.physics && parameterization_tendencies!(diagn,time,model) + model.physics || zero_tendencies!(diagn) # set tendencies to zero otherwise # occasionally reinitialize the implicit solver with new temperature profile - initialize_implicit!(model,diagn,progn,dt,i,lf2) + initialize!(model.implicit,i,dt,diagn,model.geometry,model.constants) - dynamics_tendencies!(diagn,progn,model,lf2) # dynamical core - implicit_correction!(diagn,progn,model) # semi-implicit time stepping corrections + dynamics_tendencies!(diagn,progn,model,lf2) # dynamical core + implicit_correction!(diagn,model.implicit,progn) # semi-implicit time stepping corrections # LOOP OVER ALL LAYERS for diffusion, leapfrog time integration # and progn state from spectral to grid for next time step @@ -224,56 +262,56 @@ function timestep!( progn::PrognosticVariables{NF}, # all prognostic variables horizontal_diffusion!(progn_layer,diagn_layer,model) # implicit diffusion of vor, div, temp leapfrog!(progn_layer,diagn_layer,dt,lf1,model) # time step forward for vor, div, temp gridded!(diagn_layer,progn_layer,lf2,model) # propagate spectral state to grid - else # surface level + else # surface level (;pres_grid,pres_tend) = diagn.surface pres_old = progn.surface.timesteps[1].pres pres_new = progn.surface.timesteps[2].pres pres_lf = progn.surface.timesteps[lf2].pres - leapfrog!(pres_old,pres_new,pres_tend,dt,lf1,model.constants) + leapfrog!(pres_old,pres_new,pres_tend,dt,lf1,model.time_stepping) gridded!(pres_grid,pres_lf,model.spectral_transform) end end end """ - time_stepping!( progn::PrognosticVariables, # all prognostic variables - diagn::DiagnosticVariables, # all pre-allocated diagnostic variables - model::ModelSetup) # all precalculated structs - -Main time loop that that initialises output and feedback, loops over all time steps +$(TYPEDSIGNATURES) +Main time loop that that initializes output and feedback, loops over all time steps and calls the output and feedback functions.""" function time_stepping!(progn::PrognosticVariables, # all prognostic variables diagn::DiagnosticVariables, # all pre-allocated diagnostic variables model::ModelSetup) # all precalculated structs - (; n_timesteps, Δt, Δt_sec ) = model.constants - time = model.parameters.startdate + (; Δt, Δt_sec ) = model.time_stepping + (; clock, time_stepping ) = model # SCALING: we use vorticity*radius,divergence*radius in the dynamical core - scale!(progn,model) + scale!(progn,model.spectral_grid.radius) # OUTPUT INITIALISATION AND STORING INITIAL CONDITIONS + FEEDBACK # propagate spectral state to grid variables for initial condition output + (;output,feedback) = model lf = 1 # use first leapfrog index gridded!(diagn,progn,lf,model) - outputter = initialize_netcdf_output(diagn,model) - feedback = initialize_feedback(outputter,model) + initialize!(output,feedback,time_stepping,diagn,model) + initialize!(feedback,clock,model) # FIRST TIMESTEPS: EULER FORWARD THEN 1x LEAPFROG - time = first_timesteps!(progn,diagn,time,model,feedback,outputter) - initialize_implicit!(model,diagn,2Δt) # from now on precomputed implicit terms with 2Δt + first_timesteps!(progn,diagn,clock,model,output) + initialize!(model.implicit,2Δt,diagn,model) # from now on precomputed implicit terms with 2Δt # MAIN LOOP - for i in 2:n_timesteps # start at 2 as first Δt in first_timesteps! - timestep!(progn,diagn,time,2Δt,i,model) # calculate tendencies and leapfrog forward - time += Dates.Second(Δt_sec) # time of lf=2 and diagn after timestep! + for i in 2:clock.n_timesteps # start at 2 as first Δt in first_timesteps! + + # calculate tendencies and leapfrog forward + timestep!(progn,diagn,clock.time,2Δt,i,model) + clock.time += Dates.Second(Δt_sec) # time of lf=2 and diagn after timestep! progress!(feedback,progn) # updates the progress meter bar - write_netcdf_output!(outputter,time,diagn,model) + write_output!(output,clock.time,diagn) end - unscale!(progn,model) # unscale radius-scaling from the dynamical core - write_restart_file(time,progn,outputter,model) + unscale!(progn) # undo radius-scaling for vor,div from the dynamical core + write_restart_file(clock.time,progn,output) progress_finish!(feedback) # finishes the progress meter bar return progn diff --git a/src/dynamics/vertical_coordinates.jl b/src/dynamics/vertical_coordinates.jl new file mode 100644 index 000000000..ba6d73f4c --- /dev/null +++ b/src/dynamics/vertical_coordinates.jl @@ -0,0 +1,78 @@ +@with_kw struct NoVerticalCoordinates <: VerticalCoordinates + nlev::Int = 1 +end + +@with_kw struct SigmaCoordinates <: VerticalCoordinates + nlev::Int = 8 + σ_half::Vector{Float64} = default_sigma_coordinates(nlev) + + SigmaCoordinates(nlev::Integer,σ_half::AbstractVector) = sigma_okay(nlev,σ_half) ? + new(nlev,σ_half) : error("σ_half = $σ_half cannot be used for $nlev-level SigmaCoordinates") +end + +# obtain nlev from length of predefined σ_half levels +SigmaCoordinates(σ_half::AbstractVector) = SigmaCoordinates(nlev=length(σ_half)-1;σ_half) + +"""Coefficients of the generalised logistic function to describe the vertical coordinate. +Default coefficients A,K,C,Q,B,M,ν are fitted to the old L31 configuration at ECMWF. + +Following the notation of [https://en.wikipedia.org/wiki/Generalised_logistic_function](https://en.wikipedia.org/wiki/Generalised_logistic_function) (Dec 15 2021). + +Change default parameters for more/fewer levels in the stratosphere vs troposphere vs boundary layer.""" +@with_kw struct GenLogisticCoefs + A::Float64 = -0.283 # obtained from a fit in /input_data/vertical_coordinate/vertical_resolution.ipynb + K::Float64 = 0.871 + C::Float64 = 0.414 + Q::Float64 = 6.695 + B::Float64 = 10.336 + M::Float64 = 0.602 + ν::Float64 = 5.812 +end + +"""Generalised logistic function based on the coefficients in `coefs`.""" +function generalised_logistic(x,coefs::GenLogisticCoefs) + (; A,K,C,Q,B,M,ν ) = coefs + return @. A + (K-A)/(C+Q*exp(-B*(x-M)))^inv(ν) +end + +""" +$(TYPEDSIGNATURES) +Vertical sigma coordinates defined by their nlev+1 half levels `σ_levels_half`. Sigma coordinates are +fraction of surface pressure (p/p0) and are sorted from top (stratosphere) to bottom (surface). +The first half level is at 0 the last at 1. Evaluate a generalised logistic function with +coefficients in `P` for the distribution of values in between. Default coefficients follow +the L31 configuration historically used at ECMWF.""" +function default_sigma_coordinates(nlev::Integer) + GLcoefs = GenLogisticCoefs() + z = range(0,1,nlev+1) # normalised = level/nlev + σ_half = generalised_logistic(z,GLcoefs) + σ_half .-= σ_half[1] # topmost half-level is at 0 pressure + σ_half ./= σ_half[end] # lowermost half-level is at 1, i.e. p=p_surface + return σ_half +end + +""" +$(TYPEDSIGNATURES) +Check that nlev and σ_half match.""" +function sigma_okay(nlev::Integer,σ_half::AbstractVector) + @assert σ_half[1] >= 0 "First manually specified σ_half has to be >0" + @assert σ_half[end] == 1 "Last manually specified σ_half has to be 1." + @assert nlev == (length(σ_half) - 1) "nlev has to be length of σ_half - 1" + @assert isincreasing(σ_half) "Vertical sigma coordinates are not increasing." + return true +end + +#TODO +@with_kw struct SigmaPressureCoordinates <: VerticalCoordinates + nlev::Int = 8 + A::Vector{Float64} = default_hybrid_coordinates(:A,nlev) + B::Vector{Float64} = default_hybrid_coordinates(:B,nlev) +end + +default_vertical_coordinates(::Type{<:Barotropic}) = NoVerticalCoordinates +default_vertical_coordinates(::Type{<:ShallowWater}) = NoVerticalCoordinates +default_vertical_coordinates(::Type{<:PrimitiveEquation}) = SigmaCoordinates + +default_nlev(::Type{<:Barotropic}) = 1 +default_nlev(::Type{<:ShallowWater}) = 1 +default_nlev(::Type{<:PrimitiveEquation}) = 8 diff --git a/src/output/feedback.jl b/src/output/feedback.jl index 64ff0bae9..72ec2a8c6 100644 --- a/src/output/feedback.jl +++ b/src/output/feedback.jl @@ -1,66 +1,97 @@ +""" +Feedback struct that contains options and object for command-line feedback +like the progress meter. +$(TYPEDFIELDS)""" mutable struct Feedback <: AbstractFeedback - verbose::Bool # print feedback to REPL? - debug::Bool # run nan_detection code? - output::Bool # store output (here: write to parameters and progress.txt?) + "print feedback to REPL?" + verbose::Bool - # PROGRESS - progress_meter::ProgressMeter.Progress # struct containing everything progress related - progress_txt::Union{IOStream,Nothing} # txt is a Nothing in case of no output - - # NaRS (Not-a-Real) AND OTHER MODEL STATE FEEDBACK - nars_detected::Bool # did Infs/NaNs occur in the simulation? -end + "check for NaRs in the prognostic variables" + debug::Bool + + "write a progress.txt file? State synced with OutputWriter.output" + output::Bool -"""Initialises the progress txt file.""" -function initialize_feedback(outputter::Output,M::ModelSetup) - (; output, write_restart ) = outputter - (; run_id, run_path ) = outputter + "identification of run, taken from ::OutputWriter" + id::Union{String,Int} - if output # with netcdf output write parameters.txt and progress.txt - (; NF, n_days, trunc ) = M.parameters - (; Grid, npoints, nlat_half ) = M.geometry + "path to run folder, taken from ::OutputWriter" + run_path::String - # create progress.txt file in run????/ - progress_txt = open(joinpath(run_path,"progress.txt"),"w") - s = "Starting SpeedyWeather.jl run $run_id on "* - Dates.format(Dates.now(),Dates.RFC1123Format) - write(progress_txt,s*"\n") # and in file + # PROGRESS + "struct containing everything progress related" + progress_meter::ProgressMeter.Progress - # add some information on resolution and number format - write(progress_txt,"Integrating $(n_days) days at a spectral resolution of "* - "T$trunc on a $(get_nlat(Grid,nlat_half))-ring $Grid with $npoints grid points.\n") - write(progress_txt,"Number format is "*string(NF)*".\n") - write(progress_txt,"All data will be stored in $run_path.\n") + "txt is a Nothing in case of no output" + progress_txt::Union{IOStream,Nothing} - # also export parameters into run????/parameters.txt - parameters_txt = open(joinpath(run_path,"parameters.txt"),"w") - print(parameters_txt,M.parameters) - close(parameters_txt) + "did Infs/NaNs occur in the simulation?" + nars_detected::Bool +end - else # no netcdf output - progress_txt = nothing # for no ouput, allocate dummies for Feedback struct - end +""" +$(TYPEDSIGNATURES) +Generator function for a Feedback struct.""" +function Feedback(verbose::Bool=true,debug::Bool=true) + + # the following are synced with OutputWriter in + # initialize!(::OutputWriter, ...) to avoid folder-race conditions + output = false + id = "" + run_path = "" + + # PROGRESSMETER + # show progress meter via `enabled` through verbose parameter, initialize only for 1 time step + desc = "Weather is speedy: " + progress_meter = ProgressMeter.Progress(1, enabled=verbose, showspeed=true; desc) + progress_txt = nothing # initialize with nothing, initialize in initialize!(::Feedback,...) + + nars_detected = false + return Feedback(verbose,debug, + output,id,run_path, + progress_meter,progress_txt, + nars_detected) +end - nans_detected = false # don't check again if true +""" +$(TYPEDSIGNATURES) +Initializes the a `Feedback` struct.""" +function initialize!(feedback::Feedback,clock::Clock,model::ModelSetup) - # PROGRESSMETER - (; verbose, debug ) = M.parameters - (; n_timesteps ) = M.constants - DT_IN_SEC[] = M.constants.Δt_sec # hack: redefine element in global constant dt_in_sec - # used to pass on the time step to ProgressMeter.speedstring - desc = "Weather is speedy$(output ? " run $run_id: " : ": ")" + # reinitalize progress meter, minus one to exclude first_timesteps! which contain compilation + (;enabled, showspeed, desc) = feedback.progress_meter + feedback.progress_meter = ProgressMeter.Progress(clock.n_timesteps-1;enabled,showspeed, desc) - # show progress meter via `enabled` through verbose parameter - # one more time steps for the first Euler time step in first_timesteps! - progress_meter = ProgressMeter.Progress(n_timesteps+1, enabled=verbose, showspeed=true; desc) + # hack: redefine element in global constant dt_in_sec + # used to pass on the time step to ProgressMeter.speedstring + DT_IN_SEC[] = model.time_stepping.Δt_sec - return Feedback(verbose,debug,output,progress_meter,progress_txt,nans_detected) + if feedback.output # with netcdf output write progress.txt + (; run_path, id) = feedback + SG = model.spectral_grid + L = model.time_stepping + + # create progress.txt file in run_????/ + progress_txt = open(joinpath(run_path,"progress.txt"),"w") + s = "Starting SpeedyWeather.jl run $id on "* + Dates.format(Dates.now(),Dates.RFC1123Format) + write(progress_txt,s*"\n") + write(progress_txt,"Integrating:\n") + write(progress_txt,"$SG\n") + write(progress_txt,"Time: $(clock.n_days) days at Δt = $(L.Δt_sec)s\n") + write(progress_txt,"\nAll data will be stored in $run_path\n") + feedback.progress_txt = progress_txt + end end -"""Calls the progress meter and writes every 5% progress increase to txt.""" +""" +$(TYPEDSIGNATURES) +Calls the progress meter and writes every 5% progress increase to txt.""" function progress!(feedback::Feedback) - ProgressMeter.next!(feedback.progress_meter) # update progress meter - (; counter, n ) = feedback.progress_meter # unpack counter after update + + # update progress meter and unpack counter after update + ProgressMeter.next!(feedback.progress_meter) + (; counter, n ) = feedback.progress_meter # write progress to txt file too if (counter/n*100 % 1) > ((counter+1)/n*100 % 1) @@ -86,7 +117,9 @@ function progress!( feedback::Feedback, feedback.debug && nar_detection!(feedback,progn) end -"""Finalises the progress meter and the progress txt file.""" +""" +$(TYPEDSIGNATURES) +Finalises the progress meter and the progress txt file.""" function progress_finish!(F::Feedback) ProgressMeter.finish!(F.progress_meter) @@ -98,7 +131,9 @@ function progress_finish!(F::Feedback) end end -"""Detect NaR (Not-a-Real) in the prognostic variables.""" +""" +$(TYPEDSIGNATURES) +Detect NaR (Not-a-Real) in the prognostic variables.""" function nar_detection!(feedback::Feedback,progn::PrognosticVariables) feedback.nars_detected && return nothing # escape immediately if nans already detected @@ -116,7 +151,10 @@ function nar_detection!(feedback::Feedback,progn::PrognosticVariables) feedback.nars_detected |= nars_detected_here end -# adapted from ProgressMeter.jl +""" +$(TYPEDSIGNATURES) +Estimates the remaining time from a `ProgresssMeter.Progress`. +Adapted from ProgressMeter.jl""" function remaining_time(p::ProgressMeter.Progress) elapsed_time = time() - p.tinit est_total_time = elapsed_time * (p.n - p.start) / (p.counter - p.start) @@ -128,3 +166,59 @@ function remaining_time(p::ProgressMeter.Progress) end return eta end + +""" +$(TYPEDSIGNATURES) +define a ProgressMeter.speedstring method that also takes a time step +`dt_in_sec` to translate sec/iteration to days/days-like speeds.""" +function speedstring(sec_per_iter,dt_in_sec) + if sec_per_iter == Inf + return " N/A days/day" + end + + sim_time_per_time = dt_in_sec/sec_per_iter + + for (divideby, unit) in ( (365*1_000, "millenia"), + (365, "years"), + (1, "days"), + (1/24, "hours")) + if (sim_time_per_time / divideby) > 2 + return @sprintf "%5.2f %2s/day" (sim_time_per_time / divideby) unit + end + end + return " <2 hours/days" +end + +# hack: define global constant whose element will be changed in initialize_feedback +# used to pass on the time step to ProgressMeter.speedstring via calling this +# constant from the ProgressMeter module +const DT_IN_SEC = Ref(1800) + +# overwrite the speedstring function from ProgressMeter +function ProgressMeter.speedstring(sec_per_iter,dt_in_sec=SpeedyWeather.DT_IN_SEC) + speedstring(sec_per_iter,dt_in_sec[]) +end + +""" +$(TYPEDSIGNATURES) +Returns `Dates.CompoundPeriod` rounding to either (days, hours), (hours, minutes), (minutes, +seconds), or seconds with 1 decimal place accuracy for >10s and two for less. +E.g. +```julia +julia> readable_secs(12345) +3 hours, 26 minutes +``` +""" +function readable_secs(secs::Real) + millisecs = Dates.Millisecond(round(secs * 10 ^ 3)) + if millisecs >= Dates.Day(1) + return Dates.canonicalize(round(millisecs, Dates.Hour)) + elseif millisecs >= Dates.Hour(1) + return Dates.canonicalize(round(millisecs, Dates.Minute)) + elseif millisecs >= Dates.Minute(1) + return Dates.canonicalize(round(millisecs, Dates.Second)) + elseif millisecs >= Dates.Second(10) + return Dates.canonicalize(round(millisecs, Dates.Millisecond(100))) + end + return Dates.canonicalize(round(millisecs, Dates.Millisecond(10))) +end \ No newline at end of file diff --git a/src/output/output.jl b/src/output/output.jl index 26ed0288b..faf60a4a7 100644 --- a/src/output/output.jl +++ b/src/output/output.jl @@ -1,241 +1,298 @@ -Base.@kwdef mutable struct Output{NF<:Union{Float32,Float64}} <: AbstractOutput - # NF: output only in Float32/64 +""" +Number of mantissa bits to keep for each prognostic variable when compressed for +netCDF and .jld2 data output. +$(TYPEDFIELDS)""" +@with_kw struct Keepbits + u::Int = 7 + v::Int = 7 + vor::Int = 5 + div::Int = 5 + temp::Int = 10 + pres::Int = 12 + humid::Int = 7 +end + +function Base.show(io::IO,K::Keepbits) + print(io,"$(typeof(K))(") + for key in propertynames(K) + val = getfield(K,key) + print(io,"$key=$val, ") + end + print(")") +end +# default number format for output +const DEFAULT_OUTPUT_NF = Float32 + +""" +$(TYPEDSIGNATURES) +NetCDF output writer. Contains all output options and auxiliary fields for output interpolation. +To be initialised with `OutputWriter(::SpectralGrid,::TimeStepper,kwargs...)` to pass on the +resolution/time stepping information from those structs. Options include +$(TYPEDFIELDS)""" +Base.@kwdef mutable struct OutputWriter{NF<:Union{Float32,Float64},Model<:ModelSetup} <: AbstractOutputWriter + + spectral_grid::SpectralGrid{Model} + + # FILE OPTIONS output::Bool = false # output to netCDF? - output_vars::Vector{Symbol}=[:none] # vector of output variables as Symbols - write_restart::Bool = false # also write restart file if output==true? - + path::String = pwd() # path to output folder + id::Union{String,Int} = "" # run identification number/string + run_path::String = "" # will be determined in initalize! + filename::String = "output.nc" # name of the output netcdf file + write_restart::Bool = true # also write restart file if output==true? + pkg_version::VersionNumber = pkgversion(SpeedyWeather) + + # WHAT/WHEN OPTIONS startdate::DateTime = DateTime(2000,1,1) + output_dt::Float64 = 6 # output time step [hours] + output_vars::Vector{Symbol} = default_output_vars(Model) # vector of output variables as Symbols + missing_value::NF = NaN # missing value to be used in netcdf output + + # COMPRESSION OPTIONS + compression_level::Int = 3 # compression level; 1=low but fast, 9=high but slow + keepbits::Keepbits = Keepbits() # mantissa bits to keep for every variable + + # TIME STEPS AND COUNTERS (initialize later) + output_every_n_steps::Int = 0 # output frequency timestep_counter::Int = 0 # time step counter output_counter::Int = 0 # output step counter - n_timesteps::Int = 0 # number of time steps - n_outputsteps::Int = 0 # number of time steps with output - output_every_n_steps::Int = 0 # output every n time steps - run_id::String = "-1" # run identification number - run_path::String = "" # output path plus run????/ - - # the netcdf file to be written into + # the netcdf file to be written into, will be create netcdf_file::Union{NcFile,Nothing} = nothing - # grid specifications - output_matrix::Bool = false # if true sort grid points into a matrix (interpolation-free) - # full grid for output if output_matrix == false - output_Grid::Type{<:AbstractFullGrid} = FullGaussianGrid - nlat_half::Int = 0 # size of the input/output grid - nlon::Int = 0 # number of longitude/latitude points - nlat::Int = 0 + # INPUT GRID (the one used in the dynamical core) + input_Grid::Type{<:AbstractGrid} = spectral_grid.Grid - # input grid for interpolation - input_Grid::Type{<:AbstractGrid} = FullGaussianGrid - interpolator::AbstractInterpolator = DEFAULT_INTERPOLATOR(input_Grid,0,0) + # Output as matrix (particularly for reduced grids) + as_matrix::Bool = false # if true sort grid points into a matrix (interpolation-free) + # full grid for output if output_matrix == false + quadrant_rotation::NTuple{4,Int} = (0,1,2,3) # rotation of output quadrant + # matrix of output quadrant + matrix_quadrant::NTuple{4,Tuple{Int,Int}} = ((2,2),(1,2),(1,1),(2,1)) + # OUTPUT GRID + output_Grid::Type{<:AbstractFullGrid} = RingGrids.full_grid(input_Grid) + nlat_half::Int = spectral_grid.nlat_half # default: same nlat_half for in/output + nlon::Int = as_matrix ? RingGrids.matrix_size(input_Grid,spectral_grid.nlat_half)[1] : + RingGrids.get_nlon(output_Grid,nlat_half) + nlat::Int = as_matrix ? RingGrids.matrix_size(input_Grid,spectral_grid.nlat_half)[2] : + RingGrids.get_nlat(output_Grid,nlat_half) + npoints::Int = nlon*nlat + nlev::Int = spectral_grid.nlev + interpolator::AbstractInterpolator = DEFAULT_INTERPOLATOR(input_Grid,spectral_grid.nlat_half,npoints) + # fields to output (only one layer, reuse over layers) - u::Matrix{NF} = zeros(0,0) # zonal velocity - v::Matrix{NF} = zeros(0,0) # meridional velocity - vor::Matrix{NF} = zeros(0,0) # relative vorticity - div::Matrix{NF} = zeros(0,0) # divergence - pres::Matrix{NF} = zeros(0,0) # pressure - temp::Matrix{NF} = zeros(0,0) # temperature - humid::Matrix{NF} = zeros(0,0) # humidity + const u::Matrix{NF} = fill(missing_value,nlon,nlat) + const v::Matrix{NF} = fill(missing_value,nlon,nlat) + const vor::Matrix{NF} = fill(missing_value,nlon,nlat) + const div::Matrix{NF} = fill(missing_value,nlon,nlat) + const temp::Matrix{NF} = fill(missing_value,nlon,nlat) + const pres::Matrix{NF} = fill(missing_value,nlon,nlat) + const humid::Matrix{NF} = fill(missing_value,nlon,nlat) end -Output() = Output{Float32}() - -""" - run_id = get_run_id(output, output_path) - -Checks existing `run-????` folders in output path to determine a 4-digit `run_id` number. -""" -function get_run_id(output, output_path) - - if output - # pull list of existing run???? folders via readdir - pattern = r"run-\d\d\d\d" # run-???? in regex - runlist = filter(x->startswith(x,pattern),readdir(output_path)) - runlist = filter(x->endswith( x,pattern),runlist) - existing_runs = [parse(Int,id[5:end]) for id in runlist] - - # get the run id from existing folders - if length(existing_runs) == 0 # if no runfolder exists yet - run_id = 1 # start with run0001 - else - run_id = maximum(existing_runs)+1 # next run gets id +1 - end - - return @sprintf("%04d",run_id) - else - return "-1" - end +# generator function pulling grid resolution and time stepping from ::SpectralGrid and ::TimeStepper +function OutputWriter( + spectral_grid::SpectralGrid{Model}; + NF::Type{<:Union{Float32,Float64}} = DEFAULT_OUTPUT_NF, + kwargs... +) where Model + return OutputWriter{NF,Model}(;spectral_grid,kwargs...) end -""" - run_id, run_path = get_run_id_path(P::Parameters) - -Creates a new folder `run-*` with the `run_id`. Also returns the full path -`run_path` of that folder. Returns `-1, "no runpath"` in the case of no output. -""" -function get_run_id_path(P::Parameters) - - (; output, output_path, run_id ) = P - - if output - run_path = joinpath(output_path,string("run-",run_id_string(run_id))) - @assert !(string("run-",run_id) in readdir(output_path)) "Run folder already exists, choose another run_id." - - mkdir(run_path) # actually create the folder - return run_id, run_path - else - return run_id, "no runpath" +# default variables to output by model +default_output_vars(::Type{<:Barotropic}) = [:vor,:u] +default_output_vars(::Type{<:ShallowWater}) = [:vor,:u] +default_output_vars(::Type{<:PrimitiveDry}) = [:vor,:u,:temp,:pres] +default_output_vars(::Type{<:PrimitiveWet}) = [:vor,:u,:temp,:humid,:pres] + +# print all fields with type <: Number +function Base.show(io::IO,O::AbstractOutputWriter) + print(io,"$(typeof(O)):") + for key in propertynames(O) + val = getfield(O,key) + val isa Union{Number,String,DataType,NTuple,Vector{Symbol},UnionAll,Keepbits} && + print(io,"\n $key::$(typeof(val)) = $val") end end -run_id_string(run_id::Integer) = @sprintf("%04d",run_id) -run_id_string(run_id::String) = run_id - """ - outputter = initialize_netcdf_output( progn::PrognosticVariables, - diagn::DiagnosticVariables, - M::ModelSetup) - +$(TYPEDSIGNATURES) Creates a netcdf file on disk and the corresponding `netcdf_file` object preallocated with output variables -and dimensions. `write_netcdf_output!` then writes consecuitive time steps into this file. +and dimensions. `write_output!` then writes consecuitive time steps into this file. """ -function initialize_netcdf_output( diagn::DiagnosticVariables, - M::ModelSetup) - - M.parameters.output || return Output() # escape directly when no netcdf output - +function initialize!( + output::OutputWriter{output_NF,Model}, + feedback::AbstractFeedback, + time_stepping::TimeStepper, + diagn::DiagnosticVariables, + model::Model +) where {output_NF,Model} + + output.output || return nothing # exit immediately for no output + # DEFINE NETCDF DIMENSIONS TIME - (; startdate ) = M.parameters + (;startdate) = output time_string = "seconds since $(Dates.format(startdate, "yyyy-mm-dd HH:MM:0.0"))" dim_time = NcDim("time",0,unlimited=true) var_time = NcVar("time",dim_time,t=Int64,atts=Dict("units"=>time_string,"long_name"=>"time")) - + # DEFINE NETCDF DIMENSIONS SPACE - (; output_matrix, output_Grid, output_nlat_half, nlev ) = M.parameters - - # if specified (>0) use output resolution via output_nlat_half, otherwise nlat_half from dynamical core - nlat_half = output_nlat_half > 0 ? output_nlat_half : M.geometry.nlat_half - - if output_matrix == false # interpolate onto (possibly different) output grid + (;input_Grid, output_Grid, nlat_half, nlev) = output + + if output.as_matrix == false # interpolate onto (possibly different) output grid lond = get_lond(output_Grid,nlat_half) latd = get_latd(output_Grid,nlat_half) nlon = length(lond) nlat = length(latd) lon_name, lon_units, lon_longname = "lon","degrees_east","longitude" lat_name, lat_units, lat_longname = "lat","degrees_north","latitude" - - else # output grid directly into a matrix (resort grid points, no interpolation) - (; nlat_half ) = M.geometry # don't use output_nlat_half as not supported for output_matrix - nlon,nlat = RingGrids.matrix_size(M.geometry.Grid,nlat_half) # size of the matrix output - lond = collect(1:nlon) # just enumerate grid points for lond, latd + + else # output grid directly into a matrix (resort grid points, no interpolation) + (;nlat_half) = diagn # don't use output.nlat_half as not supported for output_matrix + nlon,nlat = RingGrids.matrix_size(input_Grid,nlat_half) # size of the matrix output + lond = collect(1:nlon) # just enumerate grid points for lond, latd latd = collect(1:nlat) lon_name, lon_units, lon_longname = "i","1","horizontal index i" lat_name, lat_units, lat_longname = "j","1","horizontal index j" end - σ = M.geometry.σ_levels_full + σ = model.geometry.σ_levels_full dim_lon = NcDim(lon_name,nlon,values=lond,atts=Dict("units"=>lon_units,"long_name"=>lon_longname)) dim_lat = NcDim(lat_name,nlat,values=latd,atts=Dict("units"=>lat_units,"long_name"=>lat_longname)) dim_lev = NcDim("lev",nlev,values=σ,atts=Dict("units"=>"1","long_name"=>"sigma levels")) - + # VARIABLES, define every variable here that could be output - (; output_NF, compression_level ) = M.parameters - missing_value = convert(output_NF,M.parameters.missing_value) - + (;compression_level) = output + missing_value = convert(output_NF,output.missing_value) + # given pres the right name, depending on ShallowWaterModel or PrimitiveEquationModel - pres_name = M isa ShallowWaterModel ? "interface displacement" : "surface pressure" - pres_unit = M isa ShallowWaterModel ? "m" : "hPa" - + pres_name = Model <: ShallowWater ? "interface displacement" : "surface pressure" + pres_unit = Model <: ShallowWater ? "m" : "hPa" + all_ncvars = ( # define NamedTuple to identify the NcVars by name time = var_time, u = NcVar("u",[dim_lon,dim_lat,dim_lev,dim_time],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>"zonal wind","units"=>"m/s","missing_value"=>missing_value, - "_FillValue"=>missing_value)), + atts=Dict("long_name"=>"zonal wind","units"=>"m/s","missing_value"=>missing_value, + "_FillValue"=>missing_value)), v = NcVar("v",[dim_lon,dim_lat,dim_lev,dim_time],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>"meridional wind","units"=>"m/s","missing_value"=>missing_value, - "_FillValue"=>missing_value)), + atts=Dict("long_name"=>"meridional wind","units"=>"m/s","missing_value"=>missing_value, + "_FillValue"=>missing_value)), vor = NcVar("vor",[dim_lon,dim_lat,dim_lev,dim_time],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>"relative vorticity","units"=>"1/s","missing_value"=>missing_value, - "_FillValue"=>missing_value)), + atts=Dict("long_name"=>"relative vorticity","units"=>"1/s","missing_value"=>missing_value, + "_FillValue"=>missing_value)), pres = NcVar("pres",[dim_lon,dim_lat,dim_time],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>pres_name,"units"=>pres_unit,"missing_value"=>missing_value, - "_FillValue"=>missing_value)), + atts=Dict("long_name"=>pres_name,"units"=>pres_unit,"missing_value"=>missing_value, + "_FillValue"=>missing_value)), div = NcVar("div",[dim_lon,dim_lat,dim_lev,dim_time],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>"divergence","units"=>"1/s","missing_value"=>missing_value, - "_FillValue"=>missing_value)), + atts=Dict("long_name"=>"divergence","units"=>"1/s","missing_value"=>missing_value, + "_FillValue"=>missing_value)), temp = NcVar("temp",[dim_lon,dim_lat,dim_lev,dim_time],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>"temperature","units"=>"degC","missing_value"=>missing_value, - "_FillValue"=>missing_value)), + atts=Dict("long_name"=>"temperature","units"=>"degC","missing_value"=>missing_value, + "_FillValue"=>missing_value)), humid = NcVar("humid",[dim_lon,dim_lat,dim_lev,dim_time],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>"specific humidity","units"=>"1","missing_value"=>missing_value, - "_FillValue"=>missing_value)), + atts=Dict("long_name"=>"specific humidity","units"=>"1","missing_value"=>missing_value, + "_FillValue"=>missing_value)), orog = NcVar("orog",[dim_lon,dim_lat],t=output_NF,compress=compression_level, - atts=Dict("long_name"=>"orography","units"=>"m","missing_value"=>missing_value, - "_FillValue"=>missing_value)) + atts=Dict("long_name"=>"orography","units"=>"m","missing_value"=>missing_value, + "_FillValue"=>missing_value)) ) + + # GET RUN ID, CREATE FOLDER + # get new id only if not already specified + output.id = output.id == "" ? get_run_id(output.path) : output.id + output.run_path = create_output_folder(output.path,output.id) + + feedback.id = output.id # synchronize with feedback struct + feedback.run_path = output.run_path + feedback.progress_meter.desc = "Weather is speedy: run $(output.id) " + feedback.output = true # if output=true set feedback.output=true too! - # CREATE NETCDF FILE - (; output_filename, output_vars, output_NF ) = M.parameters - run_id,run_path = get_run_id_path(M.parameters) # create output folder and get its id and path - - # vector of NcVars for output + + # CREATE NETCDF FILE, vector of NcVars for output + (; run_path, filename, output_vars) = output vars_out = [all_ncvars[key] for key in keys(all_ncvars) if key in vcat(:time,output_vars)] - netcdf_file = NetCDF.create(joinpath(run_path,output_filename),vars_out,mode=NetCDF.NC_NETCDF4) + output.netcdf_file = NetCDF.create(joinpath(run_path,filename),vars_out,mode=NetCDF.NC_NETCDF4) - # CREATE OUTPUT STRUCT - (; write_restart, startdate ) = M.parameters - (; n_timesteps, n_outputsteps, output_every_n_steps ) = M.constants - - u = :u in output_vars ? fill(missing_value,nlon,nlat) : zeros(output_NF,0,0) - v = :v in output_vars ? fill(missing_value,nlon,nlat) : zeros(output_NF,0,0) - vor = :vor in output_vars ? fill(missing_value,nlon,nlat) : zeros(output_NF,0,0) - div = :div in output_vars ? fill(missing_value,nlon,nlat) : zeros(output_NF,0,0) - pres = :pres in output_vars ? fill(missing_value,nlon,nlat) : zeros(output_NF,0,0) - temp = :temp in output_vars ? fill(missing_value,nlon,nlat) : zeros(output_NF,0,0) - humid = :humid in output_vars ? fill(missing_value,nlon,nlat) : zeros(output_NF,0,0) - - # CREATE OUTPUT INTERPOLATOR - input_Grid = M.parameters.Grid # grid to interpolate from (the on used in dyn core) - Interpolator = M.parameters.output_Interpolator # type of interpolator - npoints = get_npoints(output_Grid,nlat_half) # number of grid points to interpolate onto - input_nlat_half = M.geometry.nlat_half - interpolator = output_matrix ? Interpolator(input_Grid,0,0) : - Interpolator(output_NF,input_Grid,input_nlat_half,npoints) + # INTERPOLATION: PRECOMPUTE LOCATION INDICES + latds, londs = RingGrids.get_latdlonds(output_Grid,output.nlat_half) + output.as_matrix || RingGrids.update_locator!(output.interpolator,latds,londs) + + # OUTPUT FREQUENCY + output.output_every_n_steps = max(1,floor(Int,output.output_dt/time_stepping.Δt_hrs)) + + # RESET COUNTERS + output.timestep_counter = 0 # time step counter + output.output_counter = 0 # output step counter - # PRECOMPUTE LOCATION INDICES - latds, londs = get_latdlonds(output_Grid,nlat_half) - output_matrix || update_locator!(interpolator,latds,londs) - - outputter = Output( output=true; output_vars, write_restart, - startdate, n_timesteps, n_outputsteps, - output_every_n_steps, run_id, run_path, - netcdf_file, - output_matrix, output_Grid, nlat_half, nlon, nlat, - input_Grid, interpolator, - u, v, vor, div, pres, temp, humid) - # WRITE INITIAL CONDITIONS TO FILE - write_netcdf_variables!(outputter,diagn,M) - write_netcdf_time!(outputter,startdate) + write_netcdf_variables!(output,diagn) + write_netcdf_time!(output,startdate) + + # also export parameters into run????/parameters.txt + parameters_txt = open(joinpath(output.run_path,"parameters.txt"),"w") + println(parameters_txt,model.spectral_grid) + println(parameters_txt,model.planet) + println(parameters_txt,model.atmosphere) + println(parameters_txt,model.time_stepping) + println(parameters_txt,model.output) + println(parameters_txt,model.initial_conditions) + println(parameters_txt,model.horizontal_diffusion) + model isa Union{ShallowWater,PrimitiveEquation} && println(parameters_txt,model.implicit) + model isa Union{ShallowWater,PrimitiveEquation} && println(parameters_txt,model.orography) + close(parameters_txt) +end - return outputter +""" +$(TYPEDSIGNATURES) +Checks existing `run_????` folders in `path` to determine a 4-digit `id` number +by counting up. E.g. if folder run_0001 exists it will return the string "0002". +Does not create a folder for the returned run id. +""" +function get_run_id(path::String) + # pull list of existing run_???? folders via readdir + pattern = r"run_\d\d\d\d" # run_???? in regex + runlist = filter(x->startswith(x,pattern),readdir(path)) + runlist = filter(x->endswith( x,pattern),runlist) + existing_runs = [parse(Int,id[5:end]) for id in runlist] + + # get the run id from existing folders + if length(existing_runs) == 0 # if no runfolder exists yet + run_id = 1 # start with run_0001 + else + run_id = maximum(existing_runs)+1 # next run gets id +1 + end + + return @sprintf("%04d",run_id) +end + +""" +$(TYPEDSIGNATURES) +Creates a new folder `run_*` with the identification `id`. Also returns the full path +`run_path` of that folder. +""" +function create_output_folder(path::String,id::Union{String,Int}) + run_id = string("run_",run_id_to_string(id)) + run_path = joinpath(path,run_id) + @assert !(run_id in readdir(path)) "Run folder $run_path already exists." + mkdir(run_path) # actually create the folder + return run_path end -"""write_netcdf_output!(outputter::Output, # contains everything for netcdf_file output - time_sec::Int, # model time [s] for output - progn::PrognosticVariables, # all prognostic variables - diagn::DiagnosticVariables, # all diagnostic variables - M::ModelSetup) # all parameters +run_id_to_string(run_id::Integer) = @sprintf("%04d",run_id) +run_id_to_string(run_id::String) = run_id -Writes the variables from `diagn` of time step `i` at time `time_sec` into `netcdf_file`. Simply escapes for no -netcdf output of if output shouldn't be written on this time step. Converts variables from `diagn` to float32 -for output, truncates the mantissa for higher compression and applies lossless compression.""" -function write_netcdf_output!( outputter::Output, # everything for netcdf output - time::DateTime, # model time for output - diagn::DiagnosticVariables, # all diagnostic variables - model::ModelSetup) # all parameters + +""" +$(TYPEDSIGNATURES) +Writes the variables from `diagn` of time step `i` at time `time` into `outputter.netcdf_file`. +Simply escapes for no netcdf output of if output shouldn't be written on this time step. +Interpolates onto output grid and resolution as specified in `outputter`, converts to output +number format, truncates the mantissa for higher compression and applies lossless compression.""" +function write_output!( outputter::OutputWriter, # everything for netcdf output + time::DateTime, # model time for output + diagn::DiagnosticVariables) # all diagnostic variables outputter.timestep_counter += 1 # increase counter (; output, output_every_n_steps, timestep_counter ) = outputter @@ -243,15 +300,18 @@ function write_netcdf_output!( outputter::Output, # everything for timestep_counter % output_every_n_steps == 0 || return nothing # escape if output not written on this step # WRITE VARIABLES - write_netcdf_variables!(outputter,diagn,model) + write_netcdf_variables!(outputter,diagn) write_netcdf_time!(outputter,time) end -function write_netcdf_time!(outputter::Output, +""" +$(TYPEDSIGNATURES) +Write the current time `time::DateTime` to the netCDF file in `output::OutputWriter`.""" +function write_netcdf_time!(output::OutputWriter, time::DateTime) - (; netcdf_file, startdate ) = outputter - i = outputter.output_counter + (; netcdf_file, startdate ) = output + i = output.output_counter time_sec = [round(Int64,Dates.value(Dates.Second(time-startdate)))] NetCDF.putvar(netcdf_file,"time",time_sec,start=[i]) # write time [sec] of next output step @@ -260,27 +320,25 @@ function write_netcdf_time!(outputter::Output, return nothing end -function write_netcdf_variables!( outputter::Output, - diagn::DiagnosticVariables, - model::ModelSetup) - - outputter.output_counter += 1 # increase output step counter - (; output_vars ) = outputter # Vector{Symbol} of variables to output - i = outputter.output_counter +""" +$(TYPEDSIGNATURES) +Write diagnostic variables from `diagn` to the netCDF file in `output::OutputWriter`.""" +function write_netcdf_variables!( output::OutputWriter, + diagn::DiagnosticVariables{NF,Grid,Model}) where {NF,Grid,Model} - (; u, v, vor, div, pres, temp, humid ) = outputter - (; output_matrix, output_Grid ) = outputter - (; interpolator ) = outputter + output.output_counter += 1 # increase output step counter + (;output_vars) = output # Vector{Symbol} of variables to output + i = output.output_counter - # output to matrix options - quadrant_rotation = model.parameters.output_quadrant_rotation - matrix_quadrant = model.parameters.output_matrix_quadrant + (;u, v, vor, div, pres, temp, humid) = output + (;output_Grid, interpolator) = output + (;quadrant_rotation, matrix_quadrant) = output for (k,diagn_layer) in enumerate(diagn.layers) (; u_grid, v_grid, vor_grid, div_grid, temp_grid, humid_grid ) = diagn_layer.grid_variables - - if output_matrix # resort gridded variables interpolation-free into a matrix + + if output.as_matrix # resort gridded variables interpolation-free into a matrix # create (matrix,grid) tuples for simultaneous grid -> matrix conversion MGs = ((M,G) for (M,G) in zip((u,v,vor,div,temp,humid), @@ -289,7 +347,7 @@ function write_netcdf_variables!( outputter::Output, RingGrids.Matrix!(MGs...; quadrant_rotation, matrix_quadrant) - else # or interpolate onto a full grid + else # or interpolate onto a full grid :u in output_vars && RingGrids.interpolate!(output_Grid(u), u_grid, interpolator) :v in output_vars && RingGrids.interpolate!(output_Grid(v), v_grid, interpolator) :vor in output_vars && RingGrids.interpolate!(output_Grid(vor), vor_grid, interpolator) @@ -299,12 +357,12 @@ function write_netcdf_variables!( outputter::Output, end # UNSCALE THE SCALED VARIABLES - unscale!(vor,model) # was vor*radius, back to vor - unscale!(div,model) # same - temp .-= 273.15 # convert to ˚C + unscale!(vor,diagn.scale[]) # was vor*radius, back to vor + unscale!(div,diagn.scale[]) # same + temp .-= 273.15 # convert to ˚C # ROUNDING FOR ROUND+LOSSLESS COMPRESSION - (; keepbits ) = model.parameters + (; keepbits ) = output :u in output_vars && round!(u, keepbits.u) :v in output_vars && round!(v, keepbits.v) :vor in output_vars && round!(vor, keepbits.vor) @@ -313,7 +371,7 @@ function write_netcdf_variables!( outputter::Output, :humid in output_vars && round!(humid,keepbits.humid) # WRITE VARIABLES TO FILE, APPEND IN TIME DIMENSION - (; netcdf_file ) = outputter + (; netcdf_file ) = output :u in output_vars && NetCDF.putvar(netcdf_file,"u", u, start=[1,1,k,i],count=[-1,-1,1,1]) :v in output_vars && NetCDF.putvar(netcdf_file,"v", v, start=[1,1,k,i],count=[-1,-1,1,1]) :vor in output_vars && NetCDF.putvar(netcdf_file,"vor", vor, start=[1,1,k,i],count=[-1,-1,1,1]) @@ -326,45 +384,40 @@ function write_netcdf_variables!( outputter::Output, if :pres in output_vars (; pres_grid ) = diagn.surface - if output_matrix + if output.as_matrix RingGrids.Matrix!(pres,diagn.surface.pres_grid; quadrant_rotation, matrix_quadrant) else RingGrids.interpolate!(output_Grid(pres),pres_grid,interpolator) end - if model isa PrimitiveEquation + if Model <: PrimitiveEquation @. pres = exp(pres)/100 # convert from log(pₛ) to surface pressure pₛ [hPa] end - round!(pres,model.parameters.keepbits.pres) + + round!(pres,output.keepbits.pres) - NetCDF.putvar(outputter.netcdf_file,"pres",pres,start=[1,1,i],count=[-1,-1,1]) + NetCDF.putvar(output.netcdf_file,"pres",pres,start=[1,1,i],count=[-1,-1,1]) end return nothing end """ - write_restart_file( time::DateTime, - progn::PrognosticVariables, - outputter::Output, - M::ModelSetup) - +$(TYPEDSIGNATURES) A restart file `restart.jld2` with the prognostic variables is written to the output folder (or current path) that can be used to restart the model. `restart.jld2` will then be used as initial conditions. The prognostic variables -are bitround for compression and the 2nd leapfrog time step is discarded. -While the dynamical core may work with scaled variables, the restart file -contains these variables unscaled.""" +are bitrounded for compression and the 2nd leapfrog time step is discarded. +Variables in restart file are unscaled.""" function write_restart_file(time::DateTime, progn::PrognosticVariables, - outputter::Output, - model::ModelSetup) + output::OutputWriter) - (; run_path, write_restart ) = outputter + (; run_path, write_restart, keepbits ) = output + output.output || return nothing # exit immediately if no output and write_restart || return nothing # exit immediately if no restart file desired # COMPRESSION OF RESTART FILE - (; keepbits ) = model.parameters for layer in progn.layers # copy over leapfrog 2 to 1 @@ -394,24 +447,22 @@ function write_restart_file(time::DateTime, jldopen(joinpath(run_path,"restart.jld2"),"w"; compress=true) do f f["prognostic_variables"] = progn f["time"] = time - f["version"] = model.parameters.version + f["version"] = output.pkg_version f["description"] = "Restart file created for SpeedyWeather.jl" end end """ - get_full_output_file_path(p::Parameters) - +$(TYPEDSIGNATURES) Returns the full path of the output file after it was created. """ -get_full_output_file_path(P::Parameters) = joinpath(P.output_path, string("run-",run_id_string(P.run_id),"/"), P.output_filename) +get_full_output_file_path(output::OutputWriter) = joinpath(output.run_path, output.filename) """ - load_trajectory(var_name::Union{Symbol, String}, M::ModelSetup) - +$(TYPEDSIGNATURES) Loads a `var_name` trajectory of the model `M` that has been saved in a netCDF file during the time stepping. """ -function load_trajectory(var_name::Union{Symbol, String}, M::ModelSetup) - @assert M.parameters.output "Output is turned off" - return NetCDF.ncread(get_full_output_file_path(M.parameters), string(var_name)) +function load_trajectory(var_name::Union{Symbol, String}, model::ModelSetup) + @assert model.output.output "Output is turned off" + return NetCDF.ncread(get_full_output_file_path(model.output), string(var_name)) end diff --git a/src/output/pretty_printing.jl b/src/output/pretty_printing.jl deleted file mode 100644 index 21f770e7d..000000000 --- a/src/output/pretty_printing.jl +++ /dev/null @@ -1,75 +0,0 @@ -function Base.show(io::IO, P::PrognosticVariables) - - ζ = P.layers[end].timesteps[1].vor # create a view on vorticity - ζ_grid = Matrix(gridded(ζ)) # to grid space - ζ_grid = ζ_grid[:,end:-1:1] # flip latitudes - - nlon,nlat = size(ζ_grid) - - plot_kwargs = pairs(( xlabel="˚E", - xfact=360/(nlon-1), - ylabel="˚N", - yfact=180/(nlat-1), - yoffset=-90, - title="Surface relative vorticity", - colormap=:viridis, - compact=true, - colorbar=true, - width=60, - height=30)) - - print(io,UnicodePlots.heatmap(ζ_grid';plot_kwargs...)) -end - -# adapted from ProgressMeter.jl -function speedstring(sec_per_iter,dt_in_sec) - if sec_per_iter == Inf - return " N/A days/day" - end - - sim_time_per_time = dt_in_sec/sec_per_iter - - for (divideby, unit) in ( (365*1_000, "millenia"), - (365, "years"), - (1, "days"), - (1/24, "hours")) - if (sim_time_per_time / divideby) > 2 - return @sprintf "%5.2f %2s/day" (sim_time_per_time / divideby) unit - end - end - return " <2 hours/days" -end - -# hack: define global constant whose element will be changed in initialize_feedback -# used to pass on the time step to ProgressMeter.speedstring via calling this -# constant from the ProgressMeter module -const DT_IN_SEC = Ref(1800) - -function ProgressMeter.speedstring(sec_per_iter,dt_in_sec=SpeedyWeather.DT_IN_SEC) - speedstring(sec_per_iter,dt_in_sec[]) -end - -""" - readable_secs(secs::Real) -> Dates.CompoundPeriod - -Returns `Dates.CompoundPeriod` rounding to either (days, hours), (hours, minutes), (minutes, -seconds), or seconds with 1 decimal place accuracy for >10s and two for less. -E.g. -```julia -julia> readable_secs(12345) -3 hours, 26 minutes -``` -""" -function readable_secs(secs::Real) - millisecs = Dates.Millisecond(round(secs * 10 ^ 3)) - if millisecs >= Dates.Day(1) - return Dates.canonicalize(round(millisecs, Dates.Hour)) - elseif millisecs >= Dates.Hour(1) - return Dates.canonicalize(round(millisecs, Dates.Minute)) - elseif millisecs >= Dates.Minute(1) - return Dates.canonicalize(round(millisecs, Dates.Second)) - elseif millisecs >= Dates.Second(10) - return Dates.canonicalize(round(millisecs, Dates.Millisecond(100))) - end - return Dates.canonicalize(round(millisecs, Dates.Millisecond(10))) -end \ No newline at end of file diff --git a/src/physics/boundary_layer.jl b/src/physics/boundary_layer.jl index fe6294269..900775dff 100644 --- a/src/physics/boundary_layer.jl +++ b/src/physics/boundary_layer.jl @@ -1,58 +1,66 @@ -"""Concrete type that disables the boundary layer scheme.""" -struct NoBoundaryLayer{NF} <: BoundaryLayer{NF} end - -"""NoBoundaryLayer scheme just passes.""" -function boundary_layer!( column::ColumnVariables, - scheme::NoBoundaryLayer, - model::PrimitiveEquation) - return nothing -end +"""Concrete type that disables the boundary layer drag scheme.""" +struct NoBoundaryLayerDrag{NF} <: BoundaryLayerDrag{NF} end """NoBoundaryLayer scheme does not need any initialisation.""" -function initialize_boundary_layer!(K::ParameterizationConstants, - scheme::NoBoundaryLayer, - P::Parameters, - G::Geometry) +function initialize!( scheme::NoBoundaryLayerDrag, + model::PrimitiveEquation) return nothing end -"""Following Held and Suarez, 1996 BAMS""" -Base.@kwdef struct LinearDrag{NF<:Real} <: BoundaryLayer{NF} - σb::NF = 0.7 # sigma coordinate below which linear drag is applied - drag_time::NF = 24.0 # [hours] time scale for linear drag coefficient at σ=1 (=1/kf in HS96) +"""NoBoundaryLayer scheme just passes.""" +function boundary_layer_drag!( column::ColumnVariables, + scheme::NoBoundaryLayerDrag, + model::PrimitiveEquation) + return nothing end -# generator so that LinearDrag(drag_time=1::Int) is still possible → Float64 -LinearDrag(;kwargs...) = LinearDrag{Float64}(;kwargs...) +"""Linear boundary layer drag Following Held and Suarez, 1996 BAMS +$(TYPEDFIELDS)""" +@with_kw struct LinearDrag{NF<:AbstractFloat} <: BoundaryLayerDrag{NF} + # PARAMETERS + σb::Float64 = 0.7 # sigma coordinate below which linear drag is applied + time_scale::Float64 = 24 # [hours] time scale for linear drag coefficient at σ=1 (=1/kf in HS96) -function boundary_layer!( column::ColumnVariables, - scheme::LinearDrag, - model::PrimitiveEquation) - (;u,v,u_tend,v_tend) = column - (;drag_coefs) = model.parameterization_constants - - @inbounds for k in eachlayer(column) - kᵥ = drag_coefs[k] - if kᵥ > 0 - u_tend[k] -= kᵥ*u[k] # Held and Suarez 1996, equation 1 - v_tend[k] -= kᵥ*v[k] - end - end + # PRECOMPUTED CONSTANTS + nlev::Int = 0 + drag_coefs::Vector{NF} = zeros(NF,nlev) end -function initialize_boundary_layer!(K::ParameterizationConstants, - scheme::LinearDrag, - P::Parameters, - G::Geometry) +""" +$(TYPEDSIGNATURES) +Generator function using `nlev` from `SG::SpectralGrid`""" +LinearDrag(SG::SpectralGrid;kwargs...) = LinearDrag{SG.NF}(nlev=SG.nlev;kwargs...) + +""" +$(TYPEDSIGNATURES) +Precomputes the drag coefficients for this `BoundaryLayerDrag` scheme.""" +function initialize!( scheme::LinearDrag, + model::PrimitiveEquation) - (;σ_levels_full,radius) = G - (;σb,drag_time) = scheme - (;drag_coefs) = K + (;σ_levels_full,radius) = model.geometry + (;σb,time_scale,drag_coefs) = scheme - kf = radius/(drag_time*3600) # scale with radius as ∂ₜu is; hrs -> sec + kf = radius/(time_scale*3600) # scale with radius as ∂ₜu is; hrs -> sec for (k,σ) in enumerate(σ_levels_full) - drag_coefs[k] = kf*max(0,(σ-σb)/(1-σb)) # drag only below σb, lin increasing to kf at σ=1 + drag_coefs[k] = -kf*max(0,(σ-σb)/(1-σb)) # drag only below σb, lin increasing to kf at σ=1 end end - \ No newline at end of file + +""" +$(TYPEDSIGNATURES) +Compute tendency for boundary layer drag of a `column` and add to its tendencies fields""" +function boundary_layer_drag!( column::ColumnVariables, + scheme::LinearDrag) + + (;u,v,u_tend,v_tend) = column + (;drag_coefs) = scheme + + @inbounds for k in eachlayer(column) + kᵥ = drag_coefs[k] + if kᵥ > 0 + u_tend[k] += kᵥ*u[k] # Held and Suarez 1996, equation 1 + v_tend[k] += kᵥ*v[k] + end + end +end \ No newline at end of file diff --git a/src/physics/column_variables.jl b/src/physics/column_variables.jl index 5bc1d9cbe..978841296 100644 --- a/src/physics/column_variables.jl +++ b/src/physics/column_variables.jl @@ -107,4 +107,4 @@ function reset_column!(column::ColumnVariables) end # iterator for convenience -eachlayer(column::ColumnVariables) = Base.OneTo(column.nlev) +eachlayer(column::ColumnVariables) = Base.OneTo(column.nlev) \ No newline at end of file diff --git a/src/physics/convection.jl b/src/physics/convection.jl index 45e2e4dc3..13af21ef7 100644 --- a/src/physics/convection.jl +++ b/src/physics/convection.jl @@ -1,9 +1,5 @@ """ - diagnose_convection!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, - ) - +$(TYPEDSIGNATURES) Check whether the convection scheme should be activated in the given atmospheric column. 1. A conditional instability exists when the saturation moist energy (MSS) decreases with @@ -28,16 +24,13 @@ boundary of the full level k. The top-of-convection (TCN) layer, or cloud-top, is the largest value of k for which condition 1 is satisfied. The cloud-top layer may be subsequently adjusted upwards by the -large-scale condensation parameterization, which is executed after this one. -""" -function diagnose_convection!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, -) where {NF<:AbstractFloat} +large-scale condensation parameterization, which is executed after this one.""" +function diagnose_convection!(column::ColumnVariables,convection::SpeedyConvection) + (; alhc,pres_ref ) = model.parameters (; pres_thresh_cnv, RH_thresh_pbl_cnv ) = model.constants (; nlev ) = column - (; humid, pres, sat_humid, dry_static_energy, moist_static_energy, + (; humid, pres, sat_humid, moist_static_energy, sat_moist_static_energy, sat_moist_static_energy_half) = column if pres[end] > pres_thresh_cnv @@ -89,6 +82,9 @@ function diagnose_convection!( return nothing end +convection!(column::ColumnVariables,model::PrimitiveDry) = nothing +convection!(column::ColumnVariables,model::PrimitiveWet) = convection!(column,model.convection) + """ convection!( column::ColumnVariables{NF}, @@ -195,3 +191,69 @@ function convection!( return nothing end + + + # # Compute the entrainment coefficients for the convection parameterization. + # (;max_entrainment) = P + # entrainment_profile = zeros(nlev) + # for k = 2:nlev-1 + # entrainment_profile[k] = max(0, (σ_levels_full[k] - 0.5)^2) + # end + + # # profile as fraction of cloud-base mass flux + # entrainment_profile /= sum(entrainment_profile) # Normalise + # entrainment_profile *= max_entrainment # fraction of max entrainment + + # # PARAMETRIZATIONS + # # Large-scale condensation (occurs when relative humidity exceeds a given threshold) + # RH_thresh_pbl_lsc::NF # Relative humidity threshold for LSC in PBL + # RH_thresh_range_lsc::NF # Vertical range of relative humidity threshold + # RH_thresh_max_lsc ::NF # Maximum relative humidity threshold + # humid_relax_time_lsc::NF # Relaxation time for humidity (hours) + + # # Convection + # pres_thresh_cnv::NF # Minimum (normalised) surface pressure for the occurrence of convection + # RH_thresh_pbl_cnv::NF # Relative humidity threshold for convection in PBL + # RH_thresh_trop_cnv::NF # Relative humidity threshold for convection in the troposphere + # humid_relax_time_cnv::NF # Relaxation time for PBL humidity (hours) + # max_entrainment::NF # Maximum entrainment as a fraction of cloud-base mass flux + # ratio_secondary_mass_flux::NF # Ratio between secondary and primary mass flux at cloud-base + + + # "For computing saturation vapour pressure" + # magnus_coefs::Coefficients = MagnusCoefs{NF}() + + # # Large-Scale Condensation (from table B10) + # "Index of atmospheric level at which large-scale condensation begins" + # k_lsc::Int = 2 + + # "Relative humidity threshold for boundary layer" + # RH_thresh_pbl_lsc::Float64 = 0.95 + + # "Vertical range of relative humidity threshold" + # RH_thresh_range_lsc::Float64 = 0.1 + + # "Maximum relative humidity threshold" + # RH_thresh_max_lsc::Float64 = 0.9 + + # "Relaxation time for humidity (hours)" + # humid_relax_time_lsc::Float64 = 4.0 + + # # Convection + # "Minimum (normalised) surface pressure for the occurrence of convection" + # pres_thresh_cnv::Float64 = 0.8 + + # "Relative humidity threshold for convection in PBL" + # RH_thresh_pbl_cnv::Float64 = 0.9 + + # "Relative humidity threshold for convection in the troposphere" + # RH_thresh_trop_cnv::Float64 = 0.7 + + # "Relaxation time for PBL humidity (hours)" + # humid_relax_time_cnv::Float64 = 6.0 + + # "Maximum entrainment as a fraction of cloud-base mass flux" + # max_entrainment::Float64 = 0.5 + + # "Ratio between secondary and primary mass flux at cloud-base" + # ratio_secondary_mass_flux::Float64 = 0.8 diff --git a/src/physics/define_column.jl b/src/physics/define_column.jl index ec96434c8..548f101c2 100644 --- a/src/physics/define_column.jl +++ b/src/physics/define_column.jl @@ -1,10 +1,9 @@ """ - column = ColumnVariables{NF<:AbstractFloat} - Mutable struct that contains all prognostic (copies thereof) and diagnostic variables in a single column needed to evaluate the physical parametrizations. For now the struct is mutable as we will reuse the struct to iterate over horizontal grid points. Every column vector has `nlev` entries, from [1] at the top to -[end] at the lowermost model level at the planetary boundary layer.""" +[end] at the lowermost model level at the planetary boundary layer. +$(TYPEDFIELDS)""" Base.@kwdef mutable struct ColumnVariables{NF<:AbstractFloat} <: AbstractColumnVariables{NF} # DIMENSIONS @@ -119,7 +118,4 @@ Base.@kwdef mutable struct ColumnVariables{NF<:AbstractFloat} <: AbstractColumnV # Shortwave radiation: shortwave_radiation rel_hum::Vector{NF} = fill(NF(NaN), nlev) # Relative humidity grad_dry_static_energy::NF = NF(NaN) # gradient of dry static energy -end - -# use Float64 if not provided -ColumnVariables(;kwargs...) = ColumnVariables{Float64}(;kwargs...) \ No newline at end of file +end \ No newline at end of file diff --git a/src/physics/pretty_printing.jl b/src/physics/pretty_printing.jl new file mode 100644 index 000000000..0f8f59005 --- /dev/null +++ b/src/physics/pretty_printing.jl @@ -0,0 +1,8 @@ +# print all fields with type <: Number +function Base.show(io::IO,P::AbstractParameterization) + print(io,"$(typeof(P)):") + for key in propertynames(P) + val = getfield(P,key) + val isa Number && print(io,"\n $key::$(typeof(val)) = $val") + end +end \ No newline at end of file diff --git a/src/physics/shortwave_radiation.jl b/src/physics/shortwave_radiation.jl index ed1c19f6e..8ccc3e5e8 100644 --- a/src/physics/shortwave_radiation.jl +++ b/src/physics/shortwave_radiation.jl @@ -1,7 +1,7 @@ """ Parameters for radiation parameterizations. """ -Base.@kwdef struct RadiationCoefs{NF<:Real} <: Coefficients +@with_kw struct RadiationCoefs{NF<:Real} <: Coefficients epslw::NF = 0.05 # Fraction of blackbody spectrum absorbed/emitted by PBL only emisfc::NF = 0.98 # Longwave surface emissivity diff --git a/src/physics/temperature_relaxation.jl b/src/physics/temperature_relaxation.jl index f4af8ea25..d2e8cda20 100644 --- a/src/physics/temperature_relaxation.jl +++ b/src/physics/temperature_relaxation.jl @@ -1,66 +1,82 @@ struct NoTemperatureRelaxation{NF} <: TemperatureRelaxation{NF} end +NoTemperatureRelaxation(SG::SpectralGrid) = NoTemperatureRelaxation{SG.NF}() -"""NoBoundaryLayer scheme just passes.""" +"""$(TYPEDSIGNATURES) just passes.""" function temperature_relaxation!( column::ColumnVariables, - scheme::NoTemperatureRelaxation, - model::PrimitiveEquation) + scheme::NoTemperatureRelaxation) return nothing end -"""NoBoundaryLayer scheme does not need any initialisation.""" -function initialize_temperature_relaxation!(K::ParameterizationConstants, - scheme::NoTemperatureRelaxation, - P::Parameters, - G::Geometry) +"""$(TYPEDSIGNATURES) just passes, does not need any initialization.""" +function initialize!( scheme::NoTemperatureRelaxation, + model::PrimitiveEquation) return nothing end -Base.@kwdef struct HeldSuarez{NF<:Real} <: TemperatureRelaxation{NF} - σb::NF = 0.7 # sigma coordinate below which linear drag is applied - - relax_time_slow::NF = 40.0*24 # [hours] time scale for slow global relaxation - relax_time_fast::NF = 4.0*24 # [hours] time scale for faster tropical surface relaxation - - Tmin::NF = 200.0 # minimum temperature [K] in equilibrium temperature - Tmax::NF = 315.0 # maximum temperature [K] in equilibrium temperature - - ΔTy::NF = 60.0 # meridional temperature gradient [K] - Δθz::NF = 10.0 # vertical temperature gradient [K] +""" +Struct that defines the temperature relaxation from Held and Suarez, 1996 BAMS +$(TYPEDFIELDS)""" +@with_kw struct HeldSuarez{NF<:AbstractFloat} <: TemperatureRelaxation{NF} + # DIMENSIONS + "number of latitude rings" + nlat::Int + + "number of vertical levels" + nlev::Int + + # OPTIONS + "sigma coordinate below which faster surface relaxation is applied" + σb::Float64 = 0.7 + + "time scale [hrs] for slow global relaxation" + relax_time_slow::Float64 = 40*24 + + "time scale [hrs] for faster tropical surface relaxation" + relax_time_fast::Float64 = 4*24 + + "minimum equilibrium temperature [K]" + Tmin::Float64 = 200 + + "maximum equilibrium temperature [K]" + Tmax::Float64 = 315 + + "meridional temperature gradient [K]" + ΔTy::Float64 = 60 + + "vertical temperature gradient [K]" + Δθz::Float64 = 10 + + # precomputed constants, allocate here, fill in initialize! + κ::Base.RefValue{NF} = Ref(zero(NF)) + p₀::Base.RefValue{NF} = Ref(zero(NF)) + + temp_relax_freq::Matrix{NF} = zeros(NF,nlev,nlat) # (inverse) relax time scale per layer and lat + temp_equil_a::Vector{NF} = zeros(NF,nlat) # terms to calc equilibrium temper func + temp_equil_b::Vector{NF} = zeros(NF,nlat) # of latitude and pressure end -# generator so that HeldSuarez(Tmin=200::Int) is still possible → Float64 -HeldSuarez(;kwargs...) = HeldSuarez{Float64}(;kwargs...) - -function temperature_relaxation!( column::ColumnVariables{NF}, - scheme::HeldSuarez, - model::PrimitiveEquation) where NF - - (;temp,temp_tend,pres,ln_pres) = column - j = column.jring[] # latitude ring index j - (;temp_relax_freq,temp_equil_a,temp_equil_b) = model.parameterization_constants - Tmin = convert(NF,scheme.Tmin) - p₀ = convert(NF,model.parameters.pres_ref*100) # [hPa] → [Pa] - (;κ) = model.constants # R/cₚ = 2/7 - - @inbounds for k in eachlayer(column) - lnp = ln_pres[k] # logarithm of pressure at level k - kₜ = temp_relax_freq[k,j] # (inverse) relaxation time scale - - # Held and Suarez 1996, equation 3 with precomputed a,b during initilisation - Teq = max(Tmin,(temp_equil_a[j] + temp_equil_b[j]*lnp)*(pres[k]/p₀)^κ) - temp_tend[k] -= kₜ*(temp[k] - Teq) # Held and Suarez 1996, equation 2 - end +""" +$(TYPEDSIGNATURES) +create a HeldSuarez temperature relaxation with arrays allocated given `spectral_grid`""" +function HeldSuarez(SG::SpectralGrid;kwargs...) + (;NF, Grid, nlat_half, nlev) = SG + nlat = RingGrids.get_nlat(Grid,nlat_half) + return HeldSuarez{NF}(;nlev,nlat,kwargs...) end -function initialize_temperature_relaxation!(K::ParameterizationConstants, - scheme::HeldSuarez, - P::Parameters, - G::Geometry) +"""$(TYPEDSIGNATURES) +initialize the HeldSuarez temperature relaxation by precomputing terms for the +equilibrium temperature Teq.""" +function initialize!( scheme::HeldSuarez, + model::PrimitiveEquation) - (;σ_levels_full,radius,coslat,sinlat) = G - (;σb,ΔTy,Δθz,relax_time_slow,relax_time_fast,Tmax) = scheme - (;temp_relax_freq,temp_equil_a,temp_equil_b) = K - p₀ = P.pres_ref*100 # [hPa] → [Pa] + (;σ_levels_full, radius, coslat, sinlat) = model.geometry + (;σb, ΔTy, Δθz, relax_time_slow, relax_time_fast, Tmax) = scheme + (;temp_relax_freq, temp_equil_a, temp_equil_b) = scheme + + p₀ = model.atmosphere.pres_ref*100 # [hPa] → [Pa] + scheme.p₀[] = p₀ + scheme.κ[] = model.constants.κ # thermodynamic kappa # slow relaxation everywhere, fast in the tropics kₐ = radius/(relax_time_slow*3600) # scale with radius as ∂ₜT is; hrs -> sec @@ -79,52 +95,83 @@ function initialize_temperature_relaxation!(K::ParameterizationConstants, end end -""" - J = JablonowskiRelaxation{NF}() +"""$(TYPEDSIGNATURES) +Apply temperature relaxation following Held and Suarez 1996, BAMS.""" +function temperature_relaxation!( column::ColumnVariables{NF}, + scheme::HeldSuarez) where NF + + (;temp, temp_tend, pres, ln_pres) = column + j = column.jring[] # latitude ring index j + + (;temp_relax_freq, temp_equil_a, temp_equil_b) = scheme + Tmin = convert(NF,scheme.Tmin) + + p₀ = scheme.p₀[] # reference surface pressure + κ = scheme.κ[] # thermodynamic kappa + + @inbounds for k in eachlayer(column) + lnp = ln_pres[k] # logarithm of pressure at level k + kₜ = temp_relax_freq[k,j] # (inverse) relaxation time scale + + # Held and Suarez 1996, equation 3 with precomputed a,b during initilisation + Teq = max(Tmin,(temp_equil_a[j] + temp_equil_b[j]*lnp)*(pres[k]/p₀)^κ) + temp_tend[k] -= kₜ*(temp[k] - Teq) # Held and Suarez 1996, equation 2 + end +end + +"""$(TYPEDSIGNATURES) HeldSuarez-like temperature relaxation, but towards the Jablonowski temperature profile with increasing temperatures in the stratosphere.""" -Base.@kwdef struct JablonowskiRelaxation{NF<:Real} <: TemperatureRelaxation{NF} - σb::NF = 0.7 # sigma coordinate below which relax_time_fast is applied +@with_kw struct JablonowskiRelaxation{NF<:AbstractFloat} <: TemperatureRelaxation{NF} + # DIMENSIONS + nlat::Int + nlev::Int - η₀::NF = 0.252 # conversion from σ to Jablonowski's ηᵥ-coordinates - u₀::NF = 35.0 # max amplitude of zonal wind [m/s] - ΔT::NF = 4.8e5 # temperature difference used for stratospheric lapse rate [K] + # OPTIONS + "sigma coordinate below which relax_time_fast is applied" + σb::Float64= 0.7 - relax_time_slow::NF = 40.0*24 # [hours] time scale for slow global relaxation - relax_time_fast::NF = 4.0*24 # [hours] time scale for faster tropical surface relaxation -end + "conversion from σ to Jablonowski's ηᵥ-coordinates" + η₀::Float64 = 0.252 -# generator so that arguments are converted to Float64 -JablonowskiRelaxation(;kwargs...) = JablonowskiRelaxation{Float64}(;kwargs...) + "max amplitude of zonal wind [m/s]" + u₀::Float64 = 35 -function temperature_relaxation!( column::ColumnVariables{NF}, - scheme::JablonowskiRelaxation, - model::PrimitiveEquation) where NF + "temperature difference used for stratospheric lapse rate [K]" + ΔT::Float64 = 4.8e5 - (;temp,temp_tend) = column - j = column.jring[] # latitude ring index j - (;temp_relax_freq,temp_equil) = model.parameterization_constants + "[hours] time scale for slow global relaxation" + relax_time_slow::NF = 40*24 + + "[hours] time scale for fast aster tropical surface relaxation" + relax_time_fast::NF = 4*24 - @inbounds for k in eachlayer(column) - kₜ = temp_relax_freq[k,j] # (inverse) relaxation time scale + # precomputed constants, allocate here, fill in initialize! + temp_relax_freq::Matrix{NF} = zeros(NF,nlev,nlat) # (inverse) relax time scale per layer and lat + temp_equil::Matrix{NF} = zeros(NF,nlev,nlat) # terms to calc equilibrium temperature as func +end - # Held and Suarez 1996, equation 2, but using temp_equil from - # Jablonowski and Williamson 2006, equation 6 - temp_tend[k] -= kₜ*(temp[k] - temp_equil[k,j]) - end +""" +$(TYPEDSIGNATURES) +create a JablonowskiRelaxation temperature relaxation with arrays allocated given `spectral_grid`""" +function JablonowskiRelaxation(SG::SpectralGrid;kwargs...) + (;NF, Grid, nlat_half, nlev) = SG + nlat = RingGrids.get_nlat(Grid,nlat_half) + return JablonowskiRelaxation{NF}(;nlev,nlat,kwargs...) end -function initialize_temperature_relaxation!(K::ParameterizationConstants, - scheme::JablonowskiRelaxation, - P::Parameters, - G::Geometry) +"""$(TYPEDSIGNATURES) +initialize the JablonowskiRelaxation temperature relaxation by precomputing terms for the +equilibrium temperature Teq and the frequency (strength of relaxation).""" +function initialize!( scheme::JablonowskiRelaxation, + model::PrimitiveEquation) - (;σ_levels_full,radius,coslat,sinlat) = G - (;σb,relax_time_slow,relax_time_fast,η₀,u₀,ΔT) = scheme - (;temp_relax_freq,temp_equil) = K - (;gravity,radius,rotation) = P.planet - (;lapse_rate,R_dry,σ_tropopause,temp_ref) = P + (;σ_levels_full, radius, coslat, sinlat) = model.geometry + (;σb, relax_time_slow, relax_time_fast, η₀, u₀, ΔT) = scheme + (;temp_relax_freq, temp_equil) = scheme + (;gravity, rotation) = model.planet + (;lapse_rate, R_dry, σ_tropopause, temp_ref) = model.atmosphere Γ = lapse_rate/1000 # from [K/km] to [K/m] aΩ = radius*rotation @@ -153,7 +200,27 @@ function initialize_temperature_relaxation!(K::ParameterizationConstants, A2 = 2u₀*cos(ηᵥ)^(3/2) # Jablonowski and Williamson, eq. (6) - temp_equil[k,j] = Tη + A1*((-2sinϕ^6*(cosϕ^2 + 1/3) + 10/63)*A2 + (8/5*cosϕ^3*(sinϕ^2 + 2/3) - π/4)*aΩ) + temp_equil[k,j] = Tη + A1*((-2sinϕ^6*(cosϕ^2 + 1/3) + 10/63)*A2 + + (8/5*cosϕ^3*(sinϕ^2 + 2/3) - π/4)*aΩ) end end -end \ No newline at end of file +end + +"""$(TYPEDSIGNATURES) +Apply HeldSuarez-like temperature relaxation to the Jablonowski and Williamson +vertical profile.""" +function temperature_relaxation!( column::ColumnVariables, + scheme::JablonowskiRelaxation) + + (;temp, temp_tend) = column + j = column.jring[] # latitude ring index j + (;temp_relax_freq, temp_equil) = scheme + + @inbounds for k in eachlayer(column) + kₜ = temp_relax_freq[k,j] # (inverse) relaxation time scale + + # Held and Suarez 1996, equation 2, but using temp_equil from + # Jablonowski and Williamson 2006, equation 6 + temp_tend[k] -= kₜ*(temp[k] - temp_equil[k,j]) + end +end diff --git a/src/physics/tendencies.jl b/src/physics/tendencies.jl index a1319b718..4d55fd81d 100644 --- a/src/physics/tendencies.jl +++ b/src/physics/tendencies.jl @@ -1,24 +1,23 @@ """ - parameterization_tendencies!( diagn::DiagnosticVariables, - time::DateTime, - M::PrimitiveEquation) - +$(TYPEDSIGNATURES) Compute tendencies for u,v,temp,humid from physical parametrizations. Extract for each vertical atmospheric column the prognostic variables (stored in `diagn` as they are grid-point transformed), loop over all grid-points, compute all parametrizations on a single-column basis, then write the tendencies back into a horizontal field of tendencies. """ -function parameterization_tendencies!( diagn::DiagnosticVariables, - time::DateTime, - model::PrimitiveEquation) - +function parameterization_tendencies!( + diagn::DiagnosticVariables, + time::DateTime, + model::PrimitiveEquation, +) + + (;boundary_layer_drag) = model + (;temperature_relaxation) = model + # (;vertical_diffusion) = model + (;static_energy_diffusion) = model + G = model.geometry - boundary_layer_scheme = model.parameters.boundary_layer - temperature_relax_scheme = model.parameters.temperature_relaxation - vertical_diffusion_scheme = model.parameters.vertical_diffusion - static_energy_diffusion_scheme = model.parameters.static_energy_diffusion - rings = eachring(G.Grid,G.nlat_half) @floop for ij in eachgridpoint(diagn) # loop over all horizontal grid points @@ -34,12 +33,12 @@ function parameterization_tendencies!( diagn::DiagnosticVariables, get_thermodynamics!(column,model) # VERTICAL DIFFUSION - vertical_diffusion!(column,vertical_diffusion_scheme,model) - static_energy_diffusion!(column,static_energy_diffusion_scheme,model) + # vertical_diffusion!(column,vertical_diffusion,model) + static_energy_diffusion!(column,static_energy_diffusion) # HELD-SUAREZ - temperature_relaxation!(column,temperature_relax_scheme,model) - boundary_layer!(column,boundary_layer_scheme,model) + temperature_relaxation!(column,temperature_relaxation) + boundary_layer_drag!(column,boundary_layer_drag) # Calculate parametrizations (order of execution is important!) # convection!(column,model) @@ -51,27 +50,33 @@ function parameterization_tendencies!( diagn::DiagnosticVariables, # vertical_diffusion!(column,M) # sum fluxes on half levels up and down for every layer - fluxes_to_tendencies!(column,model) + fluxes_to_tendencies!(column,model.geometry,model.constants) # write tendencies from parametrizations back into horizontal fields write_column_tendencies!(diagn,column,ij) end end -function fluxes_to_tendencies!( column::ColumnVariables{NF}, - model::PrimitiveEquation) where NF +""" +$(TYPEDSIGNATURES) +Convert the fluxes on half levels to tendencies on full levels.""" +function fluxes_to_tendencies!( + column::ColumnVariables, + geometry::Geometry, + constants::DynamicsConstants, +) (;nlev,u_tend,flux_u_upward,flux_u_downward) = column (;v_tend,flux_v_upward,flux_v_downward) = column (;humid_tend,flux_humid_upward,flux_humid_downward) = column (;temp_tend,flux_temp_upward,flux_temp_downward) = column - Δσ = model.geometry.σ_levels_thick + Δσ = geometry.σ_levels_thick pₛ = column.pres[end] # surface pressure # # g/pₛ and g/(pₛ*cₚ), see Fortran SPEEDY documentation eq. (3,5) - g_pₛ = convert(NF,model.constants.gravity/pₛ) - g_pₛ_cₚ = g_pₛ/convert(NF,model.parameters.cₚ) + g_pₛ = constants.gravity/pₛ + g_pₛ_cₚ = g_pₛ/constants.cₚ # fluxes are defined on half levels including top k=1/2 and surface k=nlev+1/2 @inbounds for k in 1:nlev diff --git a/src/physics/thermodynamics.jl b/src/physics/thermodynamics.jl index 041f4898c..4ed63e2fc 100644 --- a/src/physics/thermodynamics.jl +++ b/src/physics/thermodynamics.jl @@ -5,34 +5,48 @@ Parameters for computing saturation vapour pressure using the August-Roche-Magnu where T is in Kelvin and i = 1,2 for saturation with respect to water and ice, respectively. -""" -Base.@kwdef struct MagnusCoefs{NF<:Real} <: Coefficients - e₀::NF = 6.108 # Saturation vapour pressure at 0°C - T₀::NF = 273.16 # 0°C in Kelvin +$(TYPEDFIELDS)""" +@with_kw struct MagnusCoefs{NF<:AbstractFloat} + "Saturation vapour pressure at 0°C" + e₀::NF = 6.108 + + "0°C in Kelvin" + T₀::NF = 273.16 T₁::NF = 35.86 T₂::NF = 7.66 C₁::NF = 17.269 C₂::NF = 21.875 end +@with_kw struct Thermodynamics{NF} <: AbstractThermodynamics{NF} + magnus_coefs::MagnusCoefs{NF} = MagnusCoefs{NF}() + mol_ratio::NF + latent_heat_condensation::NF + latent_heat_sublimation::NF +end + +function Thermodynamics(SG::SpectralGrid,atm::AbstractAtmosphere;kwargs...) + (;latent_heat_condensation, latent_heat_sublimation) = atm + mol_ratio = atm.mol_mass_vapour/atm.mol_mass_dry_air + return Thermodynamics{SG.NF}(;mol_ratio,latent_heat_condensation,latent_heat_sublimation,kwargs...) +end + """ - get_thermodynamics!(column::ColumnVariables,model::PrimitiveWetCore) +$(TYPEDSIGNATURES) +Calculate the dry static energy for the primitive dry model.""" +function get_thermodynamics!(column::ColumnVariables,model::PrimitiveDry) + dry_static_energy!(column, model.constants) +end +""" +$(TYPEDSIGNATURES) Calculate thermodynamic quantities like saturation vapour pressure, saturation specific humidity, dry static energy, moist static energy and saturation moist static energy from the prognostic column variables.""" -function get_thermodynamics!( column::ColumnVariables, - model::PrimitiveEquation) - - # Calculate thermodynamic quantities at full levels - dry_static_energy!(column, model) - - if model isa PrimitiveWetCore - saturation_vapour_pressure!(column, model) - saturation_specific_humidity!(column, model) - moist_static_energy!(column, model) - saturation_moist_static_energy!(column, model) - end +function get_thermodynamics!(column::ColumnVariables,model::PrimitiveWet) + dry_static_energy!(column, model.constants) + saturation_humidity!(column, model.thermodynamics) + moist_static_energy!(column, model.thermodynamics) # Interpolate certain variables to half-levels # interpolate!(column, model) @@ -41,189 +55,65 @@ function get_thermodynamics!( column::ColumnVariables, end """ - interpolate!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, - ) -""" -function interpolate!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, -) where {NF<:AbstractFloat} - (; humid, humid_half ) = column - (; sat_humid, sat_humid_half ) = column - (; dry_static_energy, dry_static_energy_half ) = column - (; sat_moist_static_energy, sat_moist_static_energy_half ) = column - - for (full, half) in ( - (humid, humid_half), - (sat_humid, sat_humid_half), - (dry_static_energy, dry_static_energy_half), - (sat_moist_static_energy, sat_moist_static_energy_half), - ) - interpolate!(full, half, column, model) - end +$(TYPEDSIGNATURES) +Compute the dry static energy SE = cₚT + Φ (latent heat times temperature plus geopotential) +for the column.""" +function dry_static_energy!(column::ColumnVariables,constants::DynamicsConstants) - return nothing -end + (;cₚ) = constants + (;dry_static_energy, geopot, temp) = column -""" - interpolate!( - A_full_level::Vector{NF}, - A_half_level::Vector{NF}, - column::ColumnVariables{NF}, - model::PrimitiveEquation, - ) - -Given some generic column variable A defined at full levels, do a linear interpolation in -log(σ) to calculate its values at half-levels. -""" -function interpolate!( - A_full_level::Vector{NF}, - A_half_level::Vector{NF}, - column::ColumnVariables{NF}, - model::PrimitiveEquation, -) where {NF<:AbstractFloat} - (; nlev ) = column - (; σ_levels_full, σ_levels_half ) = model.geometry - - # For A at each full level k, compute A at the half-level below, i.e. at the boundary - # between the full levels k and k+1. - for k = 1:nlev-1 - A_half_level[k] = - A_full_level[k] + - (A_full_level[k+1] - A_full_level[k]) * - (log(σ_levels_half[k+1]) - log(σ_levels_full[k])) / - (log(σ_levels_full[k+1]) - log(σ_levels_full[k])) + @inbounds for k in eachlayer(column) + dry_static_energy[k] = cₚ * temp[k] + geopot[k] end - # Compute the values at the surface separately - A_half_level[nlev] = - A_full_level[nlev] + - (A_full_level[nlev] - A_full_level[nlev-1]) * - (log(NF(0.99)) - log(σ_levels_full[nlev])) / - (log(σ_levels_full[nlev]) - log(σ_levels_full[nlev-1])) - return nothing end - -""" - saturation_vapour_pressure!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, - ) - -Compute the saturation vapour pressure as a function of temperature using the +"""$(TYPEDSIGNATURES) +Compute (1) the saturation vapour pressure as a function of temperature using the August-Roche-Magnus formula, -eᵢ(T) = e₀ * exp(Cᵢ * (T - T₀) / (T - Tᵢ)), + eᵢ(T) = e₀ * exp(Cᵢ * (T - T₀) / (T - Tᵢ)), where T is in Kelvin and i = 1,2 for saturation with respect to water and ice, -respectively. -""" -function saturation_vapour_pressure!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, -) where {NF<:AbstractFloat} +respectively. And (2) the saturation specific humidity according to the formula, + + 0.622 * e / (p - (1 - 0.622) * e), - (; sat_vap_pres, temp ) = column - (; e₀, T₀, C₁, C₂, T₁, T₂ ) = model.parameters.magnus_coefs +where `e` is the saturation vapour pressure, `p` is the pressure, and 0.622 is the ratio of +the molecular weight of water to dry air.""" +function saturation_humidity!( + column::ColumnVariables, + thermodynamics::Thermodynamics, +) + (;sat_humid, sat_vap_pres, pres, temp) = column + (;e₀, T₀, C₁, C₂, T₁, T₂) = thermodynamics.magnus_coefs + (;mol_ratio) = thermodynamics # = mol_mass_vapour/mol_mass_dry_air = 0.622 for k in eachlayer(column) # change coefficients for water (temp > T₀) or ice (else) C, T = temp[k] > T₀ ? (C₁, T₁) : (C₂, T₂) sat_vap_pres[k] = e₀ * exp(C * (temp[k] - T₀) / (temp[k] - T)) - end - - return nothing -end - -""" - saturation_specific_humidity!( column::ColumnVariables{NF}, - model::PrimitiveEquation ) where {NF<:AbstractFloat} - -Compute the saturation specific humidity according to the formula, - -0.622 * e / (p - (1 - 0.622) * e), - -where e is the saturation vapour pressure, p is the pressure, and 0.622 is the ratio of -the molecular weight of water to dry air. -""" -function saturation_specific_humidity!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, -) where {NF<:AbstractFloat} - - (; sat_humid, sat_vap_pres, pres ) = column - (; mol_mass_vapour, mol_mass_dry_air ) = model.parameters - - mol_ratio = convert(NF, mol_mass_vapour/mol_mass_dry_air) - - @inbounds for k in eachlayer(column) sat_humid[k] = mol_ratio*sat_vap_pres[k] / (pres[k] - (1-mol_ratio)*sat_vap_pres[k]) end - - return nothing end -""" - dry_static_energy!(column::ColumnVariables,model::PrimitiveEquation) +"""$(TYPEDSIGNATURES) +Compute the moist static energy -Compute the dry static energy SE = cₚT + Φ (latent heat times temperature plus geopotential) -for the column.""" -function dry_static_energy!(column::ColumnVariables{NF}, - model::PrimitiveEquation) where NF + MSE = SE + Lc*Q = cₚT + Φ + Lc*Q - cₚ = convert(NF,model.parameters.cₚ) - (;dry_static_energy, geopot, temp) = column - - @inbounds for k in eachlayer(column) - dry_static_energy[k] = cₚ * temp[k] + geopot[k] - end - - return nothing -end - -""" - moist_static_energy!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, - ) -""" -function moist_static_energy!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, -) where {NF<:AbstractFloat} - (; alhc ) = model.parameters - (; moist_static_energy, dry_static_energy, humid ) = column +with the static energy `SE`, the latent heat of condensation `Lc`, +the geopotential `Φ`. As well as the saturation moist static energy +which replaces Q with Q_sat""" +function moist_static_energy!(column::ColumnVariables,thermodynamics::Thermodynamics) + (;latent_heat_condensation) = thermodynamics + (;sat_moist_static_energy, moist_static_energy, dry_static_energy) = column + (;humid, sat_humid) = column for k in eachlayer(column) - moist_static_energy[k] = dry_static_energy[k] + alhc * humid[k] + moist_static_energy[k] = dry_static_energy[k] + latent_heat_condensation * humid[k] + sat_moist_static_energy[k] = dry_static_energy[k] + latent_heat_condensation * sat_humid[k] end - - return nothing -end - -""" - saturation_moist_static_energy!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, - ) -""" -function saturation_moist_static_energy!( - column::ColumnVariables{NF}, - model::PrimitiveEquation, -) where {NF<:AbstractFloat} - (; alhc ) = model.parameters - (; sat_moist_static_energy, - sat_moist_static_energy_half, - dry_static_energy, - sat_humid) = column - - for k in eachlayer(column) - sat_moist_static_energy[k] = dry_static_energy[k] + alhc * sat_humid[k] - end - - return nothing end diff --git a/src/physics/vertical_diffusion.jl b/src/physics/vertical_diffusion.jl index ec4d8e521..8d3470cd1 100644 --- a/src/physics/vertical_diffusion.jl +++ b/src/physics/vertical_diffusion.jl @@ -1,124 +1,56 @@ struct NoVerticalDiffusion{NF} <: VerticalDiffusion{NF} end -NoVerticalDiffusion() = NoVerticalDiffusion{DEFAULT_NF}() +NoVerticalDiffusion(SG::SpectralGrid) = NoVerticalDiffusion{SG.NF}() -function vertical_diffusion!( column::ColumnVariables, - scheme::NoVerticalDiffusion, - model::PrimitiveEquation) - return nothing -end - -function initialize_vertical_diffusion!(K::ParameterizationConstants, - scheme::NoVerticalDiffusion, - P::Parameters, - G::Geometry) +function initialize!( scheme::NoVerticalDiffusion, + model::PrimitiveEquation) return nothing end -Base.@kwdef struct VerticalLaplacian{NF<:Real} <: VerticalDiffusion{NF} - time_scale::NF = 10.0 # [hours] time scale to control the strength of vertical diffusion - height_scale::NF = 100.0 # [m] scales for Δσ so that time_scale is sensible - - resolution_scaling::NF = 1.0 # (inverse) scaling with resolution T - nlev_scaling::NF = -2.0 # (inverse) scaling with n vertical levels -end - -# generator so that arguments are converted to Float64 -VerticalLaplacian(;kwargs...) = VerticalLaplacian{Float64}(;kwargs...) - -function vertical_diffusion!( column::ColumnVariables{NF}, - scheme::VerticalLaplacian, - model::PrimitiveEquation) where NF - - (;nlev,u_tend,v_tend,temp_tend) = column - (;u,v,temp) = column - (;time_scale, height_scale, resolution_scaling, nlev_scaling) = scheme - (;trunc) = model.parameters - ∇²_below = model.parameterization_constants.vert_diff_∇²_below - ∇²_above = model.parameterization_constants.vert_diff_∇²_above - - # GET DIFFUSION COEFFICIENT as a function of u,v,temp and surface pressure - # *3600 for [hrs] → [s], *1e3 for [km] → [m] - # include a height scale, technically not needed, but so that the dimensionless - # 1/Δσ² gets a resonable scale in meters such that the time scale is not - # counterintuitively in seconds or years - ν0 = model.geometry.radius*inv(time_scale*3600) / height_scale^2 - ν0 /= (32/(trunc+1))^resolution_scaling*(8/nlev)^nlev_scaling - ν0 = convert(NF,ν0) - - # DO DIFFUSION - @inbounds begin - - # top layer with no flux boundary conditions at k=1/2 - ν∇²_below = ν0*∇²_below[1] # diffusion operator - u_tend[1] += ν∇²_below*(u[2] - u[1]) # diffusion of u - v_tend[1] += ν∇²_below*(v[2] - v[1]) # diffusion of v - temp_tend[1] += ν∇²_below*(temp[2] - temp[1]) # diffusion of temperature - - # full Laplacian in other layers - for k in 2:nlev-1 - # diffusion coefficient ν times 1/Δσ²-like operator - ν∇²_above = ν0*∇²_above[k-1] - ν∇²_below = ν0*∇²_below[k] - ν∇²_at_k = ν∇²_above + ν∇²_below - - # discrete Laplacian, like the (1, -2, 1)-stencil but for variable Δσ - u_tend[k] += ν∇²_below*u[k+1] - ν∇²_at_k*u[k] + ν∇²_above*u[k-1] - v_tend[k] += ν∇²_below*v[k+1] - ν∇²_at_k*v[k] + ν∇²_above*v[k-1] - temp_tend[k] += ν∇²_below*temp[k+1] - ν∇²_at_k*temp[k] + ν∇²_above*temp[k-1] - end - - # bottom layer with no flux boundary conditions at k=nlev+1/2 - ν∇²_above = ν0*∇²_above[end] - u_tend[end] += ν∇²_above*(u[end-1] - u[end]) - v_tend[end] += ν∇²_above*(v[end-1] - v[end]) - temp_tend[end] += ν∇²_above*(temp[end-1] - temp[end]) - end - +function static_energy_diffusion!( column::ColumnVariables, + scheme::NoVerticalDiffusion) return nothing end -function initialize_vertical_diffusion!(K::ParameterizationConstants, - scheme::VerticalLaplacian, - P::Parameters, - G::Geometry) - - (;vert_diff_∇²_above, vert_diff_∇²_below, vert_diff_Δσ) = K - Δσ = G.σ_levels_thick +""" +Diffusion of dry static energy: A relaxation towards a reference +gradient of static energy wrt to geopotential, see Fortran SPEEDY documentation. +$(TYPEDFIELDS)""" +@with_kw struct StaticEnergyDiffusion{NF<:AbstractFloat} <: VerticalDiffusion{NF} + "time scale [hrs] for strength" + time_scale::Float64 = 6 - # thickness Δσ of half levels - @. vert_diff_Δσ = 1/2*(Δσ[2:end] + Δσ[1:end-1]) - - # 1/Δσ² but for variable Δσ on half levels - # = 1/(1/2*Δσₖ(Δσ_k-1 + Δσₖ)) - @. vert_diff_∇²_above = inv(Δσ[2:end]*vert_diff_Δσ) + "[1] ∂SE/∂Φ, vertical gradient of static energy SE with geopotential Φ" + static_energy_lapse_rate::Float64 = 0.1 - # = 1/(1/2*Δσₖ(Δσ_k+1 + Δσₖ)) - @. vert_diff_∇²_below = inv(Δσ[1:end-1]*vert_diff_Δσ) + # precomputations + Fstar::Base.RefValue{NF} = Ref(zero(NF)) # excluding the surface pressure pₛ +end - return nothing -end +StaticEnergyDiffusion(SG::SpectralGrid;kwargs...) = StaticEnergyDiffusion{SG.NF}(;kwargs...) + +"""$(TYPEDSIGNATURES) +Initialize dry static energy diffusion.""" +function initialize!( scheme::StaticEnergyDiffusion{NF}, + model::PrimitiveEquation) where NF -Base.@kwdef struct StaticEnergyDiffusion{NF<:Real} <: VerticalDiffusion{NF} - time_scale::NF = 6.0 # [hours] time scale for strength - static_energy_lapse_rate::NF = 0.1 # [1] ∂SE/∂Φ, vertical gradient of - # static energy SE with geopotential Φ + (;nlev) = model.spectral_grid + (;gravity) = model.planet + C₀ = 1/nlev # average Δσ + + # Fortran SPEEDY documentation equation (70), excluding the surface pressure pₛ + scheme.Fstar[] = convert(NF,C₀/gravity/(scheme.time_scale*3600)) end +"""$(TYPEDSIGNATURES) +Apply dry static energy diffusion.""" function static_energy_diffusion!( column::ColumnVariables{NF}, - scheme::StaticEnergyDiffusion, - model::PrimitiveEquation) where NF - - (;nlev,dry_static_energy,flux_temp_upward,geopot) = column + scheme::StaticEnergyDiffusion) where NF + + (;nlev, dry_static_energy, flux_temp_upward, geopot) = column pₛ = column.pres[end] # surface pressure + Fstar = scheme.Fstar[]*pₛ + Γˢᵉ = convert(NF,scheme.static_energy_lapse_rate) - (;time_scale, static_energy_lapse_rate) = scheme - (;gravity) = model.parameters.planet - - # Fortran SPEEDY documentation equation (70) - C₀ = 1/nlev # average Δσ - Fstar = convert(NF,C₀*pₛ/gravity/(time_scale*3600)) # [hrs] → [s] - Γˢᵉ = convert(NF,static_energy_lapse_rate) - # relax static energy profile back to a reference gradient Γˢᵉ @inbounds for k in 1:nlev-1 # Fortran SPEEDY doc eq (74) @@ -128,8 +60,86 @@ function static_energy_diffusion!( column::ColumnVariables{NF}, end end -function static_energy_diffusion!( column::ColumnVariables, - scheme::NoVerticalDiffusion, - model::PrimitiveEquation) - return nothing -end \ No newline at end of file +# @with_kw struct VerticalLaplacian{NF<:Real} <: VerticalDiffusion{NF} +# time_scale::NF = 10.0 # [hours] time scale to control the strength of vertical diffusion +# height_scale::NF = 100.0 # [m] scales for Δσ so that time_scale is sensible + +# resolution_scaling::NF = 1.0 # (inverse) scaling with resolution T +# nlev_scaling::NF = -2.0 # (inverse) scaling with n vertical levels +# end + +# # generator so that arguments are converted to Float64 +# VerticalLaplacian(;kwargs...) = VerticalLaplacian{Float64}(;kwargs...) + +# function vertical_diffusion!( column::ColumnVariables{NF}, +# scheme::VerticalLaplacian, +# model::PrimitiveEquation) where NF + +# (;nlev,u_tend,v_tend,temp_tend) = column +# (;u,v,temp) = column +# (;time_scale, height_scale, resolution_scaling, nlev_scaling) = scheme +# (;trunc) = model.parameters +# ∇²_below = model.parameterization_constants.vert_diff_∇²_below +# ∇²_above = model.parameterization_constants.vert_diff_∇²_above + +# # GET DIFFUSION COEFFICIENT as a function of u,v,temp and surface pressure +# # *3600 for [hrs] → [s], *1e3 for [km] → [m] +# # include a height scale, technically not needed, but so that the dimensionless +# # 1/Δσ² gets a resonable scale in meters such that the time scale is not +# # counterintuitively in seconds or years +# ν0 = model.geometry.radius*inv(time_scale*3600) / height_scale^2 +# ν0 /= (32/(trunc+1))^resolution_scaling*(8/nlev)^nlev_scaling +# ν0 = convert(NF,ν0) + +# # DO DIFFUSION +# @inbounds begin + +# # top layer with no flux boundary conditions at k=1/2 +# ν∇²_below = ν0*∇²_below[1] # diffusion operator +# u_tend[1] += ν∇²_below*(u[2] - u[1]) # diffusion of u +# v_tend[1] += ν∇²_below*(v[2] - v[1]) # diffusion of v +# temp_tend[1] += ν∇²_below*(temp[2] - temp[1]) # diffusion of temperature + +# # full Laplacian in other layers +# for k in 2:nlev-1 +# # diffusion coefficient ν times 1/Δσ²-like operator +# ν∇²_above = ν0*∇²_above[k-1] +# ν∇²_below = ν0*∇²_below[k] +# ν∇²_at_k = ν∇²_above + ν∇²_below + +# # discrete Laplacian, like the (1, -2, 1)-stencil but for variable Δσ +# u_tend[k] += ν∇²_below*u[k+1] - ν∇²_at_k*u[k] + ν∇²_above*u[k-1] +# v_tend[k] += ν∇²_below*v[k+1] - ν∇²_at_k*v[k] + ν∇²_above*v[k-1] +# temp_tend[k] += ν∇²_below*temp[k+1] - ν∇²_at_k*temp[k] + ν∇²_above*temp[k-1] +# end + +# # bottom layer with no flux boundary conditions at k=nlev+1/2 +# ν∇²_above = ν0*∇²_above[end] +# u_tend[end] += ν∇²_above*(u[end-1] - u[end]) +# v_tend[end] += ν∇²_above*(v[end-1] - v[end]) +# temp_tend[end] += ν∇²_above*(temp[end-1] - temp[end]) +# end + +# return nothing +# end + +# function initialize_vertical_diffusion!(K::ParameterizationConstants, +# scheme::VerticalLaplacian, +# P::Parameters, +# G::Geometry) + +# (;vert_diff_∇²_above, vert_diff_∇²_below, vert_diff_Δσ) = K +# Δσ = G.σ_levels_thick + +# # thickness Δσ of half levels +# @. vert_diff_Δσ = 1/2*(Δσ[2:end] + Δσ[1:end-1]) + +# # 1/Δσ² but for variable Δσ on half levels +# # = 1/(1/2*Δσₖ(Δσ_k-1 + Δσₖ)) +# @. vert_diff_∇²_above = inv(Δσ[2:end]*vert_diff_Δσ) + +# # = 1/(1/2*Δσₖ(Δσ_k+1 + Δσₖ)) +# @. vert_diff_∇²_below = inv(Δσ[1:end-1]*vert_diff_Δσ) + +# return nothing +# end \ No newline at end of file diff --git a/src/physics/vertical_interpolation.jl b/src/physics/vertical_interpolation.jl new file mode 100644 index 000000000..eedd91756 --- /dev/null +++ b/src/physics/vertical_interpolation.jl @@ -0,0 +1,66 @@ +""" + interpolate!( + column::ColumnVariables{NF}, + model::PrimitiveEquation, + ) +""" +function interpolate!( + column::ColumnVariables{NF}, + model::PrimitiveEquation, +) where {NF<:AbstractFloat} + (; humid, humid_half ) = column + (; sat_humid, sat_humid_half ) = column + (; dry_static_energy, dry_static_energy_half ) = column + (; sat_moist_static_energy, sat_moist_static_energy_half ) = column + + for (full, half) in ( + (humid, humid_half), + (sat_humid, sat_humid_half), + (dry_static_energy, dry_static_energy_half), + (sat_moist_static_energy, sat_moist_static_energy_half), + ) + interpolate!(full, half, column, model) + end + + return nothing +end + +""" + interpolate!( + A_full_level::Vector{NF}, + A_half_level::Vector{NF}, + column::ColumnVariables{NF}, + model::PrimitiveEquation, + ) + +Given some generic column variable A defined at full levels, do a linear interpolation in +log(σ) to calculate its values at half-levels. +""" +function interpolate!( + A_full_level::Vector{NF}, + A_half_level::Vector{NF}, + column::ColumnVariables{NF}, + model::PrimitiveEquation, +) where {NF<:AbstractFloat} + (; nlev ) = column + (; σ_levels_full, σ_levels_half ) = model.geometry + + # For A at each full level k, compute A at the half-level below, i.e. at the boundary + # between the full levels k and k+1. + for k = 1:nlev-1 + A_half_level[k] = + A_full_level[k] + + (A_full_level[k+1] - A_full_level[k]) * + (log(σ_levels_half[k+1]) - log(σ_levels_full[k])) / + (log(σ_levels_full[k+1]) - log(σ_levels_full[k])) + end + + # Compute the values at the surface separately + A_half_level[nlev] = + A_full_level[nlev] + + (A_full_level[nlev] - A_full_level[nlev-1]) * + (log(NF(0.99)) - log(σ_levels_full[nlev])) / + (log(σ_levels_full[nlev]) - log(σ_levels_full[nlev-1])) + + return nothing +end \ No newline at end of file diff --git a/src/run_speedy.jl b/src/run_speedy.jl index 4cc6015fe..009995db9 100644 --- a/src/run_speedy.jl +++ b/src/run_speedy.jl @@ -1,96 +1,66 @@ """ - progn_vars = run_speedy(NF,Model;kwargs...) or - progn_vars = run_speedy(NF;kwargs...) or - progn_vars = run_speedy(Model;kwargs...) +$(TYPEDSIGNATURES) +Run a SpeedyWeather.jl `simulation`. The `simulation.model` is assumed to be initialized, +otherwise use `initialize=true` as keyword argument.""" +function run!( simulation::Simulation; + initialize::Bool = false, + n_days::Real = 10, + startdate::Union{Nothing,DateTime} = nothing, + output::Bool = false) + + (;prognostic_variables, diagnostic_variables, model) = simulation -Runs SpeedyWeather.jl with number format `NF` and the model `Model` and any additional parameters -in the keyword arguments `kwargs...`. Any unspecified parameters will use the default values as -defined in [`Parameters`](@ref).""" -function run_speedy(::Type{NF}=DEFAULT_NF, # default number format - ::Type{Model}=DEFAULT_MODEL; # default model - kwargs... # all additional non-default parameters - ) where {NF<:AbstractFloat,Model<:ModelSetup} + # set the clock + if typeof(startdate) == DateTime model.clock.time = startdate end + model.clock.n_days = n_days + initialize!(model.clock,model.time_stepping) - # INITIALIZE MODEL - progn_vars,diagn_vars,model_setup = initialize_speedy(NF,Model;kwargs...) + model.output.output = output # enable/disable output + initialize && initialize!(model) # initialize again? - # START MODEL INTEGRATION - time_stepping!(progn_vars,diagn_vars,model_setup) - return progn_vars # return prognostic variables when finished + # run it, yeah! + time_stepping!(prognostic_variables,diagnostic_variables,model) end -# if only Model M provided, use default number format NF -run_speedy(::Type{Model};kwargs...) where {Model<:ModelSetup} = run_speedy(DEFAULT_NF,Model;kwargs...) - """ - progn_vars, diagn_vars, model_setup = initialize_speedy(NF,Model;kwargs...) or - progn_vars, diagn_vars, model_setup = initialize_speedy(NF,kwargs...) or - progn_vars, diagn_vars, model_setup = initialize_speedy(Model,kwargs...) - -Initialize the model by returning -- `progn_vars`, the initial conditions of the prognostic variables -- `diagn_vars`, the preallocated the diagnotic variables (initialised to zero) -- `model_setup`, the collected pre-calculated structs that don't change throughout integration. + progn_vars = run_speedy(NF,Model;kwargs...) or + progn_vars = run_speedy(NF;kwargs...) or + progn_vars = run_speedy(Model;kwargs...) -The keyword arguments `kwargs` are the same as for `run_speedy`. The `model_setup` contains -fields that hold the parameters, constants, geometry, spectral transform, boundaries and diffusion.""" -function initialize_speedy( ::Type{NF}=DEFAULT_NF, # default number format - ::Type{Model}=DEFAULT_MODEL; # default model - kwargs... # all additional non-default parameters - ) where {NF<:AbstractFloat,Model<:ModelSetup} +Runs SpeedyWeather.jl with number format `NF` and the model `Model` and any additional parameters +in the keyword arguments `kwargs...`. Unspecified parameters use the default values.""" +function run_speedy(::Type{NF} = DEFAULT_NF, # default number format + ::Type{Model} = DEFAULT_MODEL; # default model + spectral_grid::NamedTuple = NamedTuple(), # some keyword arguments to be + planet::NamedTuple = NamedTuple(), # passed on + atmosphere::NamedTuple = NamedTuple(), + time_stepping::NamedTuple = NamedTuple(), + feedback::NamedTuple = NamedTuple(), + output::NamedTuple = NamedTuple(), + clock::NamedTuple = NamedTuple(), + kwargs... + ) where {NF<:AbstractFloat,Model<:ModelSetup} - ConcreteModel = default_concrete_model(Model) # pick default concrete type if Model abstract - P = Parameters{ConcreteModel}(NF=NF;kwargs...) # all model parameters chosen through kwargs - Random.seed!(P.seed) # seed Julia's default RNG for reproducibility - - C = DynamicsConstants(P) # constants used in the dynamical core - G = Geometry(P) # everything grid - S = SpectralTransform(P) # everything spectral transform - B = Boundaries(P,S,G) # arrays for boundary conditions - H = HorizontalDiffusion(P.diffusion,P,C,G,S) # precomputed arrays for horizontal diffusion - D = DeviceSetup(CPUDevice()) # device the model is running on, so far only CPU - - if ConcreteModel <: Barotropic # pack all of the above into a *Model struct - M = BarotropicModel(P,C,G,S,H,D) # typeof(M) is used to dispatch dynamically - elseif ConcreteModel <: ShallowWater # to the supported model types - I = Implicit(P) # precompute arrays for semi-implicit corrections - M = ShallowWaterModel(P,C,G,S,B,H,I,D) - elseif ConcreteModel <: PrimitiveDryCore # no humidity - I = Implicit(P) - K = ParameterizationConstants(P,G) - M = PrimitiveDryCoreModel(P,C,K,G,S,B,H,I,D) - elseif ConcreteModel <: PrimitiveWetCore # with humidity - I = Implicit(P) - K = ParameterizationConstants(P,G) - M = PrimitiveWetCoreModel(P,C,K,G,S,B,H,I,D) - end + # pass on some keyword arguments to the default structs for convenience + spectral_grid = SpectralGrid{Model}(;NF,spectral_grid...) + planet = Earth(;planet...) + atmosphere = EarthAtmosphere(;atmosphere...) + time_stepping = Leapfrog(spectral_grid;time_stepping...) + clock = Clock(time_stepping;clock...) + output = OutputWriter(spectral_grid;output...) + feedback = Feedback(output;feedback...) - prognostic_vars = initial_conditions(M) # initialize prognostic variables - diagnostic_vars = DiagnosticVariables(G,S) # preallocate all diagnostic variables with zeros + # create model with mostly defaults and initalize + ConcreteModel = default_concrete_model(Model) + model = ConcreteModel(;spectral_grid,planet,atmosphere,time_stepping,clock, + output,feedback,kwargs...) + simulation = initialize!(model) - return prognostic_vars, diagnostic_vars, M + # run it, yeah! + (;prognostic_variables, diagnostic_variables, model) = simulation + time_stepping!(prognostic_variables,diagnostic_variables,model) + return simulation.prognostic_variables # return prognostic variables when finished end # if only Model M provided, use default number format NF -initialize_speedy(::Type{Model};kwargs...) where {Model<:ModelSetup} = initialize_speedy(DEFAULT_NF,Model;kwargs...) - -""" - progn = run_speedy!(progn::PrognosticVariables, - diagn::DiagnosticVariables, - M::ModelSetup) - -Convenience function that can be used in combination with `initialize_speedy(args...;kwargs...)` as - - P,D,M = initialize_speedy(kwargs...) - # possibly change P, D, M manually - run_speedy!(P,D,M) - # or investigate D, M afterwards - -to allow for access to the prognostic/diagnostic variables before the time integration is started.""" -function run_speedy!( progn::PrognosticVariables, # all prognostic variables - diagn::DiagnosticVariables, # all pre-allocated diagnostic variables - model::ModelSetup, # all precalculated structs - ) - time_stepping!(progn,diagn,model) - return progn -end \ No newline at end of file +run_speedy(::Type{Model};kwargs...) where {Model<:ModelSetup} = run_speedy(DEFAULT_NF,Model;kwargs...) diff --git a/test/column_variables.jl b/test/column_variables.jl index 22813bca7..c125884ea 100644 --- a/test/column_variables.jl +++ b/test/column_variables.jl @@ -34,7 +34,11 @@ end @testset for NF in (Float32,Float64) nlev = 8 - _,diagn,model = initialize_speedy(NF,PrimitiveEquation,nlev=nlev) + spectral_grid = SpectralGrid(NF;nlev) + model = Model(;spectral_grid) + simulation = initialize!(model) + diagn = simulation.diagnostic_variables + column = ColumnVariables{NF}(;nlev) SpeedyWeather.reset_column!(column) diff --git a/test/diffusion.jl b/test/diffusion.jl index 9568ced4d..38fee7cb4 100644 --- a/test/diffusion.jl +++ b/test/diffusion.jl @@ -1,7 +1,11 @@ @testset "Horizontal diffusion of random" begin - for T in (Float32,Float64) + for NF in (Float32,Float64) - p,d,m = initialize_speedy(T) + spectral_grid = SpectralGrid(NF) + m = Model(;spectral_grid) + simulation = initialize!(m) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables (;vor) = p.layers[1].timesteps[1] (;vor_tend) = d.layers[1].tendencies @@ -24,7 +28,7 @@ vor0 = copy(vor) vor1 = copy(vor) - SpeedyWeather.leapfrog!(vor0,vor1,vor_tend,m.constants.Δt,1,m.constants) + SpeedyWeather.leapfrog!(vor0,vor1,vor_tend,m.time_stepping.Δt,1,m.time_stepping) @test any(vor0 .!= vor1) # check that at least some coefficients are different @test any(vor0 .== vor1) # check that at least some coefficients are identical diff --git a/test/geopotential.jl b/test/geopotential.jl index c5971a89a..0a77d3b32 100644 --- a/test/geopotential.jl +++ b/test/geopotential.jl @@ -1,27 +1,30 @@ @testset "Geopotential reasonable" begin for NF in (Float32,Float64) nlev = 8 - p,d,m = initialize_speedy(NF,PrimitiveWetCore,nlev=nlev, - Grid=FullGaussianGrid) + spectral_grid = SpectralGrid(PrimitiveWet;NF,nlev,Grid=FullGaussianGrid) + model = Model(;spectral_grid) + simulation = initialize!(model) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables # give every layer some constant temperature temp = 280 # in Kelvin lf = 1 for (progn_layer,diagn_layer) in zip(p.layers,d.layers) - progn_layer.timesteps[lf].temp[1] = temp*m.spectral_transform.norm_sphere + progn_layer.timesteps[lf].temp[1] = temp*model.spectral_transform.norm_sphere fill!(progn_layer.timesteps[lf].humid,0) # dry core - SpeedyWeather.gridded!(diagn_layer,progn_layer,lf,m) # propagate spectral state to grid - SpeedyWeather.linear_virtual_temperature!(diagn_layer,progn_layer,m,lf) + SpeedyWeather.gridded!(diagn_layer,progn_layer,lf,model) # propagate spectral state to grid + SpeedyWeather.linear_virtual_temperature!(diagn_layer,progn_layer,model,lf) end - SpeedyWeather.geopotential!(d,m.boundaries,m.geometry) + SpeedyWeather.geopotential!(d,model.orography,model.constants) # approximate heights [m] for this setup heights = [27000,18000,13000,9000,6000,3700,1800,700] for k in 1:8 geopot_grid = Matrix(gridded(d.layers[k].dynamics_variables.geopot)) - height_over_ocean = geopot_grid[48,24]/m.parameters.planet.gravity # middle of pacific + height_over_ocean = geopot_grid[48,24]/model.planet.gravity # middle of pacific @test heights[k] ≈ height_over_ocean rtol=0.5 # very large error allowed end end @@ -30,8 +33,11 @@ end @testset "Add geopotential and kinetic energy, compute -∇²B term, no errors" begin for NF in (Float32,Float64) nlev = 8 - p,d,m = initialize_speedy(NF,PrimitiveWetCore,nlev=nlev, - Grid=FullGaussianGrid) + spectral_grid = SpectralGrid(PrimitiveWet;NF,nlev,Grid=FullGaussianGrid) + m = Model(;spectral_grid) + simulation = initialize!(m) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables # give every layer some constant temperature temp = 280 # in Kelvin @@ -39,7 +45,7 @@ end p.layers[k].timesteps[1].temp[1] = temp*m.spectral_transform.norm_sphere end - SpeedyWeather.geopotential!(d,m.boundaries,m.geometry) + SpeedyWeather.geopotential!(d,m.orography,m.constants) lf = 1 for (progn_layer,diagn_layer) in zip(p.layers,d.layers) @@ -52,8 +58,11 @@ end @testset "Virtual temperature calculation" begin for NF in (Float32,Float64) nlev = 8 - p,d,m = initialize_speedy(NF,PrimitiveWetCore,nlev=nlev, - Grid=FullGaussianGrid) + spectral_grid = SpectralGrid(PrimitiveWet;NF,nlev,Grid=FullGaussianGrid) + m = Model(;spectral_grid) + simulation = initialize!(m) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables # give every layer some constant temperature temp = 280 # in Kelvin diff --git a/test/initialize.jl b/test/initialize.jl deleted file mode 100644 index 0dc270efd..000000000 --- a/test/initialize.jl +++ /dev/null @@ -1,65 +0,0 @@ -@testset "Zero generators" begin - @testset for NF in (Float32,Float64) - P = Parameters{SpeedyWeather.BarotropicModel}(;NF) - G = Geometry(P) - S = SpectralTransform(P) - - P = zeros(PrognosticVariables{NF},5,5,3) - P = zeros(DiagnosticVariables,G,S) - end -end - -@testset "Initialize from rest" begin - - # BAROTROPIC MODEL - progn, diagn, model = initialize_speedy(Barotropic,initial_conditions=StartFromRest()) - for layer in progn.layers - for step in layer.timesteps - @test all(step.vor .== 0) - end - end - - # SHALLOW WATER MODEL - progn, diagn, model = initialize_speedy(ShallowWater,initial_conditions=StartFromRest()) - for layer in progn.layers - for step in layer.timesteps - @test all(step.vor .== 0) - @test all(step.div .== 0) - end - end - @test all(progn.surface.timesteps[1].pres .== 0) - @test all(progn.surface.timesteps[2].pres .== 0) - - # PRIMITIVE EQUATION MODEL - progn, diagn, model = initialize_speedy(PrimitiveDryCore,initial_conditions=StartFromRest()) - for layer in progn.layers - for step in layer.timesteps - @test all(step.vor .== 0) - @test all(step.div .== 0) - end - end - - """ - S = model.geospectral.spectral_transform - k = model.parameters.nlev # test surface layer only at the moment - lf = 1 # first leapfrog index - temp_grid = gridded(progn.temp[:,:,lf,k],S) - pres_surf_grid = gridded(progn.pres_surf[:,:,lf],S) - humid_grid = gridded(progn.humid[:,:,lf,k],S) - - # temperature between 200K and 350K everywhere - # println((sum(temp_grid)/length(temp_grid),minimum(temp_grid),maximum(temp_grid))) - @test all(temp_grid .> 200) - @test all(temp_grid .< 350) - - # surface pressure between log(300hPa) and log(2000hPa) everywhere - # println((sum(pres_surf_grid)/length(pres_surf_grid),minimum(pres_surf_grid),maximum(pres_surf_grid))) - @test all(pres_surf_grid .> log(300)) - @test all(pres_surf_grid .< log(2000)) - - # humidity non-negative everywhere - # humidity has currently values of O(1e9)... - # println((sum(humid_grid)/length(humid_grid),minimum(humid_grid),maximum(humid_grid))) - @test_skip all(humid_grid .>= 0) - """ -end \ No newline at end of file diff --git a/test/netcdf_output.jl b/test/netcdf_output.jl index a0a0a777c..c1404170e 100644 --- a/test/netcdf_output.jl +++ b/test/netcdf_output.jl @@ -1,35 +1,144 @@ -@testset "NetCDF output" begin - import NetCDF - - @testset "Time axis" begin - - function manual_time_axis(dt, n_timesteps) - - time_axis = zeros(Float64, n_timesteps+1) - for i=0:n_timesteps - time_axis[i+1] = Float64(dt) * i - end - time_axis - end +import NetCDF - tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits - - p, d, m = initialize_speedy(Float32, output=true, recalculate_implicit=1000000, output_path=tmp_output_path, n_days=1, output_dt=0, run_id="dense-output-test") - SpeedyWeather.time_stepping!(p, d, m) +@testset "Output on various grids" begin + tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits + n_days = 1 - tmp_read_path = joinpath(tmp_output_path, "run-dense-output-test", "output.nc") - t = NetCDF.ncread(tmp_read_path, "time") - @test t ≈ Int64.(manual_time_axis(m.constants.Δt_sec, m.constants.n_timesteps)) - - # this is a nonsense simulation with way too large timesteps, but it's here to test the time axis output - # 1kyrs simulation - p, d, m = initialize_speedy(Float32, trunc=31, output=true, recalculate_implicit=1000000, output_path=tmp_output_path, n_days=365000, Δt_at_T31=60*24*365*10, output_dt=24*365*10, run_id="long-output-test") - SpeedyWeather.time_stepping!(p, d, m) + # default grid, Float64, ShallowWater + spectral_grid = SpectralGrid(ShallowWater;NF=Float64) + output = OutputWriter(spectral_grid,path=tmp_output_path) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false - tmp_read_path = joinpath(tmp_output_path, "run-long-output-test", "output.nc") - t = NetCDF.ncread(tmp_read_path, "time") - @test t ≈ Int64.(manual_time_axis(m.constants.Δt_sec, m.constants.n_timesteps)) + # default grid, Float32, ShallowWater + spectral_grid = SpectralGrid(ShallowWater;NF=Float32) + output = OutputWriter(spectral_grid,path=tmp_output_path) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false - @test t ≈ SpeedyWeather.load_trajectory("time", m) - end + # FullClenshawGrid, Float32, ShallowWater + spectral_grid = SpectralGrid(ShallowWater;NF=Float32,Grid=FullClenshawGrid) + output = OutputWriter(spectral_grid,path=tmp_output_path) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false + + # OctahedralClenshawGrid, Float32, ShallowWater + spectral_grid = SpectralGrid(ShallowWater;NF=Float32,Grid=OctahedralClenshawGrid) + output = OutputWriter(spectral_grid,path=tmp_output_path) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false + + # HEALPixGrid, Float32, ShallowWater + spectral_grid = SpectralGrid(ShallowWater;NF=Float32,Grid=HEALPixGrid) + output = OutputWriter(spectral_grid,path=tmp_output_path) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false + + # OctaHEALPixGrid, Float32, ShallowWater + spectral_grid = SpectralGrid(ShallowWater;NF=Float32,Grid=OctaHEALPixGrid) + output = OutputWriter(spectral_grid,path=tmp_output_path) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false + + # OctahedralClenshawGrid, as matrix, Float32, ShallowWater + spectral_grid = SpectralGrid(ShallowWater;NF=Float32,Grid=OctahedralClenshawGrid) + output = OutputWriter(spectral_grid,path=tmp_output_path,as_matrix=true) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false + + # OctaHEALPixGrid, as matrix, Float32, PrimitiveDry + spectral_grid = SpectralGrid(PrimitiveDry;NF=Float32,Grid=OctaHEALPixGrid) + output = OutputWriter(spectral_grid,path=tmp_output_path,as_matrix=true) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false + + # OctaHEALPixGrid, as matrix, Float32, but output Float64 PrimitiveDry + spectral_grid = SpectralGrid(PrimitiveDry;NF=Float32,Grid=OctaHEALPixGrid) + output = OutputWriter(spectral_grid,path=tmp_output_path,as_matrix=true,NF=Float64) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true;n_days) + @test simulation.model.feedback.nars_detected == false +end + +@testset "Restart from output file" begin + tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits + + spectral_grid = SpectralGrid(PrimitiveDry) + output = OutputWriter(spectral_grid,path=tmp_output_path,id="restart-test") + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true,n_days=1) + + initial_conditions = StartFromFile(path=tmp_output_path,id="restart-test") + model2 = Model(;spectral_grid,initial_conditions) + simulation2 = initialize!(model2) + + p1 = simulation.prognostic_variables + p2 = simulation2.prognostic_variables + + for varname in propertynames(p1.layers[1].timesteps[1]) + if SpeedyWeather.has(p1, varname) + for (var_new, var_old) in zip(SpeedyWeather.get_var(p1, varname), SpeedyWeather.get_var(p2, varname)) + @test all(var_new .== var_old) + end + end + end + @test all(SpeedyWeather.get_pressure(p1) .== SpeedyWeather.get_pressure(p2)) end + +@testset "Time axis" begin + + function manual_time_axis(dt, n_timesteps) + + time_axis = zeros(Float64, n_timesteps+1) + for i=0:n_timesteps + time_axis[i+1] = Float64(dt) * i + end + time_axis + end + + tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits + + spectral_grid = SpectralGrid(PrimitiveDry) + output = OutputWriter(spectral_grid,path=tmp_output_path,id="dense-output-test",output_dt=0) + model = Model(;spectral_grid,output) + simulation = initialize!(model) + run!(simulation,output=true,n_days=1) + + tmp_read_path = joinpath(model.output.run_path,model.output.filename) + t = NetCDF.ncread(tmp_read_path, "time") + @test t ≈ Int64.(manual_time_axis(model.time_stepping.Δt_sec, model.clock.n_timesteps)) + + # this is a nonsense simulation with way too large timesteps, but it's here to test the time axis output + # for future tests: This simulation blows up because of too large time steps but only a warning is thrown + # at the moment, no error + # 1kyrs simulation + spectral_grid = SpectralGrid(PrimitiveDry) + time_stepping = Leapfrog(spectral_grid,Δt_at_T31=60*24*365*10) + output = OutputWriter(spectral_grid,path=tmp_output_path,id="long-output-test",output_dt=24*365*10) + model = Model(;spectral_grid,output,time_stepping) + simulation = initialize!(model) + run!(simulation,output=true,n_days=365000) + + tmp_read_path = joinpath(model.output.run_path,model.output.filename) + t = NetCDF.ncread(tmp_read_path, "time") + @test t ≈ Int64.(manual_time_axis(model.time_stepping.Δt_sec, model.clock.n_timesteps)) + @test t ≈ SpeedyWeather.load_trajectory("time", model) +end diff --git a/test/run_speedy.jl b/test/run_speedy.jl index 272983b84..1f0a6cf4f 100644 --- a/test/run_speedy.jl +++ b/test/run_speedy.jl @@ -1,17 +1,29 @@ @testset "run_speedy no errors, no blowup" begin - p = run_speedy(Barotropic) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) + # Barotropic + spectral_grid = SpectralGrid(Barotropic) + model = Model(;spectral_grid) + simulation = initialize!(model) + run!(simulation,n_days=10) + @test simulation.model.feedback.nars_detected == false - p = run_speedy(ShallowWater) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(PrimitiveDryCore) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(PrimitiveWetCore) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p,d,m = initialize_speedy() - run_speedy!(p,d,m) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) + # ShallowWater + spectral_grid = SpectralGrid(ShallowWater) + model = Model(;spectral_grid) + simulation = initialize!(model) + run!(simulation,n_days=10) + @test simulation.model.feedback.nars_detected == false + + # PrimitiveDry + spectral_grid = SpectralGrid(PrimitiveDry) + model = Model(;spectral_grid) + simulation = initialize!(model) + run!(simulation,n_days=10) + @test simulation.model.feedback.nars_detected == false + + # PrimitiveWet + spectral_grid = SpectralGrid(PrimitiveWet) + model = Model(;spectral_grid) + simulation = initialize!(model) + run!(simulation,n_days=10) + @test simulation.model.feedback.nars_detected == false end \ No newline at end of file diff --git a/test/run_speedy_with_output.jl b/test/run_speedy_with_output.jl deleted file mode 100644 index 7cadf65d4..000000000 --- a/test/run_speedy_with_output.jl +++ /dev/null @@ -1,61 +0,0 @@ -@testset "Output on various grids" begin - tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits - n_days = 1 - - p = run_speedy(Float64,output=true,output_path=tmp_output_path;n_days) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(Float32,output=true,output_path=tmp_output_path;n_days) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(Float64,Grid=FullClenshawGrid,output=true,output_path=tmp_output_path;n_days) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(Float64,Grid=OctahedralGaussianGrid,output=true,output_path=tmp_output_path;n_days) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(Float64,Grid=OctahedralClenshawGrid,output=true,output_path=tmp_output_path;n_days) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(Float64,Grid=OctahedralClenshawGrid,output_matrix=true,output=true,output_path=tmp_output_path;n_days) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) - - p = run_speedy(Float64,Grid=OctahedralClenshawGrid,output_matrix=true,output_NF=Float32,output=true,output_path=tmp_output_path;n_days) - @test all(isfinite.(p.layers[1].timesteps[1].vor)) -end - -@testset "Restart from output file" begin - tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits - - p1, d1, m1 = initialize_speedy(Float32, ShallowWater, output=true, output_path=tmp_output_path, run_id="restart-test") - run_speedy!(p1, d1, m1) - - p2, d2, m2 = initialize_speedy(Float32, ShallowWater, initial_conditions=StartFromFile(), output_path=tmp_output_path, restart_id="restart-test") - - for varname in propertynames(p1.layers[1].timesteps[1]) - if SpeedyWeather.has(p1, varname) - for (var_new, var_old) in zip(SpeedyWeather.get_var(p1, varname), SpeedyWeather.get_var(p2, varname)) - @test all(var_new .== var_old) - end - end - end - @test all(SpeedyWeather.get_pressure(p1) .== SpeedyWeather.get_pressure(p2)) -end - -@testset "Restart from PrognosticVariables" begin - - p1, d1, m1 = initialize_speedy(Float32, ShallowWater) - run_speedy!(p1, d1, m1) - - p2, d2, m2 = initialize_speedy(Float32, ShallowWater) - copy!(p2, p1) - - for varname in propertynames(p1.layers[1].timesteps[1]) - if SpeedyWeather.has(p1, varname) - for (var_new, var_old) in zip(SpeedyWeather.get_var(p1, varname), SpeedyWeather.get_var(p2, varname)) - @test all(var_new .== var_old) - end - end - end - @test all(SpeedyWeather.get_pressure(p1) .== SpeedyWeather.get_pressure(p2)) -end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index d04da771a..4ddbc9acf 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -32,9 +32,7 @@ include("column_variables.jl") # include("shortwave_radiation.jl") # INITIALIZATION AND INTEGRATION -include("initialize.jl") include("run_speedy.jl") -include("run_speedy_with_output.jl") # OUTPUT include("netcdf_output.jl") \ No newline at end of file diff --git a/test/set_vars.jl b/test/set_vars.jl index 4677995c1..d65023f67 100644 --- a/test/set_vars.jl +++ b/test/set_vars.jl @@ -1,14 +1,18 @@ @testset "Test PrognosticVariables set_vars! and get_var" begin # test setting LowerTriangularMatrices - P, D, M = initialize_speedy(Float64, PrimitiveWetCore, Grid=FullGaussianGrid, initial_conditions=StartFromRest()) - + spectral_grid = SpectralGrid(PrimitiveWet) + initial_conditions = StartFromRest() + M = Model(;spectral_grid,initial_conditions) + simulation = initialize!(M) + P = simulation.prognostic_variables + nlev = M.geometry.nlev lmax = M.spectral_transform.lmax mmax = M.spectral_transform.mmax lf = 1 - sph_data = [rand(LowerTriangularMatrix, lmax+2, mmax+1) for i=1:nlev] + sph_data = [rand(LowerTriangularMatrix{spectral_grid.NF}, lmax+2, mmax+1) for i=1:nlev] SpeedyWeather.set_vorticity!(P, sph_data) SpeedyWeather.set_divergence!(P, sph_data) @@ -35,9 +39,6 @@ @test all(P.layers[i].timesteps[lf].vor .== 0) end - # test setting grids - P, D, M = initialize_speedy(Float64, PrimitiveWetCore, Grid=FullGaussianGrid, initial_conditions=StartFromRest()) - grid_data = [gridded(sph_data[i], M.spectral_transform) for i in eachindex(sph_data)] SpeedyWeather.set_vorticity!(P, grid_data) @@ -54,8 +55,6 @@ end @test all(isapprox(P.surface.timesteps[lf].pres,sph_data[1])) - P, D, M = initialize_speedy(Float64, PrimitiveWetCore, Grid=FullGaussianGrid, initial_conditions=StartFromRest()) - grid_data = [gridded(sph_data[i], M.spectral_transform) for i in eachindex(sph_data)] SpeedyWeather.set_vorticity!(P, grid_data, M) @@ -73,8 +72,18 @@ @test all(isapprox(P.surface.timesteps[lf].pres,sph_data[1])) # test setting matrices - P, D, M = initialize_speedy(Float64, PrimitiveWetCore, Grid=FullGaussianGrid, initial_conditions=StartFromRest()) + spectral_grid = SpectralGrid(PrimitiveWet,Grid=FullGaussianGrid) + initial_conditions = StartFromRest() + M = Model(;spectral_grid,initial_conditions) + simulation = initialize!(M) + P = simulation.prognostic_variables + + nlev = M.geometry.nlev + lmax = M.spectral_transform.lmax + mmax = M.spectral_transform.mmax + lf = 1 + grid_data = [gridded(sph_data[i], M.spectral_transform) for i in eachindex(sph_data)] matrix_data = [Matrix(grid_data[i]) for i in eachindex(grid_data)] SpeedyWeather.set_vorticity!(P, grid_data) diff --git a/test/spectral_gradients.jl b/test/spectral_gradients.jl index f4da12dfb..ca9529faa 100644 --- a/test/spectral_gradients.jl +++ b/test/spectral_gradients.jl @@ -1,15 +1,18 @@ @testset "Divergence of a non-divergent flow zero?" begin @testset for NF in (Float32,Float64) - p,d,m = initialize_speedy( NF, - ShallowWater) + spectral_grid = SpectralGrid(ShallowWater;NF) + m = Model(;spectral_grid) + simulation = initialize!(m) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables fill!(p.layers[1].timesteps[1].vor,0) # make sure vorticity and divergence are 0 fill!(p.layers[1].timesteps[1].div,0) fill!(d.layers[1].tendencies.vor_tend,0) # start with some vorticity only - vor0 = randn(LowerTriangularMatrix{Complex{NF}},p.lmax+2,p.mmax+1) + vor0 = randn(LowerTriangularMatrix{Complex{NF}},p.trunc+2,p.trunc+1) p.layers[1].timesteps[1].vor .= vor0 lf = 1 @@ -45,15 +48,18 @@ end @testset "Curl of an irrotational flow zero?" begin @testset for NF in (Float32,Float64) - p,d,m = initialize_speedy( NF, - ShallowWater) + spectral_grid = SpectralGrid(ShallowWater;NF) + m = Model(;spectral_grid) + simulation = initialize!(m) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables fill!(p.layers[1].timesteps[1].vor,0) # make sure vorticity and divergence are 0 fill!(p.layers[1].timesteps[1].div,0) fill!(d.layers[1].tendencies.div_tend,0) # start with some vorticity only - div0 = randn(LowerTriangularMatrix{Complex{NF}},p.lmax+2,p.mmax+1) + div0 = randn(LowerTriangularMatrix{Complex{NF}},p.trunc+2,p.trunc+1) p.layers[1].timesteps[1].div .= div0 lf = 1 @@ -68,13 +74,14 @@ end G = m.geometry S = m.spectral_transform + C = m.constants # to evaluate ∇×(uv) use curl of vorticity fluxes (=∇×(uv(ζ+f))) with ζ=1,f=0 fill!(d.layers[1].grid_variables.vor_grid,1) - fill!(G.f_coriolis,0) + fill!(C.f_coriolis,0) # calculate uω,vω in spectral space - SpeedyWeather.vorticity_flux_divcurl!(d.layers[1],m,curl=true) + SpeedyWeather.vorticity_flux_curldiv!(d.layers[1],C,G,S,div=true) for div_lm in d.layers[1].tendencies.div_tend @test abs(div_lm) < sqrt(eps(NF)) @@ -90,10 +97,10 @@ end OctahedralClenshawGrid, HEALPixGrid) - p,d,m = initialize_speedy(NF;Grid) - G = m.geometry + SG = SpectralGrid(NF;Grid) + G = Geometry(SG) - A = Grid(randn(NF,G.npoints)) + A = Grid(randn(NF,SG.npoints)) B = copy(A) SpeedyWeather.scale_coslat⁻¹!(A,G) SpeedyWeather.scale_coslat!(A,G) @@ -111,11 +118,10 @@ end @testset "Flipsign in divergence!, curl!" begin @testset for NF in (Float32,Float64) - p,d,m = initialize_speedy( NF, - Barotropic) + SG = SpectralGrid(NF) + S = SpectralTransform(SG) - S = m.spectral_transform - lmax,mmax = p.lmax,p.mmax + lmax,mmax = S.lmax,S.mmax A1 = randn(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) A2 = randn(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) B = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) @@ -134,11 +140,10 @@ end @testset "Add in divergence!, curl!" begin @testset for NF in (Float32,Float64) - p,d,m = initialize_speedy( NF, - Barotropic) + SG = SpectralGrid(NF) + S = SpectralTransform(SG) - S = m.spectral_transform - lmax,mmax = p.lmax,p.mmax + lmax,mmax = S.lmax,S.mmax A1 = randn(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) A2 = randn(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) B = zeros(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) @@ -168,8 +173,11 @@ end @testset "D,ζ -> u,v -> D,ζ" begin @testset for NF in (Float32,Float64) - p,d,m = initialize_speedy( NF, - ShallowWater) + spectral_grid = SpectralGrid(ShallowWater;NF) + m = Model(;spectral_grid) + simulation = initialize!(m) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables # make sure vorticity and divergence are 0 fill!(p.layers[1].timesteps[1].vor,0) @@ -180,7 +188,7 @@ end fill!(d.layers[1].tendencies.div_tend,0) # create initial conditions - lmax,mmax = p.lmax,p.mmax + lmax,mmax = p.trunc,p.trunc vor0 = randn(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) div0 = randn(LowerTriangularMatrix{Complex{NF}},lmax+2,mmax+1) @@ -283,10 +291,16 @@ end @testset "∇×∇=0 and ∇⋅∇=∇²" begin for NF in (Float32,Float64) - p,d,m = initialize_speedy(NF,Grid=FullGaussianGrid) - a = randn(LowerTriangularMatrix{Complex{NF}},33,32) - SpeedyWeather.spectral_truncation!(a,31) + trunc = 31 + spectral_grid = SpectralGrid(ShallowWater;NF,trunc,Grid=FullGaussianGrid) + m = Model(;spectral_grid) + simulation = initialize!(m) + p = simulation.prognostic_variables + d = simulation.diagnostic_variables + + a = randn(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + SpeedyWeather.spectral_truncation!(a) a[:,1] .= real.(a[:,1]) dadx = zero(a) diff --git a/test/spectral_transform.jl b/test/spectral_transform.jl index 5ff3df002..9e89d81f5 100644 --- a/test/spectral_transform.jl +++ b/test/spectral_transform.jl @@ -50,10 +50,10 @@ spectral_resolutions_inexact = (127,255) FullHEALPixGrid, FullOctaHEALPixGrid) - p,d,m = initialize_speedy(NF;trunc,Grid) - S = m.spectral_transform + SG = SpectralGrid(NF;trunc,Grid) + S = SpectralTransform(SG) - alms = copy(p.layers[1].timesteps[1].vor) + alms = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) fill!(alms,0) alms[1,1] = 1 @@ -70,14 +70,15 @@ end @testset "Transform: Recompute, precompute identical results" begin for trunc in spectral_resolutions for NF in (Float32,Float64) - p1,d1,m1 = initialize_speedy(NF;trunc,recompute_legendre=false) - p2,d2,m2 = initialize_speedy(NF;trunc,recompute_legendre=true) - (;vor) = p1.layers[1].timesteps[1] - alms = randn(typeof(vor),size(vor)...) + SG = SpectralGrid(NF;trunc) + S1 = SpectralTransform(SG,recompute_legendre=true) + S2 = SpectralTransform(SG,recompute_legendre=false) - map1 = gridded(alms,m1.spectral_transform) - map2 = gridded(alms,m2.spectral_transform) + alms = randn(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + + map1 = gridded(alms,S1) + map2 = gridded(alms,S2) # is only approx as recompute_legendre may use a different precision @test map1 ≈ map2 @@ -93,13 +94,13 @@ end OctahedralGaussianGrid, OctahedralClenshawGrid) - P = Parameters{SpeedyWeather.BarotropicModel}(;NF,trunc,Grid) - S = SpectralTransform(P) + SG = SpectralGrid(NF;trunc,Grid) + S = SpectralTransform(SG,recompute_legendre=true) lmax = 3 for l in 1:lmax for m in 1:l - alms = zeros(LowerTriangularMatrix{Complex{NF}},S.lmax+2,S.mmax+1) + alms = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) alms[l,m] = 1 map = gridded(alms,S) @@ -122,13 +123,14 @@ end OctaHEALPixGrid, FullHEALPixGrid, FullOctaHEALPixGrid) - P = Parameters{SpeedyWeather.BarotropicModel}(;NF,trunc,Grid) - S = SpectralTransform(P) + + SG = SpectralGrid(NF;trunc,Grid) + S = SpectralTransform(SG,recompute_legendre=true) lmax = 3 for l in 1:lmax for m in 1:l - alms = zeros(LowerTriangularMatrix{Complex{NF}},S.lmax+2,S.mmax+1) + alms = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) alms[l,m] = 1 map = gridded(alms,S) @@ -157,22 +159,14 @@ end # clenshaw-curtis grids are only exact for cubic truncation dealiasing = Grid in (FullGaussianGrid,OctahedralGaussianGrid) ? 2 : 3 - P = Parameters{SpeedyWeather.ShallowWaterModel}(;NF,Grid,trunc,dealiasing) - S = SpectralTransform(P) - G = Geometry(P) - B = Boundaries(P,S,G) + SG = SpectralGrid(NF;trunc,Grid,dealiasing) + S = SpectralTransform(SG,recompute_legendre=false) + O = EarthOrography(SG,smoothing=true,smoothing_truncation=31) + initialize!(O,SpeedyWeather.Earth(),S,Geometry(SG)) - oro_grid = B.orography.orography + oro_grid = O.orography oro_spec = spectral(oro_grid,S) - # smooth orography - lmax = 30 - for m in 1:trunc+1 - for l in max(lmax,m):trunc+2 - oro_spec[l,m] = 0 - end - end - oro_grid1 = gridded(oro_spec,S) oro_spec1 = spectral(oro_grid1,S) oro_grid2 = gridded(oro_spec1,S) diff --git a/test/time_stepping.jl b/test/time_stepping.jl index adf432c69..66cda87c7 100644 --- a/test/time_stepping.jl +++ b/test/time_stepping.jl @@ -18,8 +18,9 @@ end # loop over different precisions @testset for NF in (Float16,Float32,Float64) - P = Parameters{SpeedyWeather.BarotropicModel}(NF=NF) - C = DynamicsConstants(P) + + spectral_grid = SpectralGrid(NF) + L = Leapfrog(spectral_grid) # INITIAL CONDITIONS lmax,mmax = 3,3 @@ -37,7 +38,7 @@ end for i in 2:n_timesteps+1 # always evaluate F with lf = 2 lf = 2 - SpeedyWeather.leapfrog!(X_old,X_new,F(X_new,NF(ω)),NF(2Δt),lf,C) + SpeedyWeather.leapfrog!(X_old,X_new,F(X_new,NF(ω)),NF(2Δt),lf,L) X_out[i] = X_old[1,1] end @@ -57,8 +58,9 @@ end # loop over different precisions @testset for NF in (Float16,Float32,Float64) - P = Parameters{SpeedyWeather.BarotropicModel}(NF=NF) - C = DynamicsConstants(P) + + spectral_grid = SpectralGrid(NF) + L = Leapfrog(spectral_grid) # INITIAL CONDITIONS lmax,mmax = 3,3 @@ -76,7 +78,7 @@ end for i in 2:n_timesteps+1 # always evaluate F with lf = 2 lf = 2 - SpeedyWeather.leapfrog!(X_old,X_new,F(X_new,NF(ω)),NF(2Δt),lf,C) + SpeedyWeather.leapfrog!(X_old,X_new,F(X_new,NF(ω)),NF(2Δt),lf,L) X_out[i] = X_old[1,1] end @@ -85,8 +87,8 @@ end @test M_RAW < 1 # CHECK THAT NO WILLIAM'S FILTER IS WORSE - P = Parameters{SpeedyWeather.BarotropicModel}(NF=NF,williams_filter=1) # Robert's filter only - C = DynamicsConstants(P) + spectral_grid = SpectralGrid(NF) + L = Leapfrog(spectral_grid,william_filter=1) # INITIAL CONDITIONS lmax,mmax = 3,3 @@ -104,7 +106,7 @@ end for i in 2:n_timesteps+1 # always evaluate F with lf = 2 lf = 2 - SpeedyWeather.leapfrog!(X_old,X_new,F(X_new,NF(ω)),NF(2Δt),lf,C) + SpeedyWeather.leapfrog!(X_old,X_new,F(X_new,NF(ω)),NF(2Δt),lf,L) X_out[i] = X_old[1,1] end diff --git a/test/vertical_levels.jl b/test/vertical_levels.jl index a59f9a523..3e831f472 100644 --- a/test/vertical_levels.jl +++ b/test/vertical_levels.jl @@ -1,20 +1,23 @@ @testset "Initialize sigma levels manually" begin - # automatic levels - p,d,m = initialize_speedy(nlev=4) - @test length(m.geometry.σ_levels_half) == 5 - @test length(m.geometry.σ_levels_full) == 4 + spectral_grid = SpectralGrid{PrimitiveDry}(nlev=4) + G = Geometry(spectral_grid) + @test length(G.σ_levels_half) == 5 + @test length(G.σ_levels_full) == 4 # manual levels - p,d,m = initialize_speedy(σ_levels_half=[0,0.4,0.6,1]) - @test m.parameters.nlev == 3 - @test length(m.geometry.σ_levels_half) == 4 - @test length(m.geometry.σ_levels_full) == 3 + σ = SigmaCoordinates([0,0.4,0.6,1]) + spectral_grid = SpectralGrid{PrimitiveDry}(vertical_coordinates=σ) + G = Geometry(spectral_grid) + @test spectral_grid.nlev == 3 + @test length(G.σ_levels_half) == 4 + @test length(G.σ_levels_full) == 3 - # specify both - p,d,m = initialize_speedy(σ_levels_half=[0,0.4,0.6,1],nlev=3) - @test m.parameters.nlev == 3 - @test length(m.geometry.σ_levels_half) == 4 - @test length(m.geometry.σ_levels_full) == 3 - + # specify both + σ = SigmaCoordinates([0,0.4,0.6,1]) + spectral_grid = SpectralGrid{PrimitiveDry}(nlev=3,vertical_coordinates=σ) + G = Geometry(spectral_grid) + @test spectral_grid.nlev == 3 + @test length(G.σ_levels_half) == 4 + @test length(G.σ_levels_full) == 3 end \ No newline at end of file