Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

New SpectralGrid, Model, initialize!, simulation, run! structure #327

Merged
merged 30 commits into from
May 29, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
c09a00d
Gregification v0
milankl May 15, 2023
b982c1a
horizontal diffusion adapted to gregification
milankl May 15, 2023
18ab33b
abstract forcing types
milankl May 15, 2023
2d7cd86
implicit half gregified
milankl May 15, 2023
511a01d
output partially gregified
milankl May 15, 2023
c6ab44d
geometry and constants gregified
milankl May 16, 2023
3fba57e
show for SpectralGrid, Geometry, Earth, EarthAtm
milankl May 18, 2023
859afdd
adiabatic and geopotential in constants
milankl May 18, 2023
76f4b03
next monster gregification
milankl May 19, 2023
ebe88c1
BarotropicModel up and running with gregification
milankl May 22, 2023
7c25f29
Renamed to OutputWriter
milankl May 22, 2023
b85522c
comments, docstrings for feedback.jl
milankl May 22, 2023
ec45751
docstrings
milankl May 22, 2023
b91c9a1
shallow water gregified, up and running
milankl May 23, 2023
fd5aeea
Clock struct introduced
milankl May 23, 2023
21ecfe3
also monolithic run_speedy
milankl May 23, 2023
3921418
Merge branch 'main' into mk/types
milankl May 23, 2023
05a6d2d
implicit for primeq reworked, temperature_average!
milankl May 23, 2023
1591ced
Merge branch 'mk/types' of https://github.com/SpeedyWeather/SpeedyWea…
milankl May 23, 2023
1146f8d
primitive dry core gregified
milankl May 24, 2023
dc31a7a
doc strings for diffusion
milankl May 24, 2023
35aa82a
PrimitiveWetModel defined
milankl May 24, 2023
f3769f8
show and parameters_txt improved
milankl May 25, 2023
23abffc
temperature relaxation gregified
milankl May 25, 2023
1dc76c0
dry static energy diffusion in
milankl May 25, 2023
de0fc64
thermodynamics reworked
milankl May 25, 2023
79b2cac
set_vars tests adapted
milankl May 26, 2023
da4d725
make tests pass with new structure
milankl May 27, 2023
e819649
Base.@kwdef or @with_kw from Parameters.jl
milankl May 28, 2023
cbfc4d3
fix doc build fail
milankl May 28, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -30,4 +30,5 @@ Manifest.toml

# no model output folders
run????/
run-*/
run-*/
run_*/
1 change: 1 addition & 0 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
```

Expand Down
4 changes: 2 additions & 2 deletions docs/src/how_to_run_speedy.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/SpeedyTransforms/SpeedyTransforms.jl
Original file line number Diff line number Diff line change
Expand Up @@ -43,5 +43,6 @@ include("spectral_transform.jl")
include("spectral_gradients.jl")
include("spectral_truncation.jl")
include("spectrum.jl")
include("show.jl")

end
2 changes: 1 addition & 1 deletion src/SpeedyTransforms/aliasing.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions src/SpeedyTransforms/show.jl
Original file line number Diff line number Diff line change
@@ -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
26 changes: 13 additions & 13 deletions src/SpeedyTransforms/spectral_transform.jl
Original file line number Diff line number Diff line change
Expand Up @@ -65,17 +65,17 @@ 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
necessary constants for the spetral transform. Also plans the Fourier transforms, retrieves the colatitudes,
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

Expand Down Expand Up @@ -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

"""
Expand All @@ -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

"""
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down
106 changes: 65 additions & 41 deletions src/SpeedyWeather.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
module SpeedyWeather

# STRUCTURE
import Parameters: @with_kw
using DocStringExtensions

# NUMERICS
Expand All @@ -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,
Expand All @@ -50,6 +66,8 @@ export LowerTriangularMatrix,
HEALPixGrid,
OctaHEALPixGrid

export Leapfrog

# EXPORT OROGRAPHIES
export NoOrography,
EarthOrography,
Expand Down Expand Up @@ -78,23 +96,22 @@ export NoVerticalDiffusion,
VerticalLaplacian

# EXPORT STRUCTS
export Parameters,
DynamicsConstants,
ParameterizationConstants,
Geometry,
export DynamicsConstants,
SpectralTransform,
Boundaries,
PrognosticVariables,
DiagnosticVariables,
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
Expand All @@ -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")
Expand Down
48 changes: 22 additions & 26 deletions src/abstract_types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
abstract type HorizontalDiffusion{NF} end
abstract type AbstractImplicit{NF} end
abstract type TimeStepper{NF} end
Loading
Loading