Restructured dynamical core, variables array-agnostic and 3-dimension…

…al, modular netCDF output (#525) * PrognosticVariables with LowerTriangularArray * WIP: initial indexing changes * WIP: indexing not working yet * WIP indexing (probably) working * WIP: zeros, ones, rand, ... * WIP: copyto! * copyto! working * LTA N-1 dim change tests passing * adjusting for size(alms) * fast cartesian indexing + broadcasting * clean up explicit arithmetics functions * code cleanup broadcasting * remove add! tests * show LowerTriangularMatrix * fix LTA plot * PrognosticVariables new structure * DiagnosticVariables new structure * zeros for Vector{Particle} * no set vars tests due to new structure * SpectralGrid with Spectral/GridVariable2D/3D types * progn/diagn prettier printing * fill!(::Tendencies, x) instead of zero_tendencies! * prognostic variables as NTuple{2, LowerTriangularArray} * BarotropicModel without layer structure * N-dimensional transform! spec->grid * n-dim transform! also grid -> spectral * CartesianIndex{0} and broadcast tests * LTA tests fixed * SpectralTransform adjusted to ndim-1 LTA, still WIP * n-dim leapfrogging * n-dimensional UV_from_vor! * BarotropicModel with n-dim formulation * define prognostic_variables(::ModelSetup) * shallow water model with new array structure * First working shallow water model with new variable structure * Spectral gradients with better size checks * horizontal diffusion with nlayers * implicit, land, ocean adapted to new structure * First half of the primitive equation dynamical core restructured * surface pressure tendency, vertical velocity and linear pres gradient restructured * vertical advection to new structure * vor, div, temp, humid tendencies, implicit and diffusion * physics_tendencies_only! renaming * Further size calls for LTA corrected * replace spectral/gridded functions with transform * transform and gradient tests passing * more tests adapted * extending tests adapted * nlayers instead of nlev; starting to adjust docs * docs rework contiues * benchmark suite nlayers spectral grid call * further docs updated * doc add section about batched transforms * typo in docs * docs: another typo * changelog updated * New 4D Variable Structure: nlayers instead of nlev; starting to adjust docs (#566) * nlayers instead of nlev; starting to adjust docs * docs rework contiues * benchmark suite nlayers spectral grid call * further docs updated * doc add section about batched transforms * typo in docs * docs: another typo * changelog updated * linting * nlayers kwarg in README * more type stable `global_diagnostics` in analysis docs * spectral! -> transform! * NF=Float32 for SpectralGrid --------- Co-authored-by: Milan <[email protected]> * all merged * rename ModelSetup to AbstractModel * NetCDF output modularised (#573) * NetCDF output modularised * docstrings everywhere * precipitation output bugs corrected * changelog added * parameters.txt pulls automatically all show strings from model components * include precip and clouds into standard output for primitive wet * show output tweaks * netcdf output tests adapted and extended * gremlins removed * particle tracker test adapted * bugs in netcdf output tests * include previous bugfix changelog * benchmarks nlayer/nlev fixed * add @inbounds back in * constrain to julia 1.10 * CI with prelease for 1.11 * add dependabot.yml * update CI.yml * move dependabot * Tuple instead of vector of different grids, latter fails in 1.9 * CI with 1.9 again * grid creation with Grid(grid.data) to avoid double nesting * define JLArray method for julia v1.9 * set compat to Julia 1.9 again * JetStreamForcing adapted * bug: wrote restart even without output * output show with restart information * RandomWaves initial conditions adapted * docs: extensions/forcing_drag adapted * transforms with singleton dimensions possible * test real to real transform * Transfroms: More number format flexibility tests * docs build errors addressed * first version NetCDFOutput new documentation * docs: particle advection with 2D shallow water model nlayers=1 * make nlayers>1 warn an error for 2D models * Improved error message for `ismatching` * docs speedytransforms * docs: add variables to netcdf * docs misc * docs: extending netCDF output * custom netcdf output * docs: custom netcdf output finished * Set! function for PrognosticVariables and ICs for new structure (#563) * export set! * reworking UV from vordiv tests with set! * move radius keyword into spectral gradient kernels * set! linting * spectral gradients: typo corrected * @inbounds in dynamics tendencies * FLoops dependency removed * allocation free vertical velocity * (almost) non-allocating vertical advection * general vertical advection stencil via ntuple --------- Co-authored-by: Max <[email protected]>
SpeedyWeather · Sep 25, 2024 · 798f9bb · 798f9bb
1 parent daeedb4
commit 798f9bb
Show file tree

Hide file tree

Showing 118 changed files with 5,245 additions and 4,061 deletions.
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
@@ -4,4 +4,4 @@ updates:
   - package-ecosystem: "github-actions"
     directory: "/" # Location of package manifests
     schedule:
-      interval: "weekly"
+      interval: "weekly"
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -15,6 +15,7 @@ jobs:
         version:
           - '1.9'
           - '1.10'
+          - '~1.11.0-0'
         os:
           - ubuntu-latest
         arch:
@@ -25,16 +26,7 @@ jobs:
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v4
-        env:
-          cache-name: cache-artifacts
-        with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-test-${{ env.cache-name }}-
-            ${{ runner.os }}-test-
-            ${{ runner.os }}-
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:

diff --git a/.gitignore b/.gitignore
@@ -46,4 +46,4 @@ default.profraw
 
 # vs code
 .vscode
-settings.json
+settings.json
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,8 +2,11 @@
 
 ## Unreleased
 
+- Introduced a new `set!` function that allows to set `PrognosticVariables` to new values with keyword arguments
+- Restructured dynamical core with prognostic/diagnostic variables array-agnostic and 3-dimensional [#525](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/525)
+- Modularised NetCDF output [#573](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/573)
+- Fixed a bug in RingGrids, now broadcasts are defined even when the dimensions mismatch in some cases [#568](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/568)
 - RingGrids: To wrap an Array with the horizontal dimension in matrix shape into a full grid, one has to use e.g. `FullGaussianGrid(map, input_as=Matrix)` now. [#572](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/572)
-* RingGrids: Fixed a bug in `RingGrids`, so that now broadcasts are defined even when the dimensions mismatch in some cases
-* CompatHelper: Allow for JLD2.jl v0.5
+- CompatHelper: Allow for JLD2.jl v0.5
 
 ## v0.11.0
diff --git a/Project.toml b/Project.toml
@@ -13,7 +13,6 @@ CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
-FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GenericFFT = "a8297547-1b15-4a5a-a998-a2ac5f1cef28"
@@ -47,7 +46,6 @@ CodecZlib = "0.7"
 Dates = "1.9"
 DocStringExtensions = "0.9"
 FFTW = "1"
-FLoops = "0.2"
 FastGaussQuadrature = "0.4, 0.5, 1"
 GPUArrays = "10"
 GenericFFT = "0.1"

diff --git a/README.md b/README.md
@@ -78,7 +78,7 @@ The interface to SpeedyWeather.jl consist of 5 steps: define the grid, create mo
 construct the model, initialize, run
 
 ```julia
-spectral_grid = SpectralGrid(trunc=31, nlev=8)          # define resolution
+spectral_grid = SpectralGrid(trunc=31, nlayers=8)          # define resolution
 orography = EarthOrography(spectral_grid)               # create non-default components
 model = PrimitiveWetModel(; spectral_grid, orography)   # construct model
 simulation = initialize!(model)                         # initialize all model components

diff --git a/benchmark/benchmark_suite.jl b/benchmark/benchmark_suite.jl
@@ -30,12 +30,12 @@ function run_benchmark_suite!(suite::BenchmarkSuite)
         Model = suite.model[i]
         NF = suite.NF[i]
         trunc = suite.trunc[i]
-        nlev = suite.nlev[i]
+        nlayers = suite.nlev[i]
         Grid = suite.Grid[i]
         dynamics = suite.dynamics[i]
         physics = suite.physics[i]
 
-        spectral_grid = SpectralGrid(;NF, trunc, Grid, nlev)
+        spectral_grid = SpectralGrid(;NF, trunc, Grid, nlayers)
         suite.nlat[i] = spectral_grid.nlat
 
         model = Model(;spectral_grid)
@@ -50,7 +50,7 @@ function run_benchmark_suite!(suite::BenchmarkSuite)
         simulation = initialize!(model)
         suite.memory[i] = Base.summarysize(simulation)
 
-        nsteps = n_timesteps(trunc, nlev)
+        nsteps = n_timesteps(trunc, nlayers)
         period = Second(round(Int,model.time_stepping.Δt_sec * (nsteps+1)))
         run!(simulation; period)
 

diff --git a/benchmark/manual_benchmarking.jl b/benchmark/manual_benchmarking.jl
@@ -36,7 +36,7 @@ write(md, "### Explanation\n\n")
 write(md, "Abbreviations in the tables below are as follows, omitted columns use defaults.\n")
 write(md, "- NF: Number format, default: $(SpeedyWeather.DEFAULT_NF)\n")
 write(md, "- T: Spectral resolution, maximum degree of spherical harmonics, default: T$(SpeedyWeather.DEFAULT_TRUNC)\n")
-write(md, "- L: Number of vertical layers, default: $(SpeedyWeather.DEFAULT_NLEV) (for 3D models)\n")
+write(md, "- L: Number of vertical layers, default: $(SpeedyWeather.DEFAULT_NLAYERS) (for 3D models)\n")
 write(md, "- Grid: Horizontal grid, default: $(SpeedyWeather.DEFAULT_GRID)\n")
 write(md, "- Rings: Grid-point resolution, number of latitude rings pole to pole\n")
 write(md, "- Dynamics: With dynamics?, default: true\n")

diff --git a/docs/make.jl b/docs/make.jl
@@ -1,4 +1,5 @@
-using Documenter, SpeedyWeather
+using Documenter
+using SpeedyWeather
 
 makedocs(
     format = Documenter.HTML(prettyurls=get(ENV, "CI", nothing)=="true",
@@ -25,6 +26,7 @@ makedocs(
                 "Orography"=>"orography.md",
                 "Land-Sea Mask"=>"land_sea_mask.md",
                 "Ocean"=>"ocean.md",
+                "NetCDF output variables"=>"custom_netcdf_output.md",
                 "Callbacks"=>"callbacks.md",
             ],
             "Dynamics" => [

diff --git a/docs/src/analysis.md b/docs/src/analysis.md
@@ -68,7 +68,7 @@ or wavenumber 0, see [Spherical Harmonic Transform](@ref)) encodes the global av
 
 ```@example analysis
 using SpeedyWeather
-spectral_grid = SpectralGrid(trunc=31, nlev=1)
+spectral_grid = SpectralGrid(trunc=31, nlayers=1)
 model = ShallowWaterModel(;spectral_grid)
 simulation = initialize!(model)
 ```
@@ -77,19 +77,19 @@ Now we check ``\eta_{0,0}`` the ``l = m = 0`` coefficent of the inital condition
 of that simulation with
 
 ```@example analysis
-simulation.prognostic_variables.surface.timesteps[1].pres[1]
+simulation.prognostic_variables.pres[1][1]
 ```
 
-`[1]` pulls the first element of the underlying [LowerTriangularMatrix](@ref lowertriangularmatrices)
-which is the coefficient of the ``l = m = 0`` mode.
-Its imaginary part is always zero (which is true for any zonal harmonic ``m=0`` as its
+`[1][1]` pulls the first Leapfrog time step and of that the first element of the underlying
+[LowerTriangularMatrix](@ref lowertriangularmatrices) which is the coefficient of the ``l = m = 0``
+harmonic. Its imaginary part is always zero (which is true for any zonal harmonic ``m=0`` as its
 imaginary part would just unnecessarily rotate something zonally constant in zonal direction),
 so you can `real` it. Also for spherical harmonic transforms there is a norm of the sphere
 by which you have to divide to get your mean value in the original units
 
 ```@example analysis
 a = model.spectral_transform.norm_sphere    # = 2√π = 3.5449078
-η_mean = real(simulation.prognostic_variables.surface.timesteps[1].pres[1]) / a
+η_mean = real(simulation.prognostic_variables.pres[1][1]) / a
 ```
 
 So the initial conditions in this simulation are such that the global mean interface displacement
@@ -105,7 +105,7 @@ model.feedback.verbose = false # hide
 run!(simulation, period=Day(10))
 
 # now we check η_mean again
-η_mean_later = real(simulation.prognostic_variables.surface.timesteps[1].pres[1]) / a
+η_mean_later = real(simulation.prognostic_variables.pres[1][1]) / a
 ```
 
 which is _exactly_ the same. So mass is conserved, woohoo. 
@@ -146,7 +146,7 @@ function total_energy(u, v, η, model)
     E = @. h/2*(u^2 + v^2) + g*h^2  # vertically-integrated mechanical energy
 
     # transform to spectral, take l=m=0 mode at [1] and normalize for mean
-    return E_mean = real(spectral(E)[1]) / model.spectral_transform.norm_sphere
+    return E_mean = real(transform(E)[1]) / model.spectral_transform.norm_sphere
 end
 ```
 
@@ -155,9 +155,9 @@ So at the current state of our simulation we have a total energy
 
 ```@example analysis
 # flat copies for convenience
-u = simulation.diagnostic_variables.layers[1].grid_variables.u_grid
-v = simulation.diagnostic_variables.layers[1].grid_variables.v_grid
-η = simulation.diagnostic_variables.surface.pres_grid
+u = simulation.diagnostic_variables.grid.u_grid[:, 1]
+v = simulation.diagnostic_variables.grid.v_grid[:, 1]
+η = simulation.diagnostic_variables.grid.pres_grid
 
 TE = total_energy(u, v, η, model)
 ```
@@ -209,11 +209,11 @@ as
 
 ```@example analysis
 # vorticity
-ζ = simulation.diagnostic_variables.layers[1].grid_variables.vor_grid
+ζ = simulation.diagnostic_variables.grid.vor_grid[:,1]
 f = coriolis(ζ)     # create f on that grid
 
 # layer thickness
-η = simulation.diagnostic_variables.surface.pres_grid
+η = simulation.diagnostic_variables.grid.pres_grid
 H = model.atmosphere.layer_thickness
 Hb = model.orography.orography
 h = @. η + H - Hb
@@ -225,12 +225,21 @@ nothing # hide
 
 and we can compare the relative vorticity field to
 ```@example analysis
-plot(ζ)
+using CairoMakie
+heatmap(ζ, title="Relative vorticity [1/s]")
+save("analysis_vor.png", ans) # hide
+nothing # hide
 ```
+![Relative vorticity](analysis_vor.png)
+
+
 the potential vorticity
 ```@example analysis
-plot(q)
+heatmap(ζ, title="Potential vorticity [1/ms]")
+save("analysis_pv.png", ans) # hide
+nothing # hide
 ```
+![Potential vorticity](analysis_pv.png)
 
 ## Absolute angular momentum
 
@@ -262,7 +271,7 @@ function total_angular_momentum(u, η, model)
     Λ = @. (u*r + Ω*r^2) * h    # vertically-integrated AAM
 
     # transform to spectral, take l=m=0 mode at [1] and normalize for mean
-    return Λ_mean = real(spectral(Λ)[1]) / model.spectral_transform.norm_sphere
+    return Λ_mean = real(transform(Λ)[1]) / model.spectral_transform.norm_sphere
 end
 ```
 
@@ -299,7 +308,7 @@ function total_circulation(ζ, model)
     f = coriolis(ζ)         # create f on the grid of ζ
     C = ζ .+ f              # absolute vorticity
     # transform to spectral, take l=m=0 mode at [1] and normalize for mean
-    return C_mean = real(spectral(C)[1]) / model.spectral_transform.norm_sphere
+    return C_mean = real(transform(C)[1]) / model.spectral_transform.norm_sphere
 end
 
 total_circulation(ζ, model)
@@ -336,7 +345,7 @@ function total_enstrophy(ζ, η, model)
     Q = @. q^2 / 2      # Potential enstrophy
 
     # transform to spectral, take l=m=0 mode at [1] and normalize for mean
-    return Q_mean = real(spectral(Q)[1]) / model.spectral_transform.norm_sphere
+    return Q_mean = real(transform(Q)[1]) / model.spectral_transform.norm_sphere
 end
 ```
 
@@ -373,14 +382,14 @@ to show how to global integral ``\iint dV`` can be written more efficiently
 # define a global integral, reusing a precomputed SpectralTransform S
 # times surface area of sphere omitted
 function ∬dA(v, h, S::SpectralTransform)
-    return real(spectral(v .* h, S)[1]) / S.norm_sphere
+    return real(transform(v .* h, S)[1]) / S.norm_sphere
 end
 
 # use SpectralTransform from model
-∬dA(v, h, model::ModelSetup) = ∬dA(v, h, model.spectral_transform)
+∬dA(v, h, model::AbstractModel) = ∬dA(v, h, model.spectral_transform)
 ```
 By reusing `model.spectral_transform` we do not have to re-precompute
-the spectral tranform on every call to `spectral`. Providing
+the spectral tranform on every call to `transform`. Providing
 the spectral transform from model as the 2nd argument simply reuses
 a previously precomputed spectral transform which is much faster
 and uses less memory.
@@ -389,22 +398,24 @@ Now the `global_diagnostics` function is defined as
 
 ```@example analysis
 function global_diagnostics(u, v, ζ, η, model)
+    
     # constants from model
+    NF = model.spectral_grid.NF     # number format used
     H = model.atmosphere.layer_thickness
     Hb = model.orography.orography
     R = model.spectral_grid.radius
     Ω = model.planet.rotation
     g = model.planet.gravity
 
-    r = R * cos.(model.geometry.lats)   # create r on that grid
-    f = coriolis(u)                     # create f on that grid
+    r = NF.(R * cos.(model.geometry.lats))  # create r on that grid
+    f = coriolis(u)                         # create f on that grid
     
     h = @. η + H - Hb           # thickness
     q = @. (ζ + f) / h          # potential vorticity
-    λ = @. u * r + Ω * r^2      # angular momentum
-    k = @. 1/2 * (u^2 + v^2)    # kinetic energy
-    p = @. 1/2 * g * h          # potential energy
-    z = @. q^2/2                # potential enstrophy
+    λ = @. u * r + Ω * r^2      # angular momentum (in the right number format NF)
+    k = @. (u^2 + v^2) / 2      # kinetic energy
+    p = @. g * h / 2            # potential energy
+    z = @. q^2 / 2              # potential enstrophy
 
     M = ∬dA(1, h, model)        # mean mass
     C = ∬dA(q, h, model)        # mean circulation
@@ -417,11 +428,11 @@ function global_diagnostics(u, v, ζ, η, model)
 end
 
 # unpack diagnostic variables and call global_diagnostics from above
-function global_diagnostics(diagn::DiagnosticVariables, model::ModelSetup)
-    u = diagn.layers[1].grid_variables.u_grid
-    v = diagn.layers[1].grid_variables.v_grid
-    ζR = diagn.layers[1].grid_variables.vor_grid
-    η = diagn.surface.pres_grid
+function global_diagnostics(diagn::DiagnosticVariables, model::AbstractModel)
+    u = diagn.grid.u_grid[:, 1]
+    v = diagn.grid.v_grid[:, 1]
+    ζR = diagn.grid.vor_grid[:, 1]
+    η = diagn.grid.pres_grid
     
     # vorticity during simulation is scaled by radius R, unscale here
     ζ = ζR ./ diagn.scale[]
@@ -465,7 +476,7 @@ function SpeedyWeather.initialize!(
     callback::GlobalDiagnostics,
     progn::PrognosticVariables,
     diagn::DiagnosticVariables,
-    model::ModelSetup,
+    model::AbstractModel,
 )
     # replace with vector of correct length
     n = progn.clock.n_timesteps + 1    # +1 for initial conditions
@@ -497,7 +508,7 @@ function SpeedyWeather.callback!(
     callback::GlobalDiagnostics,
     progn::PrognosticVariables,
     diagn::DiagnosticVariables,
-    model::ModelSetup,
+    model::AbstractModel,
 )
     callback.timestep_counter += 1  
     i = callback.timestep_counter
@@ -521,7 +532,7 @@ function SpeedyWeather.finish!(
     callback::GlobalDiagnostics,
     progn::PrognosticVariables,
     diagn::DiagnosticVariables,
-    model::ModelSetup,
+    model::AbstractModel,
 )
     n_timesteps = callback.timestep_counter