diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 700707ced..ff6499d68 100644 --- 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" \ No newline at end of file diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index d1fa9702b..5992c5c30 100644 --- 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 index ab9e86bfd..31d4a6c3a 100644 --- a/.gitignore +++ b/.gitignore @@ -46,4 +46,4 @@ default.profraw # vs code .vscode -settings.json \ No newline at end of file +settings.json diff --git a/CHANGELOG.md b/CHANGELOG.md index f6e925f3b..10f5d5ba7 100644 --- 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 index 1178acdc6..ff27d564f 100644 --- 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 index e6ff2a9e6..2b5b07aa8 100644 --- 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 index 16dd879f6..8ca53b7f1 100644 --- 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 index 45de8ba7c..44a2aa64f 100644 --- 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 index 13de6c680..e4646cc67 100644 --- 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 index a13cddee2..0357558c4 100644 --- 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 diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index 80f871a40..b94eb37e5 100644 --- a/docs/src/callbacks.md +++ b/docs/src/callbacks.md @@ -68,14 +68,15 @@ function SpeedyWeather.initialize!( callback::StormChaser, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) # allocate recorder: number of time steps (incl initial conditions) in simulation callback.maximum_surface_wind_speed = zeros(progn.clock.n_timesteps + 1) # where surface (=lowermost model layer) u, v on the grid are stored - (; u_grid, v_grid) = diagn.layers[diagn.nlev].grid_variables - + u_grid = diagn.grid.u_grid[:, diagn.nlayers] + v_grid = diagn.grid.u_grid[:, diagn.nlayers] + # maximum wind speed of initial conditions callback.maximum_surface_wind_speed[1] = max_2norm(u_grid, v_grid) @@ -116,7 +117,7 @@ function SpeedyWeather.callback!( callback::StormChaser, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) # increase counter @@ -124,7 +125,8 @@ function SpeedyWeather.callback!( i = callback.timestep_counter # where surface (=lowermost model layer) u, v on the grid are stored - (; u_grid, v_grid) = diagn.layers[diagn.nlev].grid_variables + u_grid = diagn.grid.u_grid[:, diagn.nlayers] + v_grid = diagn.grid.u_grid[:, diagn.nlayers] # maximum wind speed at current time step callback.maximum_surface_wind_speed[i] = max_2norm(u_grid, v_grid) @@ -317,7 +319,7 @@ function SpeedyWeather.callback!( callback::MyScheduledCallback, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) # scheduled callbacks start with this line to execute only when scheduled! # else escape immediately @@ -325,7 +327,8 @@ function SpeedyWeather.callback!( # Just print the North Pole surface temperature to screen (;time) = progn.clock - temp_at_north_pole = diagn.layers[end].grid_variables.temp_grid[1] + temp_at_north_pole = diagn.grid.temp_grid[1,end] + @info "North pole has a temperature of $temp_at_north_pole on $time." end @@ -349,7 +352,7 @@ is no periodically reoccuring schedule, only `schedule.times` would include some for events that are scheduled. Now let's create a primitive equation model with that callback ```@example schedule -spectral_grid = SpectralGrid(trunc=31, nlev=5) +spectral_grid = SpectralGrid(trunc=31, nlayers=5) model = PrimitiveWetModel(;spectral_grid) model.feedback.verbose = false # hide to progress meter add!(model.callbacks, north_pole_temp_at_noon_jan9) diff --git a/docs/src/custom_netcdf_output.md b/docs/src/custom_netcdf_output.md new file mode 100644 index 000000000..abca410f0 --- /dev/null +++ b/docs/src/custom_netcdf_output.md @@ -0,0 +1,175 @@ +# Customizing netCDF output + +SpeedyWeather's [NetCDF output](@ref) is modularised for the output variables, +meaning you can add relatively easy new variables to be outputted +alongside the default variables in the netCDF file. We explain here +how to define a new output variable largely following the logic +of [Extending SpeedyWeather](@ref). + +## New output variable + +Say we want to output the [Vertical velocity](@ref). In [Sigma coordinates](@ref) +on every time step, one has to integrate the divergence vertically to +know where the flow is not divergence-free, meaning that the horizontally +converging or diverging motion is balanced by a vertical velocity. +This leads to the variable ``\partial \sigma / \partial t``, which +is the equivalent of [Vertical velocity](@ref) in the [Sigma coordinates](@ref). +This variable is calculated and stored at every time step in + +```julia +simulation.diagnostic_variables.dynamics.σ_tend +``` + +So how do we access it and add it the netCDF output? + +First we define `VerticalVelocityOutput` as a new `struct` subtype of +`SpeedyWeather.AbstractOutputVariable` we add the required fields +`name::String`, `unit::String`, `long_name::String` and +`dims_xyzt::NTuple{4, Bool}` (we skip the optional fields +for `missing_value` or compression). + +```@example netcdf_custom +using SpeedyWeather + +@kwdef struct VerticalVelocityOutput <: SpeedyWeather.AbstractOutputVariable + name::String = "w" + unit::String = "s^-1" + long_name::String = "vertical velocity dσ/dt" + dims_xyzt::NTuple{4, Bool} = (true, true, true, true) +end +``` + +By default (using the `@kwdef` macro) we set the dimensions in `dims_xyzt` +to 4D because the vertical velocity is a 3D variable that we want to output +on every time step. So while `dims_xyzt` is a required field for every output variable +you should not actually change it as it is an inherent property of the output +variable. + +You can now add this variable to the `NetCDFOutput` as already described in +[Output variables](@ref) + +```@example netcdf_custom +spectral_grid = SpectralGrid() +output = NetCDFOutput(spectral_grid) +add!(output, VerticalVelocityOutput()) +``` + +Note that here we skip the `SpeedyWeather.` prefix which would point to the +SpeedyWeather scope but we have defined `VerticalVelocityOutput` in +the global scope. + +## Extend the `output!` function + +While we have defined a new output variable we have not actually +defined how to output it. Because in the end we will need to +write that variable into the netcdf file in `NetCDFOutput`, +which we describe now. We have to extend extend +SpeedyWeather's `output!` function with the following +function signature + +```@example netcdf_custom +function SpeedyWeather.output!( + output::NetCDFOutput, + variable::VerticalVelocityOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + # INTERPOLATION + w = output.grid3D # scratch grid to interpolate into + (; σ_tend) = diagn.dynamics # point to data in diagnostic variables + RingGrids.interpolate!(w, σ_tend , output.interpolator) + + # WRITE TO NETCDF + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, :, i] = w + return nothing +end +``` + +The first argument has to be `::NetCDFOutput` as this is +the argument we write into (i.e. mutate). The second argument +has to be `::VerticalVelocityOutput` so that Julia's multiple +dispatch calls this `output!` method for our new variable. +Then the prognostic, diagnostic variables and the model +follows which allows us generally to read any data and use +it to write into the netCDF file. + +In most cases you will need to interpolate any gridded variables +inside the model (which can be on a reduced grd) onto the output grid +(which has to be a full grid, see [Output grid](@ref)). For that +the `NetCDFOutput` has two scratch arrays `grid3D` and `grid2D` +which are of type and size as defined by the `output_Grid` and +`nlat_half` arguments when creating the `NetCDFOutput`. +So the three lines for interpolation are essentially those in +which your definition of a new output variable is linked +with where to find that variable in `diagnostic_variables`. +You can, in principle, also do any kind of computation here, +for example adding two variables, normalising data and so on. +In the end it has to be on the `output_Grid` hence you +probably do not want to skip the interpolation step but you +are generally allowed to do much more here before or after +the interpolation. + +The last two lines are then just about actually writing to +netcdf. For any variable that is written on every output +time step you can use the output counter `i` to point to the +correct index `i` in the netcdf file as shown here. +For 2D variables (horizontal+time) the indexing would be +`[:, :, i]`. 2D variables without time you only want to write +once (because they do not change) the indexing would change to +`[:, :]` and you then probably want to add a line at the top +like `output.output_counter > 1 || return nothing` to escape +immediately after the first output time step. But you could +also check for a specific condition (e.g. a new temperature +record in a given location) and only then write to netcdf. +Just some ideas how to customize this even further. + +## Reading the new variable + +Now let's try this in a primitive dry model + +```@example netcdf_custom +model = PrimitiveDryModel(;spectral_grid, output) +model.output.variables[:w] +``` + +By passing on `output` to the model constructor the output variables +now contain `w` and we see it here as we have defined it earlier. + +```@example netcdf_custom +simulation = initialize!(model) +run!(simulation, period=Day(5), output=true) + +# read netcdf data +using NCDatasets +path = joinpath(model.output.run_path, model.output.filename) +ds = NCDataset(path) +ds["w"] +``` + +Fantastic, it's all there. We wrap this back into a `FullGaussianGrid` +but ignore the mask (there are no masked values) in the netCDF file +which causes a `Union{Missing, Float32}` element type by reading out +the raw data with `.var`. And visualise the vertical velocity +in sigma coordinates (remember this is actually ``\partial \sigma / \partial t``) +of the last time step (index `end`) stored on layer ``k=4`` (counted from the top) + +```@example netcdf_custom +w = FullGaussianGrid(ds["w"].var[:, :, :, :], input_as=Matrix) + +using CairoMakie +heatmap(w[:, 4, end], title="vertical velocity dσ/dt at k=4") +save("sigma_tend.png", ans) # hide +nothing # hide +``` +![Sigma tendency](sigma_tend.png) + +This is now the vertical velocity between layer ``k=4`` and ``k=5``. +You can check that the vertical velocity on layer ``k=8`` is actually +zero (because that is the boundary condition at the surface) +and so would be the velocity between ``k=0`` and ``k=1`` at the top +of the atmosphere, which however is not explicitly stored. +The vertical velocity is strongest on the wind and leeward side of +mountains which is reassuring and all the analysis we want to +do here for now. diff --git a/docs/src/examples_2D.md b/docs/src/examples_2D.md index 952392af5..b821a2c78 100644 --- a/docs/src/examples_2D.md +++ b/docs/src/examples_2D.md @@ -11,7 +11,7 @@ See also [Examples 3D](@ref) for examples with the primitive equation models. !!! info "Setup script to copy and paste" ```julia using SpeedyWeather - spectral_grid = SpectralGrid(trunc=63, nlev=1) + spectral_grid = SpectralGrid(trunc=63, nlayers=1) still_earth = Earth(spectral_grid, rotation=0) initial_conditions = StartWithRandomVorticity() model = BarotropicModel(; spectral_grid, initial_conditions, planet=still_earth) @@ -22,11 +22,11 @@ See also [Examples 3D](@ref) for examples with the primitive equation models. We want to use the barotropic model to simulate some free-decaying 2D turbulence on the sphere without rotation. We start by defining the `SpectralGrid` object. To have a resolution of about 200km, we choose a spectral resolution of -T63 (see [Available horizontal resolutions](@ref)) and `nlev=1` vertical levels. +T63 (see [Available horizontal resolutions](@ref)) and `nlayers=1` vertical levels. The `SpectralGrid` object will provide us with some more information ```@example barotropic_setup using SpeedyWeather -spectral_grid = SpectralGrid(trunc=63, nlev=1) +spectral_grid = SpectralGrid(trunc=63, nlayers=1) ``` Next step we create a planet that's like Earth but not rotating. As a convention, we always pass on the spectral grid object as the first argument to every other @@ -70,7 +70,7 @@ with default settings. More options on output in [NetCDF output](@ref). !!! info "Setup script to copy and past" ```julia using SpeedyWeather - spectral_grid = SpectralGrid(trunc=63, nlev=1) + spectral_grid = SpectralGrid(trunc=63, nlayers=1) orography = NoOrography(spectral_grid) initial_conditions = ZonalJet() model = ShallowWaterModel(; spectral_grid, orography, initial_conditions) @@ -83,7 +83,7 @@ water equations with and without mountains. As the shallow water system has also one level, we can reuse the `SpectralGrid` from Example 1. ```@example galewsky_setup using SpeedyWeather -spectral_grid = SpectralGrid(trunc=63, nlev=1) +spectral_grid = SpectralGrid(trunc=63, nlayers=1) ``` Now as a first simulation, we want to disable any orography, so we create a `NoOrography` ```@example galewsky_setup @@ -141,7 +141,7 @@ Here, we have unpacked the netCDF file using [NCDatasets.jl](https://github.com/ and then plotted via `heatmap(lon, lat, vor)`. While you can do that to give you more control on the plotting, SpeedyWeather.jl also defines an extension for Makie.jl, see [Extensions](@ref). Because if our matrix `vor` here was an `AbstractGrid` (see [RingGrids](@ref)) then all -its geographic information (which grid point is where) would directly be encoded in the type. +its geographic information (which grid point is where) would be implicitly known from the type. From the netCDF file, however, you would need to use the longitude and latitude dimensions. So we can also just do (`input_as=Matrix` here as all our grids use and expect a horizontal dimension @@ -160,15 +160,17 @@ nothing # hide Note that here you need to know which grid the data comes on (an error is thrown if `FullGaussianGrid(vor)` is not size compatible). By default the output will be on the FullGaussianGrid, but if you play around with other grids, you'd need to change this here, -see [NetCDF output on other grids](@ref output_grid). +see [NetCDF output](@ref) and [Output grid](@ref). We did want to showcase the usage of [NetCDF output](@ref) here, but from now on we will use `heatmap` to plot data on our grids directly, without storing output first. So for our current simulation, that means at time = 12 days, vorticity on the grid is stored in the diagnostic variables and can be visualised with +(`[:, 1]` is horizontal x vertical dimension, so all grid points on the first and +only vertical layer) ```@example galewsky_setup -vor = simulation.diagnostic_variables.layers[1].grid_variables.vor_grid +vor = simulation.diagnostic_variables.grid.vor_grid[:, 1] heatmap(vor, title="Relative vorticity [1/s]") save("galewsky2.png", ans) # hide nothing # hide @@ -208,15 +210,15 @@ You could plot the [NetCDF output](@ref) now as before, but we'll be plotting di from the current state of the `simulation` ```@example galewsky_setup -vor = simulation.diagnostic_variables.layers[1].grid_variables.vor_grid +vor = simulation.diagnostic_variables.grid.vor_grid[:, 1] # 1 to index surface heatmap(vor, title="Relative vorticity [1/s]") save("galewsky3.png", ans) # hide nothing # hide ``` ![Galewsky jet](galewsky3.png) -Interesting! The initial conditions have zero velocity in the southern hemisphere, but still, one can see -some imprint of the orography on vorticity. You can spot the coastline of Antarctica; the Andes and +Interesting! One can clearly see some imprint of the orography on vorticity and there is especially +more vorticity in the southern hemisphere. You can spot the coastline of Antarctica; the Andes and Greenland are somewhat visible too. Mountains also completely changed the flow after 12 days, probably not surprising! @@ -225,7 +227,7 @@ probably not surprising! Setup script to copy and paste: ```@example jet_stream_setup using SpeedyWeather -spectral_grid = SpectralGrid(trunc=63, nlev=1) +spectral_grid = SpectralGrid(trunc=63, nlayers=1) forcing = JetStreamForcing(spectral_grid, latitude=60) drag = QuadraticDrag(spectral_grid) @@ -238,7 +240,7 @@ nothing # hide ``` We want to simulate polar jet streams in the shallow water model. We add a `JetStreamForcing` -that adds momentum at 60˚N to inject kinetic energy into the model. This energy needs to be removed +that adds momentum at 60˚N and 60˚S an to inject kinetic energy into the model. This energy needs to be removed (the [diffusion](@ref diffusion) is likely not sufficient) through a drag, we have implemented a `QuadraticDrag` and use the default drag coefficient. Then visualize zonal wind after 40 days with @@ -246,7 +248,7 @@ a `QuadraticDrag` and use the default drag coefficient. Then visualize zonal win ```@example jet_stream_setup using CairoMakie -u = simulation.diagnostic_variables.layers[1].grid_variables.u_grid +u = simulation.diagnostic_variables.grid.u_grid[:, 1] heatmap(u, title="Zonal wind [m/s]") save("polar_jets.png", ans) # hide nothing # hide @@ -261,7 +263,7 @@ Setup script to copy and paste: using Random # hide Random.seed!(1234) # hide using SpeedyWeather -spectral_grid = SpectralGrid(trunc=127, nlev=1) +spectral_grid = SpectralGrid(trunc=127, nlayers=1) # model components time_stepping = SpeedyWeather.Leapfrog(spectral_grid, Δt_at_T31=Minute(30)) @@ -308,7 +310,7 @@ using CairoMakie H = model.atmosphere.layer_thickness Hb = model.orography.orography -η = simulation.diagnostic_variables.surface.pres_grid +η = simulation.diagnostic_variables.grid.pres_grid h = @. η + H - Hb # @. to broadcast grid + scalar - grid heatmap(h, title="Dynamic layer thickness h", colormap=:oslo) diff --git a/docs/src/examples_3D.md b/docs/src/examples_3D.md index 3fc556c61..d764a2dd2 100644 --- a/docs/src/examples_3D.md +++ b/docs/src/examples_3D.md @@ -12,7 +12,7 @@ See also [Examples 2D](@ref Examples) for examples with the ```@example jablonowski using SpeedyWeather -spectral_grid = SpectralGrid(trunc=31, nlev=8, Grid=FullGaussianGrid, dealiasing=3) +spectral_grid = SpectralGrid(trunc=31, nlayers=8, Grid=FullGaussianGrid, dealiasing=3) orography = ZonalRidge(spectral_grid) initial_conditions = InitialConditions( @@ -41,7 +41,7 @@ off a wave propagating eastward. This wave becomes obvious when visualised with ```@example jablonowski using CairoMakie -vor = simulation.diagnostic_variables.layers[end].grid_variables.vor_grid +vor = simulation.diagnostic_variables.grid.vor_grid[:, end] heatmap(vor, title="Surface relative vorticity") save("jablonowski.png", ans) # hide nothing # hide @@ -52,7 +52,7 @@ nothing # hide ```@example heldsuarez using SpeedyWeather -spectral_grid = SpectralGrid(trunc=31, nlev=8) +spectral_grid = SpectralGrid(trunc=31, nlayers=8) # construct model with only Held-Suarez forcing, no other physics model = PrimitiveDryModel(; @@ -99,7 +99,7 @@ Visualising surface temperature with ```@example heldsuarez using CairoMakie -temp = simulation.diagnostic_variables.layers[end].grid_variables.temp_grid +temp = simulation.diagnostic_variables.grid.temp_grid[:, end] heatmap(temp, title="Surface temperature [K]", colormap=:thermal) save("heldsuarez.png", ans) # hide @@ -113,7 +113,7 @@ nothing # hide using SpeedyWeather # components -spectral_grid = SpectralGrid(trunc=31, nlev=8) +spectral_grid = SpectralGrid(trunc=31, nlayers=8) ocean = AquaPlanet(spectral_grid, temp_equator=302, temp_poles=273) land_sea_mask = AquaPlanetMask(spectral_grid) orography = NoOrography(spectral_grid) @@ -149,7 +149,7 @@ of the convection scheme, causing updrafts and downdrafts in both humidity and t ```@example aquaplanet using CairoMakie -humid = simulation.diagnostic_variables.layers[end].grid_variables.humid_grid +humid = simulation.diagnostic_variables.grid.humid_grid[:, end] heatmap(humid, title="Surface specific humidity [kg/kg]", colormap=:oslo) save("aquaplanet.png", ans) # hide @@ -180,7 +180,7 @@ simulation = initialize!(model) model.feedback.verbose = false # hide run!(simulation, period=Day(20)) -humid = simulation.diagnostic_variables.layers[end].grid_variables.humid_grid +humid = simulation.diagnostic_variables.grid.humid_grid[:, end] heatmap(humid, title="No deep convection: Surface specific humidity [kg/kg]", colormap=:oslo) save("aquaplanet_nodeepconvection.png", ans) # hide nothing # hide @@ -201,7 +201,7 @@ simulation = initialize!(model) model.feedback.verbose = false # hide run!(simulation, period=Day(20)) -humid = simulation.diagnostic_variables.layers[end].grid_variables.humid_grid +humid = simulation.diagnostic_variables.grid.humid_grid[:, end] heatmap(humid, title="No convection: Surface specific humidity [kg/kg]", colormap=:oslo) save("aquaplanet_noconvection.png", ans) # hide nothing # hide @@ -218,7 +218,7 @@ And the comparison looks like using SpeedyWeather # components -spectral_grid = SpectralGrid(trunc=31, nlev=8) +spectral_grid = SpectralGrid(trunc=31, nlayers=8) large_scale_condensation = ImplicitCondensation(spectral_grid) convection = SimplifiedBettsMiller(spectral_grid) @@ -241,7 +241,7 @@ the precipitation that comes from these parameterizations ```@example precipitation using CairoMakie -(; precip_large_scale, precip_convection) = simulation.diagnostic_variables.surface +(; precip_large_scale, precip_convection) = simulation.diagnostic_variables.physics m2mm = 1000 # convert from [m] to [mm] heatmap(m2mm*precip_large_scale, title="Large-scale precipiation [mm]: Accumulated over 10 days", colormap=:dense) save("large-scale_precipitation_acc.png", ans) # hide @@ -259,8 +259,8 @@ precipitation only in the period. ```@example precipitation # reset accumulators and simulate 6 hours -simulation.diagnostic_variables.surface.precip_large_scale .= 0 -simulation.diagnostic_variables.surface.precip_convection .= 0 +simulation.diagnostic_variables.physics.precip_large_scale .= 0 +simulation.diagnostic_variables.physics.precip_convection .= 0 run!(simulation, period=Hour(6)) # visualise, precip_* arrays are flat copies, no need to read them out again! diff --git a/docs/src/extensions.md b/docs/src/extensions.md index 5f91bed1b..32cb4fb6c 100644 --- a/docs/src/extensions.md +++ b/docs/src/extensions.md @@ -46,7 +46,7 @@ end In Julia this introduces a new (so-called compound) type that is a subtype of `AbstractForcing`, we have a bunch of these abstract super types defined (see [Abstract model components](@ref)) and you want to piggy-back on them because of multiple-dispatch. This new type could also be -a `mutable struct`, could have keywords defined with `Base.@kwdef` and can +a `mutable struct`, could have keywords defined with `@kwdef` and can also be parametric with respect to the number format `NF` or grid, but let's skip those details for now. Conceptually you include into the type any parameters (example the float `a` here) that you may need and especially those that you want to change @@ -59,9 +59,9 @@ to define a mask here that somehow includes the information of your region. For a more concrete example see [Custom forcing and drag](@ref). To define the new type's initialization, at the most basic level you need -to extend the `initialize!` function for this new type +to extend the `initialize!` function for this new type. A dummy example: ```@example extending -function initialize!(forcing::MyForcing, model::ModelSetup) +function initialize!(forcing::MyForcing, model::AbstractModel) # fill in/change any fields of your new forcing here forcing.v[1] = 1 # you can use information from other model components too @@ -84,24 +84,26 @@ As the last step we have to extend the `forcing!` function which is the function that is called on _every_ step of the time integration. This new method for `forcing!` needs to have the following function signature ```@example extending -function forcing!( diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - forcing::MyForcing, - time::DateTime, - model::ModelSetup) +function forcing!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + forcing::MyForcing, + model::AbstractModel, + lf::Integer, +) # whatever the forcing is supposed to do, in the end you want # to write into the tendency fields - diagn.tendencies.u_tend_grid[1] = forcing.a - diagn.tendencies.v_tend_grid[1] = forcing.a - diagn.tendencies.vor_tend[1] = forcing.a + diagn.tendencies.u_tend_grid = forcing.a + diagn.tendencies.v_tend_grid = forcing.a + diagn.tendencies.vor_tend = forcing.a end ``` -`DiagnosticVariablesLayer` is the type of the first argument, because it contains +`DiagnosticVariables` is the type of the first argument, because it contains the tendencies you will want to change, so this is supposed to be read and write. -The other arguments should be treated read-only. You make use of the `time` -or anything else in the `model`, but the latter likely comes with a performance -penalty which is why we often unpack the model in a function barrier. But let's -skip that detail for now. Generally, try to precompute what you can in +The other arguments should be treated read-only. You can make use of anything else +in `model`, but often we unpack the model in a function barrier (which can help with +type inference and therefore performance). But let's skip that detail for now. +Generally, try to precompute what you can in `initialize!`. For the forcing you will need to force the velocities `u, v` in grid-point space or the vorticity `vor`, divergence `div` in spectral space. This is not a constrain in most applications we came across, but in case it @@ -121,7 +123,7 @@ The `initialize!` function is a function *inside* the SpeedyWeather module, as we want to define a new method for it *outside* that can be called *inside* we actually need to write ```@example extending -function SpeedyWeather.initialize!(forcing::MyForcing, model::SpeedyWeather.ModelSetup) +function SpeedyWeather.initialize!(forcing::MyForcing, model::SpeedyWeather.AbstractModel) # how to initialize it end ``` @@ -131,9 +133,9 @@ You also probably want to make use of functions that are already defined inside SpeedyWeather or its submodules `SpeedyTransforms`, or `RingGrids`. If something does not seem to be defined, although you can see it in the documentation or directly in the code, you probably need to specify its module too! Alternatively, -note that you can also always do `import SpeedWeather: ModelSetup` to bring +note that you can also always do `import SpeedWeather: AbstractModel` to bring a given variable into global scope which removes the necessity to write -`SpeedyWeather.ModelSetup`. +`SpeedyWeather.AbstractModel`. ## Abstract model components @@ -156,13 +158,11 @@ does not mean it is not possible. Just reach out by creating an issue in this case. Similarly, `AbstractParameterization` has several -subtypes that define conceptual classes of parameterizations, -namely +subtypes that define conceptual classes of parameterizations, namely ```@example extending subtypes(SpeedyWeather.AbstractParameterization) ``` -but these are discussed in more detail in -[Parameterizations](@ref). For a more concrete example -of how to define a new forcing for the 2D models, +but these are discussed in more detail in [Parameterizations](@ref). +For a more concrete example of how to define a new forcing for the 2D models, see [Custom forcing and drag](@ref). \ No newline at end of file diff --git a/docs/src/forcing_drag.md b/docs/src/forcing_drag.md index 9c06eeb73..c72fb72d8 100644 --- a/docs/src/forcing_drag.md +++ b/docs/src/forcing_drag.md @@ -122,7 +122,7 @@ end ``` Which allows us to do ```@example extend -spectral_grid = SpectralGrid(trunc=42, nlev=1) +spectral_grid = SpectralGrid(trunc=42, nlayers=1) stochastic_stirring = StochasticStirring(spectral_grid, latitude=30, decorrelation_time=Day(5)) ``` So the respective resolution parameters and the number format are just pulled from the `SpectralGrid` @@ -136,7 +136,7 @@ Now let us have a closer look at the details of the `initialize!` function, in o actually do ```@example extend function SpeedyWeather.initialize!( forcing::StochasticStirring, - model::ModelSetup) + model::AbstractModel) # precompute forcing strength, scale with radius^2 as is the vorticity equation (; radius) = model.spectral_grid @@ -164,7 +164,7 @@ As we want to add a method for the `StochasticStirring` to the `initialize!` fun within `SpeedyWeather` we add the `SpeedyWeather.` to add this method in the right [Scope of variables](@ref). The `initialize!` function _must_ have that function signature, instance of your new type `StochasticStirring` first, then the second argument a -`model` of type `ModelSetup` or, if your forcing (and in general component) _only makes +`model` of type `AbstractModel` or, if your forcing (and in general component) _only makes sense_ in a specific model, you could also write `model::Barotropic` for example, to be more restrictive. Inside the `initialize!` method we are defining we can use parameters from other components. For example, the definition of the `S` term @@ -193,24 +193,38 @@ defined for our new `StochasticStirring` forcing. But if you define it as follow then this will be called automatically with multiple dispatch. ```@example extend -function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - forcing::StochasticStirring, - time::DateTime, - model::ModelSetup) +function SpeedyWeather.forcing!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + forcing::StochasticStirring, + model::AbstractModel, + lf::Integer, +) # function barrier only forcing!(diagn, forcing, model.spectral_transform) end ``` -The function has to be as outlined above. The first argument has to be of type -`DiagnosticVariablesLayer` as it acts on a layer of the diagnostic variables, -where the current model state in grid-point space and the tendencies (in spectral space) -are defined. The second argument has to be a `PrognosticVariablesLayer` because, -in general, the forcing may use the prognostic variables in spectral space. -The third argument has to be of the type of our new forcing, -the third argument is time which you may use or skip, the last element is a `ModelSetup`, -but as before you can be more restrictive to define a forcing only for the -`BarotropicModel` for example, use ``model::Barotropic`` in that case. +The function signature (types and number of its arguments) has to be as outlined above. +The first argument has to be of type `DiagnosticVariables` as the diagnostic variables, +are the ones you want to change (likely the tendencies within) to apply a forcing. +But technically you can change anything else too, although the results may be unexpected. +The diagnostic variables contain the current model state in grid-point space and the +tendencies (in grid and spectral space). The second argument has to be of type +`PrognosticVariables` because, in general, the forcing may use (information from) +the prognostic variables in spectral space, which includes in `progn.clock.time` the current +time for time-dependent forcing. But all prognostic variables should be considered read-only. +The third argument has to be of the type of our new custom forcing, here `StochasticStirring`, +so that multiple dispatch calls the correct method of `forcing!`. The forth argument is of type +`AbstractModel`, so that the forcing can also make use of anything inside `model`, e.g. +`model.geometry` or `model.planet` etc. But you can be more restrictive to define a forcing only +for the `BarotropicModel` for example, use ``model::Barotropic`` in that case. +Or you could define two methods, one for `Barotropic` one for all other models with +`AbstractModel` (not `Barotropic` as a more specific method is prioritised with multiple +dispatch). The 5th argument is the leapfrog index `lf` which after the first time step will +be `lf=2` to denote that tendencies are evaluated at the current time not at the previous time +(how leapfrogging works). Unless you want to read the prognostic variables, for which +you need to know whether to read `lf=1` or `lf=2`, you can ignore this (but need to include +it as argument). As you can see, for now not much is actually happening inside this function, this is what is often called a function barrier, the only thing we do in here @@ -222,9 +236,11 @@ makes that possible. And it also tells you more clearly what a function depends So we define the actual `forcing!` function that's then called as follows ```@example extend -function forcing!( diagn::DiagnosticVariablesLayer, - forcing::StochasticStirring{NF}, - spectral_transform::SpectralTransform) where NF +function forcing!( + diagn::DiagnosticVariables, + forcing::StochasticStirring{NF}, + spectral_transform::SpectralTransform +) where NF # noise and auto-regressive factors a = forcing.a[] # = sqrt(1 - exp(-2dt/τ)) @@ -238,15 +254,15 @@ function forcing!( diagn::DiagnosticVariablesLayer, end # to grid-point space - S_grid = diagn.dynamics_variables.a_grid - SpeedyTransforms.gridded!(S_grid, S, spectral_transform) + S_grid = diagn.dynamics.a_grid # use scratch array "a" + transform!(S_grid, S, spectral_transform) # mask everything but mid-latitudes RingGrids._scale_lat!(S_grid, forcing.lat_mask) # back to spectral space (; vor_tend) = diagn.tendencies - SpeedyTransforms.spectral!(vor_tend, S_grid, spectral_transform) + transform!(vor_tend, S_grid, spectral_transform) return nothing end @@ -293,7 +309,7 @@ modular interface that you can create instances of individual model components and just put them together as you like, and as long as you follow some rules. ```@example extend -spectral_grid = SpectralGrid(trunc=85, nlev=1) +spectral_grid = SpectralGrid(trunc=85, nlayers=1) stochastic_stirring = StochasticStirring(spectral_grid, latitude=-45) initial_conditions = StartFromRest() model = BarotropicModel(; spectral_grid, initial_conditions, forcing=stochastic_stirring) @@ -303,7 +319,7 @@ run!(simulation) # visualisation using CairoMakie -vor = simulation.diagnostic_variables.layers[1].grid_variables.vor_grid +vor = simulation.diagnostic_variables.grid.vor_grid[:, 1] heatmap(vor, title="Stochastically stirred vorticity") save("stochastic_stirring.png", ans) # hide nothing # hide diff --git a/docs/src/gradients.md b/docs/src/gradients.md index 8109fde36..b96a0ecac 100644 --- a/docs/src/gradients.md +++ b/docs/src/gradients.md @@ -60,7 +60,7 @@ using SpeedyWeather, CairoMakie trunc = 64 # 1-based maximum degree of spherical harmonics L = randn(LowerTriangularMatrix{ComplexF32}, trunc, trunc) spectral_truncation!(L, 5) # remove higher wave numbers -G = gridded(L) +G = transform(L) heatmap(G, title="Some fake data G") # requires `using CairoMakie` save("gradient_data.png", ans) # hide nothing # hide @@ -109,7 +109,7 @@ SpeedyTransforms? Let us start by generating some data ```@example gradient -spectral_grid = SpectralGrid(trunc=31, nlev=1) +spectral_grid = SpectralGrid(trunc=31, nlayers=1) forcing = SpeedyWeather.JetStreamForcing(spectral_grid) drag = QuadraticDrag(spectral_grid) model = ShallowWaterModel(; spectral_grid, forcing, drag) @@ -122,8 +122,8 @@ nothing # hide Now pretend you only have `u, v` to get vorticity (which is actually the prognostic variable in the model, so calculated anyway...). ```@example gradient -u = simulation.diagnostic_variables.layers[1].grid_variables.u_grid -v = simulation.diagnostic_variables.layers[1].grid_variables.v_grid +u = simulation.diagnostic_variables.grid.u_grid[:, 1] +v = simulation.diagnostic_variables.grid.v_grid[:, 1] vor = curl(u, v, radius = spectral_grid.radius) nothing # hide ``` @@ -136,8 +136,8 @@ RingGrids.scale_coslat⁻¹!(u) RingGrids.scale_coslat⁻¹!(v) S = SpectralTransform(u, one_more_degree=true) -us = spectral(u, S) -vs = spectral(v, S) +us = transform(u, S) +vs = transform(v, S) vor = curl(us, vs, radius = spectral_grid.radius) ``` @@ -165,6 +165,13 @@ additional degree, but in the returned lower triangular matrix this row is set t Scalar quantities contain this degree too for size compatibility but they should not make use of it. Use `spectral_truncation` to add or remove this degree manually. +You may also generally assume that a `SpectralTransform` struct precomputed for +some truncation, say ``l_{max} = m_{max} = T`` could also be used for smaller +lower triangular matrices. While this is mathematically true, this does not work +here in practice because [`LowerTriangularMatrices`](@ref lowertriangularmatrices) +are implemented as a vector. So always use a `SpectralTransform` struct that +fits matches your resolution exactly (otherwise an error will be thrown). + ## Example: Geostrophy (continued) Now we transfer `vor` into grid-point space, but specify that we want it on the grid @@ -172,7 +179,7 @@ that we also used in `spectral_grid`. The Coriolis parameter for a grid like `vo is obtained, and we do the following for ``f\zeta/g``. ```@example gradient -vor_grid = gridded(vor, Grid=spectral_grid.Grid) +vor_grid = transform(vor, Grid=spectral_grid.Grid) f = coriolis(vor_grid) # create Coriolis parameter f on same grid with default rotation g = model.planet.gravity fζ_g = @. vor_grid * f / g # in-place and element-wise @@ -181,11 +188,11 @@ nothing # hide Now we need to apply the inverse Laplace operator to ``f\zeta/g`` which we do as follows ```@example gradient -fζ_g_spectral = spectral(fζ_g, one_more_degree=true) +fζ_g_spectral = transform(fζ_g, one_more_degree=true) R = spectral_grid.radius η = SpeedyTransforms.∇⁻²(fζ_g_spectral) * R^2 -η_grid = gridded(η, Grid=spectral_grid.Grid) +η_grid = transform(η, Grid=spectral_grid.Grid) nothing # hide ``` Note the manual scaling with the radius ``R^2`` here. We now compare the results @@ -200,7 +207,7 @@ nothing # hide Which is the interface displacement assuming geostrophy. The actual interface displacement contains also ageostrophy ```@example gradient -η_grid2 = simulation.diagnostic_variables.surface.pres_grid +η_grid2 = simulation.diagnostic_variables.grid.pres_grid heatmap(η_grid2, title="Interface displacement η [m] with ageostrophy") save("eta_ageostrophic.png", ans) # hide nothing # hide diff --git a/docs/src/how_to_run_speedy.md b/docs/src/how_to_run_speedy.md index 0d0d1966a..abf077507 100644 --- a/docs/src/how_to_run_speedy.md +++ b/docs/src/how_to_run_speedy.md @@ -67,10 +67,10 @@ spectral_grid = SpectralGrid(trunc=85, dealiasing=3, Grid=HEALPixGrid) ## Vertical coordinates and resolution The number of vertical layers or levels (we use both terms often interchangeably) -is determined through the `nlev` argument. Especially for the +is determined through the `nlayers` argument. Especially for the `BarotropicModel` and the `ShallowWaterModel` you want to set this to ```@example howto -spectral_grid = SpectralGrid(nlev=1) +spectral_grid = SpectralGrid(nlayers=1) ``` For a single vertical level the type of the vertical coordinates does not matter, but in general you can change the spacing of the sigma coordinates @@ -108,10 +108,11 @@ just ignore those. But the `Leapfrog` time stepper comes with `Δt_at_T31` which is the parameter used to scale the time step automatically. This means at a spectral resolution of T31 it would use 30min steps, at T63 it would be ~half that, 15min, etc. Meaning that if you want to have a shorter or longer time step you can create a new -`Leapfrog` time stepper. All time inputs are supposed to be given with the help of `Dates` (e.g. `Minute()`, `Hour()`, ...). But remember that every model component depends on a -`SpectralGrid` as first argument. +`Leapfrog` time stepper. All time inputs are supposed to be given with the help of +`Dates` (e.g. `Minute()`, `Hour()`, ...). But remember that (almost) every model component +depends on a `SpectralGrid` as first argument. ```@example howto -spectral_grid = SpectralGrid(trunc=63, nlev=1) +spectral_grid = SpectralGrid(trunc=63, nlayers=1) time_stepping = Leapfrog(spectral_grid, Δt_at_T31=Minute(15)) ``` The actual time step at the given resolution (here T63) is then `Δt_sec`, there's @@ -196,10 +197,10 @@ simulation.model.output.output_dt = Second(3600) ``` Now, if there's output, it will be every hour. Furthermore the initial conditions can be set with the `initial_conditions` model component -which are then set during `initialize!(::ModelSetup)`, but you can also -change them now, before the model runs +which are then set during `initialize!(::AbstractModel)`, but you can also +change them now, before the model runs ```@example howto -simulation.prognostic_variables.layers[1].timesteps[1].vor[1] = 0 +simulation.prognostic_variables.vor[1][1, 1] = 0 ``` So with this we have set the zero mode of vorticity of the first (and only) layer in the shallow water to zero. Because the leapfrogging is a 2-step diff --git a/docs/src/installation.md b/docs/src/installation.md index 2c6644f86..a5f1dcc78 100644 --- a/docs/src/installation.md +++ b/docs/src/installation.md @@ -25,7 +25,7 @@ julia> ] add https://github.com/SpeedyWeather/SpeedyWeather.jl#main ## Compatibility with Julia versions -SpeedyWeather.jl requires Julia v1.9 or later. The package is tested on Julia 1.9, and 1.10. +SpeedyWeather.jl requires Julia v1.9 or later. The package is tested on Julia 1.9, 1.10, and 1.11. ## Extensions diff --git a/docs/src/land_sea_mask.md b/docs/src/land_sea_mask.md index d0a148943..3b634b030 100644 --- a/docs/src/land_sea_mask.md +++ b/docs/src/land_sea_mask.md @@ -21,7 +21,7 @@ You can create the default land-sea mask as follows ```@example landseamask using SpeedyWeather -spectral_grid = SpectralGrid(trunc=31, nlev=8) +spectral_grid = SpectralGrid(trunc=31, nlayers=8) land_sea_mask = LandSeaMask(spectral_grid) ``` @@ -124,7 +124,7 @@ function SpeedyWeather.initialize!( callback::MilleniumFlood, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) initialize!(callback.schedule, progn.clock) end @@ -133,7 +133,7 @@ function SpeedyWeather.callback!( callback::MilleniumFlood, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) # escape immediately if not scheduled yet isscheduled(callback.schedule, progn.clock) || return nothing diff --git a/docs/src/lowertriangularmatrices.md b/docs/src/lowertriangularmatrices.md index 8f948f5ea..88f4489e0 100644 --- a/docs/src/lowertriangularmatrices.md +++ b/docs/src/lowertriangularmatrices.md @@ -234,7 +234,7 @@ L + L ## GPU `LowerTriangularArray{T, N, ArrayType}` wraps around an array of type `ArrayType`. -If this array is a GPU array (e.g. `CuArray`), all operations are performed on GPU as well. +If this array is a GPU array (e.g. `CuArray`), all operations are performed on GPU as well (work in progress). The implementation was written so that scalar indexing is avoided in almost all cases, so that GPU operation should be performant. To use `LowerTriangularArray` on GPU you can e.g. just `adapt` an existing `LowerTriangularArray`. diff --git a/docs/src/orography.md b/docs/src/orography.md index 8626a2da9..26fb084af 100644 --- a/docs/src/orography.md +++ b/docs/src/orography.md @@ -127,7 +127,7 @@ you want to reflect this in the surface geopotential `geopot_surf` which is used in the primitive equations by ```@example orography -spectral!(orography.geopot_surf, orography.orography, model.spectral_transform) +transform!(orography.geopot_surf, orography.orography, model.spectral_transform) orography.geopot_surf .*= model.planet.gravity spectral_truncation!(orography.geopot_surf) ``` @@ -163,13 +163,13 @@ my_orography = MyOrography(spectral_grid, constant_height=200) Now we have to extend the `initialize!` function. The first argument has to be `::MyOrography` i.e. the new type we just defined, the second argument has to be -`::ModelSetup` although you could constrain it to `::ShallowWater` for example +`::AbstractModel` although you could constrain it to `::ShallowWater` for example but then it cannot be used for primitive equations. ```@example orography function SpeedyWeather.initialize!( orog::MyOrography, # first argument as to be ::MyOrography, i.e. your new type - model::ModelSetup, # second argument, use anything from model read-only + model::AbstractModel, # second argument, use anything from model read-only ) (; orography, geopot_surf) = orog # unpack @@ -181,7 +181,7 @@ function SpeedyWeather.initialize!( # then also calculate the surface geopotential for primitive equations # given orography we just set - spectral!(geopot_surf, orography, model.spectral_transform) + transform!(geopot_surf, orography, model.spectral_transform) geopot_surf .*= model.planet.gravity spectral_truncation!(geopot_surf) return nothing diff --git a/docs/src/output.md b/docs/src/output.md index 56ff7a29e..4499615fa 100644 --- a/docs/src/output.md +++ b/docs/src/output.md @@ -4,32 +4,47 @@ SpeedyWeather.jl uses NetCDF to output the data of a simulation. The following describes the details of this and how to change the way in which the NetCDF output is written. There are many options to this available. -## Accessing the NetCDF output writer - -The output writer is a component of every Model, i.e. `BarotropicModel`, `ShallowWaterModel`, `PrimitiveDryModel` and `PrimitiveWetModel`, hence a non-default output writer can be passed on as a keyword argument to the model constructor +## Creating `NetCDFOutput` ```@example netcdf using SpeedyWeather spectral_grid = SpectralGrid() -output = OutputWriter(spectral_grid, ShallowWater) +output = NetCDFOutput(spectral_grid) +``` + +With `NetCDFOutput(::SpectralGrid, ...)` one creates a `NetCDFOutput` writer with several options, +which are explained in the following. By default, the `NetCDFOutput` is created when constructing +the model, i.e. + +```@example netcdf +model = ShallowWaterModel(;spectral_grid) +model.output +``` + +The output writer is a component of every Model, i.e. `BarotropicModel`, `ShallowWaterModel`, `PrimitiveDryModel` +and `PrimitiveWetModel`, and they only differ in their default `output.variables` (e.g. the primitive +models would by default output temperature which does not exist in the 2D models `BarotropicModel` or `ShallowWaterModel`). +But any `NetCDFOutput` can be passed onto the model constructor with the `output` keyword argument. + +```@example netcdf +output = NetCDFOutput(spectral_grid, Barotropic) model = ShallowWaterModel(; spectral_grid, output=output) nothing # hide ``` -So after we have defined the grid through the `SpectralGrid` object we can use and change -the implemented `OutputWriter` by passing on additional arguments. -The `spectral_grid` has to be the first argument then the model type -(`Barotropic`, `ShallowWater`, `PrimitiveDry`, or `PrimitiveWet`) -which helps the output writer to make default choices on which variables to output. -Then we can also pass on further keyword arguments. So let's start with an example. +Here, we created `NetCDFOutput` for the model class `Barotropic` (2nd positional argument, outputting only vorticity and velocity) +but use it in the `ShallowWaterModel`. By default the `NetCDFOutput` is set to inactive, i.e. +`output.active` is `false`. It is only turned on (and initialized) with `run!(simulation, output=true)`. +So you may change the `NetCDFOutput` as you like but only calling `run!(simulation)` will not +trigger it as `output=false` is the default here. -## Example 1: NetCDF output every hour +## Output frequency If we want to increase the frequency of the output we can choose `output_dt` (default `=Hour(6)`) like so ```@example netcdf -output = OutputWriter(spectral_grid, ShallowWater, output_dt=Hour(1)) +output = NetCDFOutput(spectral_grid, ShallowWater, output_dt=Hour(1)) model = ShallowWaterModel(; spectral_grid, output=output) -nothing # hide +model.output ``` which will now output every hour. It is important to pass on the new output writer `output` to the model constructor, otherwise it will not be part of your model and the default is used instead. @@ -37,14 +52,14 @@ Note that the choice of `output_dt` can affect the actual time step that is used integration, which is explained in the following. Example, we run the model at a resolution of T42 and the time step is going to be ```@example netcdf -spectral_grid = SpectralGrid(trunc=42, nlev=1) +spectral_grid = SpectralGrid(trunc=42, nlayers=1) time_stepping = Leapfrog(spectral_grid) time_stepping.Δt_sec ``` seconds. Depending on the output frequency (we chose `output_dt = Hour(1)` above) this will be slightly adjusted during model initialization: ```@example netcdf -output = OutputWriter(spectral_grid, ShallowWater, output_dt=Hour(1)) +output = NetCDFOutput(spectral_grid, ShallowWater, output_dt=Hour(1)) model = ShallowWaterModel(; spectral_grid, time_stepping, output) simulation = initialize!(model) model.time_stepping.Δt_sec @@ -76,7 +91,7 @@ ds["time"][:] which is a bit ugly, that's why `adjust_with_output=true` is the default. In that case we would have ```@example netcdf time_stepping = Leapfrog(spectral_grid, adjust_with_output=true) -output = OutputWriter(spectral_grid, ShallowWater, output_dt=Hour(1)) +output = NetCDFOutput(spectral_grid, ShallowWater, output_dt=Hour(1)) model = ShallowWaterModel(; spectral_grid, time_stepping, output) simulation = initialize!(model) run!(simulation, period=Day(1), output=true) @@ -86,15 +101,15 @@ ds["time"][:] ``` very neatly hourly output in the NetCDF file! -## [Example 2: Output onto a higher/lower resolution grid](@id output_grid) +## Output grid Say we want to run the model at a given horizontal resolution but want to output on another resolution, -the `OutputWriter` takes as argument `output_Grid<:AbstractFullGrid` and `nlat_half::Int`. +the `NetCDFOutput` takes as argument `output_Grid<:AbstractFullGrid` and `nlat_half::Int`. So for example `output_Grid=FullClenshawGrid` and `nlat_half=48` will always interpolate onto a regular 192x95 longitude-latitude grid of 1.875˚ resolution, regardless the grid and resolution used for the model integration. -```julia -my_output_writer = OutputWriter(spectral_grid, ShallowWater, output_Grid=FullClenshawGrid, nlat_half=48) +```@example netcdf +my_output_writer = NetCDFOutput(spectral_grid, output_Grid=FullClenshawGrid, nlat_half=48) ``` Note that by default the output is on the corresponding full type of the grid type used in the dynamical core so that interpolation only happens at most in the zonal direction as they share the location of the @@ -119,7 +134,52 @@ The grids `FullHEALPixGrid`, `FullOctaHEALPixGrid` share the same latitude rings but have always as many longitude points as they are at most around the equator. These grids are not tested in the dynamical core (but you may use them experimentally) and mostly designed for output purposes. -## Example 3: Changing the output path or identification +## Output variables + +One can easily add or remove variables from being output with the `NetCDFOut` writer. The following +variables are predefined (note they are not exported so you have to prefix `SpeedyWeather.`) + +```@example netcdf +using InteractiveUtils # hide +subtypes(SpeedyWeather.AbstractOutputVariable) +``` + +"Defined" here means that every such type contains information about a variables (long) name, +its units, dimensions, any missing values and compression options. For `HumidityOutput` for example +we have + +```@example netcdf +SpeedyWeather.HumidityOutput() +``` + +You can choose name and unit as you like, e.g. `SpeedyWeather.HumidityOutput(unit = "1")` or change +the compression options, e.g. `SpeedyWeather.HumidityOutput(keepbits = 5)` but more customisation +is discussed in [Customizing netCDF output](@ref). + +We can add new output variables with `add!` + +```@example netcdf +output = NetCDFOutput(spectral_grid) # default variables +add!(output, SpeedyWeather.DivergenceOutput()) # output also divergence +output +``` + +If you didn't create a `NetCDFOutput` separately, you can also apply this directly to `model`, +either `add!(model, SpeedyWeather.DivergenceOutput())` or `add!(model.output, args...)`, +which technically also just forwards to `add!(model.output.variables, args...)`. +`output.variables` is a dictionary were the variable names (as `Symbol`s) are used as keys, +so `output.variables[:div]` just returns the `SpeedyWeather.DivergenceOutput()` we have +just created using `:div` as key. With those keys one can also `delete!` a variable +from netCDF output + +```@example netcdf +delete!(output, :div) +``` + +If you change the `name` of an output variable, i.e. `SpeedyWeather.DivergenceOutput(name="divergence")` +the key would change accordingly to `:divergence`. + +## Output path and identification That's easy by passing on `path="/my/favourite/path/"` and the folder `run_*` with `*` the identification of the run (that's the `id` keyword, which can be manually set but is also automatically determined as a @@ -127,11 +187,11 @@ number counting up depending on which folders already exist) will be created wit ```julia julia> path = pwd() "/Users/milan" -julia> my_output_writer = OutputWriter(spectral_grid, PrimitiveDry, path=path) +julia> my_output_writer = NetCDFOutput(spectral_grid, path=path) ``` This folder must already exist. If you want to give your run a name/identification you can pass on `id` ```julia -julia> my_output_writer = OutputWriter(spectral_grid, PrimitiveDry, id="diffusion_test"); +julia> my_output_writer = NetCDFOutput(spectral_grid, id="diffusion_test"); ``` which will be used instead of a 4 digit number like 0001, 0002 which is automatically determined if `id` is not provided. You will see the id of the run in the progress bar @@ -148,11 +208,11 @@ run_diffusion_test ## Further options -Further options are described in the `OutputWriter` docstring, (also accessible via `julia>?OutputWriter` for example). +Further options are described in the `NetCDFOutput` docstring, (also accessible via `julia>?NetCDFOutput` for example). Note that some fields are actual options, but others are derived from the options you provided or are arrays/objects the output writer needs, but shouldn't be passed on by the user. The actual options are declared as `[OPTION]` in the following ```@example netcdf -@doc OutputWriter +@doc NetCDFOutput ``` diff --git a/docs/src/parameterizations.md b/docs/src/parameterizations.md index e84a8803b..4a92d029b 100644 --- a/docs/src/parameterizations.md +++ b/docs/src/parameterizations.md @@ -98,7 +98,7 @@ all other arguments can then be passed on as keyword arguments with defaults defined. Creating the default convection parameterization for example would be ```@example parameterization using SpeedyWeather -spectral_grid = SpectralGrid(trunc=31, nlev=8) +spectral_grid = SpectralGrid(trunc=31, nlayers=8) convection = SimplifiedBettsMiller(spectral_grid, time_scale=Hour(4)) ``` Further keyword arguments can be added or omitted all together (using the default diff --git a/docs/src/particles.md b/docs/src/particles.md index 3f14debd9..54945a8df 100644 --- a/docs/src/particles.md +++ b/docs/src/particles.md @@ -186,9 +186,9 @@ would not just convert from `Float64` to `Float32` but also from an active to an inactive particle. In SpeedyWeather all particles can be activated or deactivated at any time. -First, you create a [`SpectralGrid`](@ref) with the `n_particles` keyword +First, you create a [`SpectralGrid`](@ref) with the `nparticles` keyword ```@example particle -spectral_grid = SpectralGrid(n_particles = 3) +spectral_grid = SpectralGrid(nparticles = 3) ``` Then the particles live as `Vector{Particle}` inside the prognostic variables ```@example particle @@ -246,7 +246,7 @@ the particle locations via netCDF. We can create it like ```@example particle_tracker using SpeedyWeather -spectral_grid = SpectralGrid(n_particles = 100) +spectral_grid = SpectralGrid(nparticles = 100, nlayers=1) particle_tracker = ParticleTracker(spectral_grid, schedule=Schedule(every=Hour(3))) ``` @@ -285,18 +285,18 @@ ds["lat"] ``` where the last two lines are lazy loading a matrix with each row a particle and each column a time step. You may do `ds["lon"][:,:]` to obtain the full `Matrix`. We had specified -`spectral_grid.n_particles` above and we will have time steps in this file +`spectral_grid.nparticles` above and we will have time steps in this file depending on the `period` the simulation ran for and the `particle_tracker.Δt` output frequency. We can visualise the particles' trajectories with ```@example particle_tracker lon = ds["lon"][:,:] lat = ds["lat"][:,:] -n_particles = size(lon,1) +nparticles = size(lon,1) using CairoMakie fig = lines(lon[1, :], lat[1, :]) # first particle only -[lines!(fig.axis, lon[i,:], lat[i,:]) for i in 2:n_particles] # add lines for other particles +[lines!(fig.axis, lon[i,:], lat[i,:]) for i in 2:nparticles] # add lines for other particles # display updated figure fig @@ -319,7 +319,7 @@ using GeoMakie, CairoMakie fig = Figure() ga = GeoAxis(fig[1, 1]; dest = "+proj=ortho +lon_0=45 +lat_0=45") -[lines!(ga, lon[i,:], lat[i,:]) for i in 1:n_particles] +[lines!(ga, lon[i,:], lat[i,:]) for i in 1:nparticles] fig save("particles_geomakie.png", fig) # hide nothing # hide diff --git a/docs/src/speedytransforms.md b/docs/src/speedytransforms.md index 4bb14cf05..47b347c25 100644 --- a/docs/src/speedytransforms.md +++ b/docs/src/speedytransforms.md @@ -17,6 +17,16 @@ Combined with the spectral transform, you could for example start with a velocit field in grid-point space, transform to spectral, compute its divergence and transform back to obtain the divergence in grid-point space. Examples are outlined in [Gradient operators](@ref). +## Notation: Spectral resolution + +There are different ways to describe the spectral resolution, the truncation wavenumber (e.g. T31), +the maximum degree ``l`` and order ``m`` of the spherical harmonics (e.g. ``l_{max}=31``, ``m_{max} = 31``), +or the size of the lower triangular matrix, e.g. 32x32. In this example, they are all equivalent. +We often use the truncation, i.e. T31, for brevity but sometimes it is important to describe +degree and order independently (see for example [One more degree for spectral fields](@ref)). +Note also how truncation, degree and order are 0-based, but matrix sizes are 1-based. + + ## Example transform Lets start with a simple transform. We could be `using SpeedyWeather` but to be more verbose @@ -29,20 +39,22 @@ using SpeedyWeather.SpeedyTransforms ``` As an example, we want to transform the ``l=m=1`` spherical harmonic from spectral space in `alms` -to grid-point space. +to grid-point space. Note, the ``+1`` on both degree (first index) and order (second index) for +0-based harmonics versus 1-based matrix indexing, see [Size of `LowerTriangularArray`](@ref). +Create a `LowerTriangularMatrix` for T5 resolution, i.e. 6x6 matrix size ```@example speedytransforms -alms = zeros(LowerTriangularMatrix{ComplexF64}, 6, 6) # spectral coefficients -alms[2, 2] = 1 # only l=1, m=1 harmonic +alms = zeros(LowerTriangularMatrix{ComplexF64}, 6, 6) # spectral coefficients T5 +alms[2, 2] = 1 # only l=1, m=1 harmonic alms ``` -Now `gridded` is the function that takes spectral coefficients `alms` and converts -them a grid-point space `map` +Now `transform` is the function that takes spectral coefficients `alms` and converts +them a grid-point space `map` (or vice versa) ```@example speedytransforms -map = gridded(alms) +map = transform(alms) ``` -By default, the `gridded` transforms onto a [`FullGaussianGrid`](@ref FullGaussianGrid) unravelled here +By default, the `transforms` transforms onto a [`FullGaussianGrid`](@ref FullGaussianGrid) unravelled here into a vector west to east, starting at the prime meridian, then north to south, see [RingGrids](@ref). We can visualize `map` quickly with a UnicodePlot via `plot` (see [Visualising RingGrid data](@ref)) ```@example speedytransforms @@ -50,9 +62,9 @@ import SpeedyWeather.RingGrids: plot # not necessary when `using SpeedyWeathe plot(map) ``` Yay! This is the what the ``l=m=1`` spherical harmonic is supposed to look like! -Now let's go back to spectral space with `spectral` +Now let's go back to spectral space with `transform` ```@example speedytransforms -alms2 = spectral(map) +alms2 = transform(map) ``` Comparing with `alms` from above you can see that the transform is exact up to a typical rounding error from `Float64`. @@ -70,15 +82,15 @@ a transform error that is larger than the rounding error from floating-point ari While the default grid for [SpeedyTransforms](@ref) is the [`FullGaussianGrid`](@ref FullGaussianGrid) we can transform onto other grids by specifying `Grid` too ```@example speedytransforms -map = gridded(alms, Grid=HEALPixGrid) +map = transform(alms, Grid=HEALPixGrid) plot(map) ``` which, if transformed back, however, can yield a larger transform error as discussed above ```@example speedytransforms -spectral(map) +transform(map) ``` On such a coarse grid the transform error (absolute and relative) is about ``10^{-2}``, this decreases -for higher resolution. The `gridded` and `spectral` functions will choose a corresponding +for higher resolution. The `transform` function will choose a corresponding grid-spectral resolution (see [Matching spectral and grid resolution](@ref)) following quadratic truncation, but you can always truncate/interpolate in spectral space with `spectral_truncation`, `spectral_interpolation` which takes `trunc` = ``l_{max} = m_{max}`` as second argument @@ -90,20 +102,23 @@ order 5 before. If the second argument in `spectral_truncation` is larger than `alms` then it will automatically call `spectral_interpolation` and vice versa. Also see [Interpolation on RingGrids](@ref) to interpolate directly between grids. If you want to control directly the resolution of the -grid `gridded` is supposed to transform onto you have to provide a `SpectralTransform` instance. +grid you want to `transform` onto, use the keyword `dealiasing` (default: 2 for quadratic, +see [Matching spectral and grid resolution](@ref)). +But you can also provide a `SpectralTransform` instance to reuse a precomputed spectral transform. More on that now. ## [The `SpectralTransform` struct](@id SpectralTransform) -Both `spectral` and `gridded` create an instance of `SpectralTransform` under the hood. This object -contains all precomputed information that is required for the transform, either way: +The function `transform` only with arguments as shown above, +will create an instance of `SpectralTransform` under the hood. +This object contains all precomputed information that is required for the transform, either way: The Legendre polynomials, pre-planned Fourier transforms, precomputed gradient, divergence and curl operators, the spherical harmonic eigenvalues among others. Maybe the most intuitive way to create a `SpectralTransform` is to start with a `SpectralGrid`, which already defines which spectral resolution is supposed to be combined with a given grid. ```@example speedytransforms using SpeedyWeather -spectral_grid = SpectralGrid(Float32, trunc=5, Grid=OctahedralGaussianGrid, dealiasing=3) +spectral_grid = SpectralGrid(NF=Float32, trunc=5, Grid=OctahedralGaussianGrid, dealiasing=3) ``` (We `using SpeedyWeather` here as `SpectralGrid` is exported therein). We also specify the number format `Float32` here to be used for the transform although this @@ -119,31 +134,30 @@ recomputation or pre-computation of the Legendre polynomials, fore more informat [(P)recompute Legendre polynomials](@ref). Passing on `S` the `SpectralTransform` now allows us to transform directly on the grid -defined therein. +defined therein. Note that we recreate `alms` to be of size 7x6 instead of 6x6 for T5 +spectral resolution because SpeedyWeather uses internally [One more degree for spectral fields](@ref) +meaning also that's the default when creating a `SpectralTransform` from a `SpectralGrid`. +But results don't change if the last degree (row) contains only zeros. + ```@example speedytransforms -map = gridded(alms, S) +alms = zeros(LowerTriangularMatrix{ComplexF64}, 7, 6) # spectral coefficients +alms[2, 2] = 1 # only l=1, m=1 harmonic + +map = transform(alms, S) plot(map) ``` + Yay, this is again the ``l=m=1`` harmonic, but this time on a slightly higher resolution `OctahedralGaussianGrid` as specified in the `SpectralTransform` `S`. Note that also the number format was converted on the fly to Float32 because that is the number format we specified in `S`! And from grid to spectral ```@example speedytransforms -alms2 = spectral(map, S) +alms2 = transform(map, S) ``` As you can see the rounding error is now more like ``10^{-8}`` as we are using Float32 (the [OctahedralGaussianGrid](@ref OctahedralGaussianGrid) is another _exact_ grid). -Note, however, that the returned `LowerTriangularMatrix` is of size 7x6, not 6x6 -what we started from. The underlying reason is that internally SpeedyWeather uses -`LowerTriangularMatrix`s of size ``l_{max} + 2 \times m_{max} + 1``. -One ``+1`` on both degree and order for 0-based harmonics versus 1-based matrix sizes, -but an additional ``+1`` for the degrees which is required by the meridional derivative. -For consistency, all `LowerTriangularMatrix`s in SpeedyWeather.jl carry this additional degree -but only the vector quantities explicitly make use of it. -See [Meridional derivative](@ref) for details. - -For this interface to SpeedyTransforms this means that on a grid-to-spectral transform you will get -one more degree than orders of the spherical harmonics by default. You can, however, always truncate +While for this interface to SpeedyTransforms this means that on a grid-to-spectral transform you will +get one more degree than orders of the spherical harmonics by default. You can, however, always truncate this additional degree, say to T5 (hence matrix size is 6x6) ```@example speedytransforms spectral_truncation(alms2, 5, 5) @@ -166,17 +180,19 @@ Say you have some global data in a matrix `m` that looks, for example, like ```@example speedytransforms alms = randn(LowerTriangularMatrix{Complex{Float32}}, 32, 32) # hide spectral_truncation!(alms, 10) # hide -map = gridded(alms, Grid=FullClenshawGrid) # hide +map = transform(alms, Grid=FullClenshawGrid) # hide m = Matrix(map) # hide m ``` You hopefully know which grid this data comes on, let us assume it is a regular -latitude-longitude grid, which we call the `FullClenshawGrid`. Note that for the spectral -transform this should not include the poles, so the 96x47 matrix size here corresponds -to - -We now wrap this matrix -therefore to associate it with the necessary grid information +latitude-longitude grid, which we call the `FullClenshawGrid` (in analogy to the Gaussian grid based +on the Gaussian quadrature). Note that for the spectral transform this should not include the poles, +so the 96x47 matrix size here corresponds to 23 latitudes north and south of the Equator respectively +plus the equator (=47). + +We now wrap this matrix into a `FullClenshawGrid` (`input_as=Matrix` is required because all +grids organise their data as vectors, see [Creating data on a RingGrid](@ref)) +therefore to associate it with the necessary grid information like its coordinates ```@example speedytransforms map = FullClenshawGrid(m, input_as=Matrix) @@ -189,7 +205,7 @@ nothing # hide Now we transform into spectral space and call `power_spectrum(::LowerTriangularMatrix)` ```@example speedytransforms -alms = spectral(map) +alms = transform(map) power = SpeedyTransforms.power_spectrum(alms) nothing # hide ``` @@ -232,7 +248,7 @@ is correctly applied across dimensions of `A` and then convert to a Awesome. For higher degrees and orders the amplitude clearly decreases! Now to grid-point space and let us visualize the result ```@example speedytransforms -map = gridded(alms) +map = transform(alms) using CairoMakie heatmap(map, title="k⁻²-distributed noise") @@ -272,6 +288,19 @@ to kilobytes SpectralTransform(spectral_grid, recompute_legendre=true) ``` +## Batched Transforms + +SpeedyTransforms also supports batched transforms. With batched input data the `transform` is performed along the leading dimension, and all further dimensions are interpreted as batch dimensions. Take for example + +```@example speedytransforms +alms = randn(LowerTriangularMatrix{Complex{Float32}}, 32, 32, 5) +grids = transform(alms) +``` + +In this case we first randomly generated five (32x32) `LowerTriangularMatrix` that hold the +coefficients and then transformed all five matrices batched to the grid space with the +transform command, yielding 5 `RingGrids` with each 48-rings. + ## Functions and type index ```@autodocs diff --git a/docs/src/structure.md b/docs/src/structure.md index 00575a91c..dd1cce9c3 100644 --- a/docs/src/structure.md +++ b/docs/src/structure.md @@ -33,7 +33,7 @@ But let's start at the top. When creating a `Simulation`, its fields are ```@example structure -spectral_grid = SpectralGrid(nlev = 1) +spectral_grid = SpectralGrid(nlayers = 1) model = BarotropicModel(; spectral_grid) simulation = initialize!(model) ``` @@ -55,6 +55,8 @@ use temperature for example (but you could use nevertheless). tree(simulation.prognostic_variables) ``` +The prognostic variable struct can be mutated (e.g. to set new initial conditions) with the [`SpeedyWeather.set!`](@ref) function. + ## Diagnostic variables Similar for the diagnostic variables, regardless the model type, they contain the @@ -126,7 +128,7 @@ And with `max_level` you can truncate the tree to go down at most that many leve a higher resolution `PrimitiveWetModel` would use ```@example structure -spectral_grid = SpectralGrid(trunc=127, nlev=8) +spectral_grid = SpectralGrid(trunc=127, nlayers=8) model = PrimitiveWetModel(;spectral_grid) simulation = initialize!(model) tree(simulation, max_level=1, with_size=true) diff --git a/src/LowerTriangularMatrices/LowerTriangularMatrices.jl b/src/LowerTriangularMatrices/LowerTriangularMatrices.jl index c7c176876..e1144b9eb 100644 --- a/src/LowerTriangularMatrices/LowerTriangularMatrices.jl +++ b/src/LowerTriangularMatrices/LowerTriangularMatrices.jl @@ -12,7 +12,6 @@ import LinearAlgebra: tril! # VISUALISATION import UnicodePlots -# export plot export LowerTriangularMatrix, LowerTriangularArray export eachharmonic, eachmatrix diff --git a/src/LowerTriangularMatrices/lower_triangular_matrix.jl b/src/LowerTriangularMatrices/lower_triangular_matrix.jl index 1f107f720..79f3b002f 100644 --- a/src/LowerTriangularMatrices/lower_triangular_matrix.jl +++ b/src/LowerTriangularMatrices/lower_triangular_matrix.jl @@ -231,7 +231,7 @@ Base.@propagate_inbounds Base.getindex(L::LowerTriangularArray{T,N}, i::Integer) Base.@propagate_inbounds Base.getindex(L::LowerTriangularArray{T,1,V}, i::Integer) where {T,V<:AbstractVector{T}} = getindex(L.data, i) Base.@propagate_inbounds Base.getindex(L::LowerTriangularArray{T,1,V}, I::CartesianIndex{M}) where {T,V<:AbstractVector{T},M} = getindex(L, Tuple(I)...) Base.@propagate_inbounds Base.getindex(L::LowerTriangularArray{T,1,V}, i::Integer, I::CartesianIndex{0}) where {T,V<:AbstractVector{T}} = getindex(L, i) - +Base.@propagate_inbounds Base.getindex(L::LowerTriangularArray{T,1,V}, i::Integer, I::CartesianIndices{0}) where {T,V<:AbstractVector{T}} = getindex(L, i) # setindex with lm, .. @inline Base.setindex!(L::LowerTriangularArray{T,N}, x, I::Vararg{Any, N}) where {T, N} = setindex!(L.data, x, I...) @@ -320,7 +320,7 @@ function lowertriangular_match(L1::LowerTriangularArray, Ls::LowerTriangularArra end """$(TYPEDSIGNATURES) -Returns a tuple like (1,2,3) as string "1×2×3". To be used with size2x_string(size()""" +Returns a tuple like `(1,2,3)` as string "1×2×3". To be used with `size2x_string(size(a))` with `a` some array.""" function size2x_string(t::Tuple) s = "$(t[1])" for i in t[2:end] @@ -401,7 +401,7 @@ function Base.copyto!( Base.OneTo(minimum(size.((L1, L2), 2; as=Matrix)))) L1.data .= convert.(T, L2.data) - L1 + return L1 end # CPU version @@ -427,7 +427,7 @@ function Base.copyto!( end end - L1 + return L1 end # Fallback / GPU version (the two versions _copyto! and copyto! are there to enable tests of this function with regular Arrays) @@ -482,19 +482,21 @@ function _copyto_core!( L1.data[ind_L1,[Colon() for i=1:(N-1)]...] = T.(L2.data[ind_L2,[Colon() for i=1:(N-1)]...]) - L1 + return L1 end + +# copyto! using matrix indexing from Matrix/Array function Base.copyto!( L::LowerTriangularArray{T}, # copy to L M::AbstractArray) where T # copy from M @boundscheck size(L, as=Matrix) == size(M) || throw(BoundsError) L.data .= convert.(T, M[lowertriangle_indices(M)]) - - L + return L end function Base.copyto!( M::AbstractArray{T}, # copy to M L::LowerTriangularArray) where T # copy from L + @boundscheck size(L, as=Matrix) == size(M) || throw(BoundsError) lower_triangle_indices = lowertriangle_indices(M) @@ -503,17 +505,15 @@ function Base.copyto!( M::AbstractArray{T}, # copy to M M[upper_triangle_indices] .= zero(T) M[lower_triangle_indices] = convert.(T, L.data) - M + return M end -# copyto! from Vector to LA +# copyto! from Vector/Array to using vector indexing function Base.copyto!( L::LowerTriangularArray{T,N}, # copy to L V::AbstractArray{S,N}) where {T,S,N}# copy from V @boundscheck size(L, as=Vector) == size(V) || throw(BoundsError) - L.data .= convert.(T, V) - - L + return L end function LowerTriangularMatrix{T}(M::LowerTriangularMatrix{T2}) where {T,T2} @@ -568,6 +568,8 @@ Base.isapprox(L1::LowerTriangularArray, L2::LowerTriangularArray; kwargs...) = Base.all(L::LowerTriangularArray) = all(L.data) Base.any(L::LowerTriangularArray) = any(L.data) +Base.repeat(L::LowerTriangularArray, counts...) = LowerTriangularArray(repeat(L.data, counts...), L.m, L.n) + # Broadcast CPU/GPU import Base.Broadcast: BroadcastStyle, Broadcasted, DefaultArrayStyle import LinearAlgebra: isstructurepreserving, fzeropreserving diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 993201d1a..984dce256 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -223,6 +223,13 @@ function (::Type{Grid})( return Grid(Array{Float64}(undef, get_npoints2D(Grid, nlat_half), k...), nlat_half) end +function Base.convert( + ::Type{Grid}, + grid::AbstractGridArray, +) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} + return Grid(ArrayType(grid.data)) +end + ## COORDINATES """$(TYPEDSIGNATURES) Latitudes (in degrees, -90˚-90˚N) and longitudes (0-360˚E) for @@ -300,6 +307,12 @@ for k in eachgrid(grid) grid[ij, k]""" @inline eachgrid(grid::AbstractGridArray) = CartesianIndices(size(grid)[2:end]) +# several arguments to check for matching grids +function eachgrid(grid1::AbstractGridArray, grids::AbstractGridArray...; kwargs...) + grids_match(grid1, grids...; kwargs...) || throw(DimensionMismatch(grid1, grids...)) + return eachgrid(grid1) +end + """ $(TYPEDSIGNATURES) Vector{UnitRange} `rings` to loop over every ring of grid `grid` @@ -324,40 +337,81 @@ function eachring(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) end """$(TYPEDSIGNATURES) Same as `eachring(grid)` but performs a bounds check to assess -that all `grids` are of same size.""" -function eachring(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArray} - n = length(grid1) - Base._all_match_first(X->length(X), n, grid1, grids...) || throw(BoundsError) +that all `grids` according to `grids_match` (non-parametric grid type, nlat_half and length).""" +function eachring(grid1::AbstractGridArray, grids::AbstractGridArray...) + # for eachring grids only need to match in the horizontal, but can different vertical (or other) dimensions + grids_match(grid1, grids...; horizontal_only = true) || throw(DimensionMismatch(grid1, grids...)) return eachring(grid1) end +function Base.DimensionMismatch(grid1::AbstractGridArray, grids::AbstractGridArray...) + s = "AbstractGridArrays do not match; $(size(grid1)) $(nonparametric_type(grid1))" + for grid in grids + s *= ", $(size(grid))-$(nonparametric_type(grid))" + end + return DimensionMismatch(s) +end + # equality and comparison, somehow needed as not covered by broadcasting Base.:(==)(G1::AbstractGridArray, G2::AbstractGridArray) = grids_match(G1, G2) && G1.data == G2.data Base.all(G::AbstractGridArray) = all(G.data) Base.any(G::AbstractGridArray) = any(G.data) -"""$(TYPEDSIGNATURES) True if both `A` and `B` are of the same type -(regardless type parameter `T` or underyling array type `ArrayType`) and -of same size.""" -function grids_match(A::AbstractGridArray, B::AbstractGridArray) - length(A) == length(B) && return grids_match(typeof(A), typeof(B)) - return false +"""$(TYPEDSIGNATURES) True if both `A` and `B` are of the same nonparametric grid type +(e.g. OctahedralGaussianArray, regardless type parameter `T` or underyling array type `ArrayType`) +and of same resolution (`nlat_half`) and total grid points (`length`). Sizes of `(4,)` and `(4,1)` +would match for example, but `(8,1)` and `(4,2)` would not (`nlat_half` not identical).""" +function grids_match( + A::AbstractGridArray, + B::AbstractGridArray; + horizontal_only::Bool = false, + vertical_only::Bool = false, +) + @assert ~(horizontal_only && vertical_only) "Conflicting options: horizontal_only = $horizontal_ony and vertical_only = $vertical_only" + + horizontal_match = get_nlat_half(A) == get_nlat_half(B) + vertical_match = size(A)[2:end] == size(B)[2:end] + type_match = grids_match(typeof(A), typeof(B)) + + if horizontal_only + # type also has to match as two different grid types can have the same nlat_half + return horizontal_match && type_match + elseif vertical_only + return vertical_match + else + return horizontal_match && vertical_match && type_match + end end +# eltypes can be different and also array types of underlying data function grids_match(A::Type{<:AbstractGridArray}, B::Type{<:AbstractGridArray}) - # eltypes can be different and also array types of underlying data return nonparametric_type(A) == nonparametric_type(B) end +"""$(TYPEDSIGNATURES) True if all grids `A, B, C, ...` provided as arguments +match according to `grids_match` wrt to `A` (and therefore all).""" +function grids_match(A::AbstractGridArray, B::AbstractGridArray...; kwargs...) + match = true # single grid A always matches itself + for Bi in B # check for all matching respectively with A + match &= grids_match(A, Bi; kwargs...) + end + return match +end + """$(TYPEDSIGNATURES) UnitRange to access data on grid `grid` on ring `j`.""" function each_index_in_ring(grid::Grid, j::Integer) where {Grid<:AbstractGridArray} return each_index_in_ring(Grid, j, grid.nlat_half) end -""" $(TYPEDSIGNATURES) UnitRange to access each grid point on grid `grid`.""" +""" $(TYPEDSIGNATURES) UnitRange to access each horizontal grid point on grid `grid`. +For a `NxM` (`N` horizontal grid points, `M` vertical layers) `OneTo(N)` is returned.""" eachgridpoint(grid::AbstractGridArray) = Base.OneTo(get_npoints(grid)) + +""" $(TYPEDSIGNATURES) Like `eachgridpoint(::AbstractGridArray)` but checks for +equal size between input arguments first.""" function eachgridpoint(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArray} n = length(grid1) + # TODO check only nonparametric_type and nlat_half identical! Base._all_match_first(X->length(X), n, grid1, grids...) || throw(BoundsError) return eachgridpoint(grid1) end diff --git a/src/RingGrids/interpolation.jl b/src/RingGrids/interpolation.jl index 06480b492..e27674344 100644 --- a/src/RingGrids/interpolation.jl +++ b/src/RingGrids/interpolation.jl @@ -141,7 +141,7 @@ struct AnvilInterpolator{NF<:AbstractFloat, G<:AbstractGrid} <: AbstractInterpol locator::AnvilLocator{NF} end -function Base.show(io::IO,L::AnvilInterpolator{NF,Grid}) where {NF,Grid} +function Base.show(io::IO,L::AnvilInterpolator{NF, Grid}) where {NF, Grid} println(io,"$(typeof(L))") println(io,"├ from: $(L.geometry.nlat)-ring $Grid, $(L.geometry.npoints) grid points") print(io,"└ onto: $(L.locator.npoints) points") @@ -151,39 +151,40 @@ end Locator(::Type{<:AnvilInterpolator}) = AnvilLocator function (::Type{I})( ::Type{NF}, - ::Type{Grid}, - nlat_half::Integer, # size of input grid - npoints::Integer # number of points to interpolate onto - ) where {I<:AbstractInterpolator, NF<:AbstractFloat, Grid<:AbstractGrid} + ::Type{Grid}, # 2D or 3D+ + nlat_half::Integer, # size of input grid + npoints::Integer # number of points to interpolate onto + ) where {I<:AbstractInterpolator, NF<:AbstractFloat, Grid<:AbstractGridArray} + Grid2D = horizontal_grid_type(Grid) Loc = Locator(I) # L is the to Interpolator I corresponding locator - geometry = GridGeometry(Grid, nlat_half) # general coordinates and indices for grid - locator = Loc(NF, npoints) # preallocate work arrays for interpolation - return I{NF, Grid}(geometry, locator) # assemble geometry and locator to interpolator + geometry = GridGeometry(Grid2D, nlat_half) # general coordinates and indices for grid + locator = Loc(NF, npoints) # preallocate work arrays for interpolation + return I{NF, Grid2D}(geometry, locator) # assemble geometry and locator to interpolator end function (::Type{I})( ::Type{Grid}, nlat_half::Integer, npoints::Integer, - ) where {I<:AbstractInterpolator, Grid<:AbstractGrid} + ) where {I<:AbstractInterpolator, Grid<:AbstractGridArray} return I(Float64, Grid, nlat_half, npoints) end const DEFAULT_INTERPOLATOR = AnvilInterpolator function interpolator( ::Type{NF}, - Aout::AbstractGrid, - A::AbstractGrid, + Aout::AbstractGridArray, + A::AbstractGridArray, Interpolator::Type{<:AbstractInterpolator}=DEFAULT_INTERPOLATOR ) where {NF<:AbstractFloat} latds, londs = get_latdlonds(Aout) # coordinates of new grid - I = Interpolator(NF, typeof(A), A.nlat_half, length(Aout)) + I = Interpolator(NF, typeof(A), A.nlat_half, get_npoints2D(Aout)) update_locator!(I, latds, londs, unsafe=false) return I end -function interpolator( Aout::AbstractGrid, - A::AbstractGrid, +function interpolator( Aout::AbstractGridArray, + A::AbstractGridArray, Interpolator::Type{<:AbstractInterpolator}=DEFAULT_INTERPOLATOR) return interpolator(Float64, Aout, A, Interpolator) # use Float64 as default end @@ -194,9 +195,9 @@ function interpolate(latd::Real, lond::Real, A::AbstractGrid) return Ai[1] end -function interpolate( latds::Vector{NF}, # latitudes to interpolate onto (90˚N...-90˚N) - londs::Vector{NF}, # longitudes to interpolate into (0˚...360˚E) - A::AbstractGrid, # gridded field to interpolate from +function interpolate( latds::AbstractVector{NF}, # latitudes to interpolate onto (90˚N...-90˚N) + londs::AbstractVector{NF}, # longitudes to interpolate into (0˚...360˚E) + A::AbstractGrid, # gridded field to interpolate from Interpolator::Type{<:AbstractInterpolator}=DEFAULT_INTERPOLATOR, ) where NF # number format used for interpolation n = length(latds) @@ -217,18 +218,18 @@ function interpolate( A::AbstractGrid{NF}, # field to interpolate interpolate!(Aout, A, I) # perform interpolation, store in As end -function interpolate!( Aout::Vector, # Out: interpolated values - A::Grid, # gridded values to interpolate from - interpolator::AnvilInterpolator{NF, Grid}, # geometry info and work arrays - ) where {NF<:AbstractFloat, Grid<:AbstractGrid} - +function interpolate!( + Aout::AbstractVector, # Out: interpolated values + A::AbstractVector, # gridded values to interpolate from + interpolator::AnvilInterpolator, # geometry info and work arrays +) (; ij_as, ij_bs, ij_cs, ij_ds, Δabs, Δcds, Δys ) = interpolator.locator (; npoints ) = interpolator.geometry # 1) Aout's length must match the interpolator - # 2) input grid A must match the interpolator's geometry (Grids are checked with dispatch) + # 2) input A must match the interpolator's geometry points (do not check grids for view support) @boundscheck length(Aout) == length(ij_as) || throw(BoundsError) - @boundscheck A.nlat_half == interpolator.geometry.nlat_half || throw(BoundsError) + @boundscheck length(A) == npoints || throw(BoundsError) A_northpole, A_southpole = average_on_poles(A, interpolator.geometry.rings) @@ -251,14 +252,29 @@ function interpolate!( Aout::Vector, # Out: interpolated values return Aout end -function interpolate!( Aout::AbstractGrid, # Out: grid to interpolate onto - A::AbstractGrid, # In: gridded data to interpolate from - interpolator::AbstractInterpolator) - - grids_match(Aout, A) && return copyto!(Aout.data, A.data) # if grids match just copy data over (eltypes might differ) +function interpolate!( + Aout::AbstractGrid, # Out: grid to interpolate onto + A::AbstractGrid, # In: gridded data to interpolate from + interpolator::AnvilInterpolator, +) + # if grids match just copy data over (eltypes might differ) + grids_match(Aout, A) && return copyto!(Aout.data, A.data) interpolate!(Aout.data, A, interpolator) end +function interpolate!( + Aout::AbstractGridArray, # Out: grid to interpolate onto + A::AbstractGridArray, # In: gridded data to interpolate from + interpolator::AnvilInterpolator, +) + # if grids match just copy data over (eltypes might differ) + grids_match(Aout, A) && return copyto!(Aout.data, A.data) + + for k in eachgrid(Aout, A, vertical_only=true) + interpolate!(view(Aout.data, :, k), view(A.data, :, k), interpolator) + end +end + function interpolate!( ::Type{NF}, Aout::AbstractGrid, A::AbstractGrid, @@ -324,7 +340,7 @@ function find_rings!( js::Vector{<:Integer}, # Out: ring indices j @assert θmax <= 90 "Latitudes θs are expected to be within [-90˚, 90˚]; θ=$(θmax)˚ given." @assert isdecreasing(latd) "Latitudes latd are expected to be strictly decreasing." - @assert latd[1] == 90 "Latitudes latd are expected to contain 90˚C, the north pole." + @assert latd[1] == 90 "Latitudes latd are expected to contain 90˚N, the north pole." # Hack: for intervals between rings to be one-sided open [j, j+1) the last element in # latd has to be prevfloat(-90) for the <=, > comparisons @@ -429,7 +445,7 @@ $(TYPEDSIGNATURES) Computes the average at the North and South pole from a given grid `A` and it's precomputed ring indices `rings`. The North pole average is an equally weighted average of all grid points on the northern-most ring. Similar for the South pole.""" -function average_on_poles( A::AbstractGrid{NF}, +function average_on_poles( A::AbstractVector{NF}, rings::Vector{<:UnitRange{<:Integer}} ) where {NF<:AbstractFloat} diff --git a/src/RingGrids/scaling.jl b/src/RingGrids/scaling.jl index ba6f63a8a..f6a54b221 100644 --- a/src/RingGrids/scaling.jl +++ b/src/RingGrids/scaling.jl @@ -23,8 +23,9 @@ Generic latitude scaling applied to `A` in-place with latitude-like vector `v`." function _scale_lat!(grid::AbstractGridArray{T}, v::AbstractVector) where T @boundscheck get_nlat(grid) == length(v) || throw(BoundsError) + rings = eachring(grid) @inbounds for k in eachgrid(grid) - for (j, ring) in enumerate(eachring(grid)) + for (j, ring) in enumerate(rings) vj = convert(T, v[j]) for ij in ring grid[ij, k] *= vj diff --git a/src/SpeedyTransforms/SpeedyTransforms.jl b/src/SpeedyTransforms/SpeedyTransforms.jl index c69e0166d..b50c23236 100644 --- a/src/SpeedyTransforms/SpeedyTransforms.jl +++ b/src/SpeedyTransforms/SpeedyTransforms.jl @@ -18,10 +18,8 @@ const DEFAULT_GRID = FullGaussianGrid # TRANSFORM export SpectralTransform, - gridded, - gridded!, - spectral, - spectral! + transform!, + transform # ALIASING export get_nlat_half diff --git a/src/SpeedyTransforms/aliasing.jl b/src/SpeedyTransforms/aliasing.jl index 7281a3a79..49e006e95 100644 --- a/src/SpeedyTransforms/aliasing.jl +++ b/src/SpeedyTransforms/aliasing.jl @@ -27,7 +27,7 @@ function get_truncation(nlat_half::Integer, end # unpack nlat_half from provided map -get_truncation(map::AbstractGrid, dealiasing::Real=DEFAULT_DEALIASING) = +get_truncation(map::AbstractGridArray, dealiasing::Real=DEFAULT_DEALIASING) = get_truncation(map.nlat_half, dealiasing) """ diff --git a/src/SpeedyTransforms/legendrepolarray.jl b/src/SpeedyTransforms/legendrepolarray.jl index d0c0ae779..03bb6054a 100644 --- a/src/SpeedyTransforms/legendrepolarray.jl +++ b/src/SpeedyTransforms/legendrepolarray.jl @@ -1,11 +1,11 @@ """ - AssociatedLegendrePolArray{T,N,M,V} <: AbstractArray{T,N} + AssociatedLegendrePolArray{T, N, M, V} <: AbstractArray{T,N} Type that wraps around a `LowerTriangularArray{T,M,V}` but is a subtype of `AbstractArray{T,M+1}`. This enables easier use with AssociatedLegendrePolynomials.jl which otherwise couldn't use the "matrix-style" (l, m) indexing of `LowerTriangularArray`. This type however doesn't support any other operations than indexing and is purerly intended for internal purposes. -""" +$(TYPEDFIELDS)""" struct AssociatedLegendrePolArray{T, N, M, V} <: AbstractArray{T, N} data::LowerTriangularArray{T, M, V} end diff --git a/src/SpeedyTransforms/spectral_gradients.jl b/src/SpeedyTransforms/spectral_gradients.jl index 7edc7b337..d5889c3c6 100644 --- a/src/SpeedyTransforms/spectral_gradients.jl +++ b/src/SpeedyTransforms/spectral_gradients.jl @@ -4,44 +4,47 @@ const DEFAULT_RADIUS = 1 $(TYPEDSIGNATURES) Curl of a vector `u, v` written into `curl`, `curl = ∇×(u, v)`. `u, v` are expected to have a 1/coslat-scaling included, otherwise `curl` is scaled. -Acts on the unit sphere, i.e. it omits 1/radius scaling as all inplace gradient operators. -`flipsign` option calculates -∇×(u, v) instead. `add` option calculates `curl += ∇×(u, v)` instead. -`flipsign` and `add` can be combined. This functions only creates the kernel and calls the generic -divergence function _divergence! subsequently with flipped u, v -> v, u for the curl.""" -function curl!( curl::LowerTriangularMatrix, - u::LowerTriangularMatrix, - v::LowerTriangularMatrix, - S::SpectralTransform; - flipsign::Bool=false, - add::Bool=false, - ) - +Acts on the unit sphere, i.e. it omits 1/radius scaling as all gradient operators +unless the `radius` keyword argument is provided. `flipsign` option calculates -∇×(u, v) instead. +`add` option calculates `curl += ∇×(u, v)` instead. `flipsign` and `add` can be combined. +This functions only creates the kernel and calls the generic divergence function _divergence! +subsequently with flipped u, v -> v, u for the curl.""" +function curl!( + curl::LowerTriangularArray, + u::LowerTriangularArray, + v::LowerTriangularArray, + S::SpectralTransform; + flipsign::Bool=false, + add::Bool=false, + kwargs..., +) # = -(∂λ - ∂θ) or (∂λ - ∂θ), adding or overwriting the output curl kernel(o, a, b, c) = flipsign ? (add ? o-(a+b-c) : -(a+b-c)) : (add ? o+(a+b-c) : a+b-c ) - _divergence!(kernel, curl, v, u, S) # flip u, v -> v, u + _divergence!(kernel, curl, v, u, S; kwargs...) # flip u, v -> v, u end """ $(TYPEDSIGNATURES) Divergence of a vector `u, v` written into `div`, `div = ∇⋅(u, v)`. `u, v` are expected to have a 1/coslat-scaling included, otherwise `div` is scaled. -Acts on the unit sphere, i.e. it omits 1/radius scaling as all inplace gradient operators. -`flipsign` option calculates -∇⋅(u, v) instead. `add` option calculates `div += ∇⋅(u, v)` instead. -`flipsign` and `add` can be combined. This functions only creates the kernel and calls -the generic divergence function _divergence! subsequently.""" -function divergence!( div::LowerTriangularMatrix, - u::LowerTriangularMatrix, - v::LowerTriangularMatrix, - S::SpectralTransform; - flipsign::Bool=false, - add::Bool=false, - ) - +Acts on the unit sphere, i.e. it omits 1/radius scaling as all gradient operators, +unless the `radius` keyword argument is provided. `flipsign` option calculates -∇⋅(u, v) instead. +`add` option calculates `div += ∇⋅(u, v)` instead. `flipsign` and `add` can be combined. +This functions only creates the kernel and calls the generic divergence function _divergence! subsequently.""" +function divergence!( + div::LowerTriangularArray, + u::LowerTriangularArray, + v::LowerTriangularArray, + S::SpectralTransform; + flipsign::Bool=false, + add::Bool=false, + kwargs..., +) # = -(∂λ + ∂θ) or (∂λ + ∂θ), adding or overwriting the output div kernel(o, a, b, c) = flipsign ? (add ? o-(a-b+c) : -(a-b+c)) : (add ? o+(a-b+c) : a-b+c ) - _divergence!(kernel, div, u, v, S) + _divergence!(kernel, div, u, v, S; kwargs...) end """ @@ -49,48 +52,53 @@ $(TYPEDSIGNATURES) Generic divergence function of vector `u`, `v` that writes into the output into `div`. Generic as it uses the kernel `kernel` such that curl, div, add or flipsign options are provided through `kernel`, but otherwise a single function is used. -Acts on the unit sphere, i.e. it omits 1/radius scaling as all inplace gradient operators. -""" -function _divergence!( kernel, - div::LowerTriangularMatrix{Complex{NF}}, - u::LowerTriangularMatrix{Complex{NF}}, - v::LowerTriangularMatrix{Complex{NF}}, - S::SpectralTransform{NF} - ) where {NF<:AbstractFloat} - - @boundscheck size(u) == size(div) || throw(BoundsError) - @boundscheck size(v) == size(div) || throw(BoundsError) - +Acts on the unit sphere, i.e. it omits 1/radius scaling as all gradient operators, +unless the `radius` keyword argument is provided.""" +function _divergence!( + kernel, + div::LowerTriangularArray, + u::LowerTriangularArray, + v::LowerTriangularArray, + S::SpectralTransform; + radius = DEFAULT_RADIUS, +) (; grad_y_vordiv1, grad_y_vordiv2 ) = S - @boundscheck size(grad_y_vordiv1) == size(div) || throw(BoundsError) - @boundscheck size(grad_y_vordiv2) == size(div) || throw(BoundsError) - lmax, mmax = size(div, as=Matrix) .- (2, 1) # 0-based lmax, mmax + + @boundscheck ismatching(S, div) || throw(DimensionMismatch(S, div)) + lmax, mmax = size(div, OneBased, as=Matrix) - lm = 0 - @inbounds for m in 1:mmax+1 # 1-based l, m - - # DIAGONAL (separate to avoid access to v[l-1, m]) - lm += 1 - ∂u∂λ = ((m-1)*im)*u[lm] - ∂v∂θ1 = zero(Complex{NF}) # always above the diagonal - ∂v∂θ2 = grad_y_vordiv2[lm]*v[lm+1] - div[lm] = kernel(div[lm], ∂u∂λ, ∂v∂θ1, ∂v∂θ2) - - # BELOW DIAGONAL (but skip last row) - for l in m+1:lmax+1 + for k in eachmatrix(div, u, v) # also checks size compatibility + lm = 0 + @inbounds for m in 1:mmax # 1-based l, m + + # DIAGONAL (separate to avoid access to v[l-1, m]) + lm += 1 + ∂u∂λ = ((m-1)*im)*u[lm, k] + ∂v∂θ1 = 0 # always above the diagonal + ∂v∂θ2 = grad_y_vordiv2[lm] * v[lm+1, k] + div[lm, k] = kernel(div[lm, k], ∂u∂λ, ∂v∂θ1, ∂v∂θ2) + + # BELOW DIAGONAL (but skip last row) + for l in m+1:lmax-1 + lm += 1 + ∂u∂λ = ((m-1)*im)*u[lm, k] + ∂v∂θ1 = grad_y_vordiv1[lm] * v[lm-1, k] + ∂v∂θ2 = grad_y_vordiv2[lm] * v[lm+1, k] # this pulls in data from the last row though + div[lm, k] = kernel(div[lm, k], ∂u∂λ, ∂v∂θ1, ∂v∂θ2) + end + + # Last row, only vectors make use of the lmax+1 row, set to zero for scalars div, curl lm += 1 - ∂u∂λ = ((m-1)*im)*u[lm] - ∂v∂θ1 = grad_y_vordiv1[lm]*v[lm-1] - ∂v∂θ2 = grad_y_vordiv2[lm]*v[lm+1] # this pulls in data from the last row though - div[lm] = kernel(div[lm], ∂u∂λ, ∂v∂θ1, ∂v∂θ2) + div[lm, k] = 0 end + end - # Last row, only vectors make use of the lmax+1 row, set to zero for scalars div, curl - lm += 1 - div[lm] = zero(Complex{NF}) + # /radius scaling if not unit sphere + if radius != 1 + div .*= inv(radius) end - return nothing + return div end """ @@ -104,26 +112,26 @@ An example usage is therefore RingGrids.scale_coslat⁻¹!(u_grid) RingGrids.scale_coslat⁻¹!(v_grid) - u = spectral(u_grid, one_more_degree=true) - v = spectral(v_grid, one_more_degree=true) + u = transform(u_grid, one_more_degree=true) + v = transform(v_grid, one_more_degree=true) div = divergence(u, v, radius = 6.371e6) - div_grid = gridded(div) + div_grid = transform(div) """ -function divergence(u::LowerTriangularMatrix, - v::LowerTriangularMatrix; - radius = DEFAULT_RADIUS) - - @assert size(u) == size(v) "Size $(size(u)) and $(size(v)) incompatible." +function divergence(u::LowerTriangularArray, + v::LowerTriangularArray; + kwargs...) S = SpectralTransform(u) - div = similar(u) - divergence!(div, u, v, S, add=false, flipsign=false) - - if radius != 1 - div .*= inv(radius) - end + return divergence(u, v, S; kwargs...) +end - return div +# use SpectralTransform if provided +function divergence(u::LowerTriangularArray, + v::LowerTriangularArray, + S::SpectralTransform; + kwargs...) + div = similar(u) + return divergence!(div, u, v, S; add=false, flipsign=false, kwargs...) end # called by divergence or curl @@ -131,10 +139,8 @@ function _div_or_curl( kernel!, u::Grid, v::Grid; - radius = DEFAULT_RADIUS -) where {Grid<:AbstractGrid} - - @assert size(u) == size(v) "Size $(size(u)) and $(size(v)) incompatible." + kwargs..., +) where {Grid<:AbstractGridArray} u_grid = copy(u) v_grid = copy(v) @@ -143,16 +149,11 @@ function _div_or_curl( RingGrids.scale_coslat⁻¹!(v_grid) S = SpectralTransform(u_grid, one_more_degree=true) - us = spectral(u_grid, S) - vs = spectral(v_grid, S) + us = transform(u_grid, S) + vs = transform(v_grid, S) div_or_vor = similar(us) - kernel!(div_or_vor, us, vs, S, add=false, flipsign=false) - - if radius != 1 - div_or_vor .*= inv(radius) - end - + kernel!(div_or_vor, us, vs, S; add=false, flipsign=false, kwargs...) return div_or_vor end @@ -163,7 +164,8 @@ Applies 1/coslat scaling, transforms to spectral space and returns the spectral divergence. Acts on the unit sphere, i.e. it omits 1/radius scaling unless `radius` keyword argument is provided. """ -divergence(u::Grid, v::Grid; kwargs...) where {Grid<:AbstractGrid} = _div_or_curl(divergence!, u, v; kwargs...) +divergence(u::Grid, v::Grid; kwargs...) where {Grid<:AbstractGridArray} = + _div_or_curl(divergence!, u, v; kwargs...) """ $(TYPEDSIGNATURES) @@ -171,7 +173,8 @@ Curl (∇×) of two vector components `u, v` on a grid. Applies 1/coslat scaling, transforms to spectral space and returns the spectral curl. Acts on the unit sphere, i.e. it omits 1/radius scaling unless `radius` keyword argument is provided.""" -curl(u::Grid, v::Grid; kwargs...) where {Grid<:AbstractGrid} = _div_or_curl(curl!, u, v; kwargs...) +curl(u::Grid, v::Grid; kwargs...) where {Grid<:AbstractGridArray} = + _div_or_curl(curl!, u, v; kwargs...) """ $(TYPEDSIGNATURES) @@ -183,25 +186,26 @@ requires both `u, v` to be transforms of fields that are scaled with RingGrids.scale_coslat⁻¹!(u_grid) RingGrids.scale_coslat⁻¹!(v_grid) - u = spectral(u_grid) - v = spectral(v_grid) + u = transform(u_grid) + v = transform(v_grid) vor = curl(u, v, radius=6.371e6) - vor_grid = gridded(div) + vor_grid = transform(div) """ -function curl( u::LowerTriangularMatrix, - v::LowerTriangularMatrix; - radius = DEFAULT_RADIUS) - - @assert size(u) == size(v) "Size $(size(u)) and $(size(v)) incompatible." +function curl( u::LowerTriangularArray, + v::LowerTriangularArray; + kwargs...) S = SpectralTransform(u) - vor = similar(u) - curl!(vor, u, v, S, add=false, flipsign=false) - - if radius != 1 - vor .*= inv(radius) - end + return curl(u, v, S; kwargs...) +end +# use SpectralTransform if provided +function curl( u::LowerTriangularArray, + v::LowerTriangularArray, + S::SpectralTransform; + kwargs...) + vor = similar(u) + curl!(vor, u, v, S; add=false, flipsign=false, kwargs...) return vor end @@ -211,67 +215,72 @@ Get U, V (=(u, v)*coslat) from vorticity ζ spectral space (divergence D=0) Two operations are combined into a single linear operation. First, invert the spherical Laplace ∇² operator to get stream function from vorticity. Then compute zonal and meridional gradients to get U, V. -Acts on the unit sphere, i.e. it omits any radius scaling as all inplace gradient operators. -""" -function UV_from_vor!( U::LowerTriangularMatrix{Complex{NF}}, - V::LowerTriangularMatrix{Complex{NF}}, - vor::LowerTriangularMatrix{Complex{NF}}, - S::SpectralTransform{NF} - ) where {NF<:AbstractFloat} - +Acts on the unit sphere, i.e. it omits any radius scaling as all inplace gradient operators, +unless the `radius` keyword argument is provided.""" +function UV_from_vor!( + U::LowerTriangularArray, + V::LowerTriangularArray, + vor::LowerTriangularArray, + S::SpectralTransform; + radius = DEFAULT_RADIUS, +) (; vordiv_to_uv_x, vordiv_to_uv1, vordiv_to_uv2 ) = S - lmax, mmax = size(vor, as=Matrix) .- (2, 1) # 0-based lmax, mmax + @boundscheck ismatching(S, U) || throw(DimensionMismatch(S, U)) - @boundscheck lmax == mmax || throw(BoundsError) - @boundscheck size(U) == size(vor) || throw(BoundsError) - @boundscheck size(V) == size(vor) || throw(BoundsError) - @boundscheck size(vordiv_to_uv_x) == size(vor) || throw(BoundsError) - @boundscheck size(vordiv_to_uv1) == size(vor) || throw(BoundsError) - @boundscheck size(vordiv_to_uv2) == size(vor) || throw(BoundsError) + # maximum degree l, order m of spherical harmonics (1-based) + lmax, mmax = size(U, OneBased, as=Matrix) + + for k in eachmatrix(U, V, vor) # also checks size compatibility + lm = 0 + @inbounds for m in 1:mmax-1 # 1-based l, m, exclude last column - lm = 0 - @inbounds for m in 1:mmax # 1-based l, m, exclude last column + # DIAGONAL (separated to avoid access to l-1, m which is above the diagonal) + lm += 1 - # DIAGONAL (separated to avoid access to l-1, m which is above the diagonal) - lm += 1 + # U = -∂/∂lat(Ψ) and V = V = ∂/∂λ(Ψ) combined with Laplace inversion ∇⁻², omit radius R scaling + U[lm, k] = vordiv_to_uv2[lm] * vor[lm+1, k] # - vordiv_to_uv1[lm]*vor[l-1, m] <- is zero + V[lm, k] = im*vordiv_to_uv_x[lm] * vor[lm, k] - # U = -∂/∂lat(Ψ) and V = V = ∂/∂λ(Ψ) combined with Laplace inversion ∇⁻², omit radius R scaling - U[lm] = vordiv_to_uv2[lm]*vor[lm+1] # - vordiv_to_uv1[lm]*vor[l-1, m] <- is zero - V[lm] = im*vordiv_to_uv_x[lm]*vor[lm] + # BELOW DIAGONAL + for l in m+1:lmax-2 # skip last two rows + lm += 1 - # BELOW DIAGONAL - for l in m+1:lmax # skip last two rows + # U = -∂/∂lat(Ψ) and V = V = ∂/∂λ(Ψ) combined with Laplace inversion ∇⁻², omit radius R scaling + # U[lm] = vordiv_to_uv2[lm]*vor[lm+1] - vordiv_to_uv1[lm]*vor[lm-1] + U[lm, k] = muladd(vordiv_to_uv2[lm], vor[lm+1, k], -vordiv_to_uv1[lm]*vor[lm-1, k]) + V[lm, k] = im*vordiv_to_uv_x[lm] * vor[lm, k] + end + + # SECOND LAST ROW lm += 1 + U[lm, k] = -vordiv_to_uv1[lm] * vor[lm-1, k] # meridional gradient again (but only 2nd term from above) + V[lm, k] = im*vordiv_to_uv_x[lm] * vor[lm, k] # zonal gradient again (as above) - # U = -∂/∂lat(Ψ) and V = V = ∂/∂λ(Ψ) combined with Laplace inversion ∇⁻², omit radius R scaling - # U[lm] = vordiv_to_uv2[lm]*vor[lm+1] - vordiv_to_uv1[lm]*vor[lm-1] - U[lm] = muladd(vordiv_to_uv2[lm], vor[lm+1], -vordiv_to_uv1[lm]*vor[lm-1]) - V[lm] = im*vordiv_to_uv_x[lm]*vor[lm] + # LAST ROW (separated to avoid out-of-bounds access to l+2, m) + lm += 1 + U[lm, k] = -vordiv_to_uv1[lm] * vor[lm-1] # meridional gradient again (but only 2nd term from above) + V[lm, k] = 0 # set explicitly to 0 as Ψ does not contribute to last row of V end - # SECOND LAST ROW - lm += 1 - U[lm] = -vordiv_to_uv1[lm]*vor[lm-1] # meridional gradient again (but only 2nd term from above) - V[lm] = im*vordiv_to_uv_x[lm]*vor[lm] # zonal gradient again (as above) + # LAST COLUMN + @inbounds begin + lm += 1 # second last row + U[lm, k] = 0 + V[lm, k] = im*vordiv_to_uv_x[lm] * vor[lm, k] - # LAST ROW (separated to avoid out-of-bounds access to l+2, m) - lm += 1 - U[lm] = -vordiv_to_uv1[lm]*vor[lm-1] # meridional gradient again (but only 2nd term from above) - V[lm] = zero(Complex{NF}) # set explicitly to 0 as Ψ does not contribute to last row of V + lm += 1 # last row + U[lm, k] = -vordiv_to_uv1[lm] * vor[lm-1, k] + V[lm, k] = 0 + end end - # LAST COLUMN - @inbounds begin - lm += 1 # second last row - U[lm] = zero(Complex{NF}) - V[lm] = im*vordiv_to_uv_x[lm]*vor[lm] - - lm += 1 # last row - U[lm] = -vordiv_to_uv1[lm]*vor[lm-1] - V[lm] = zero(Complex{NF}) + # *radius scaling if not unit sphere (*radius² for ∇⁻² then /radius to get from stream function to velocity) + if radius != 1 + U .*= radius + V .*= radius end - return nothing + return U, V end """ @@ -283,81 +292,88 @@ velocity potential from divergence. Then compute zonal and meridional gradients to get U, V. Acts on the unit sphere, i.e. it omits any radius scaling as all inplace gradient operators. """ -function UV_from_vordiv!( U::LowerTriangularMatrix{Complex{NF}}, - V::LowerTriangularMatrix{Complex{NF}}, - vor::LowerTriangularMatrix{Complex{NF}}, - div::LowerTriangularMatrix{Complex{NF}}, - S::SpectralTransform{NF} - ) where {NF<:AbstractFloat} - +function UV_from_vordiv!( + U::LowerTriangularArray, + V::LowerTriangularArray, + vor::LowerTriangularArray, + div::LowerTriangularArray, + S::SpectralTransform; + radius = DEFAULT_RADIUS, +) (; vordiv_to_uv_x, vordiv_to_uv1, vordiv_to_uv2 ) = S - lmax, mmax = size(vor, as=Matrix) .- (2, 1) # 0-based lmax, mmax - @boundscheck lmax == mmax || throw(BoundsError) - @boundscheck size(div) == size(vor) || throw(BoundsError) - @boundscheck size(U) == size(vor) || throw(BoundsError) - @boundscheck size(V) == size(vor) || throw(BoundsError) - @boundscheck size(vordiv_to_uv_x) == size(vor) || throw(BoundsError) - @boundscheck size(vordiv_to_uv1) == size(vor) || throw(BoundsError) - @boundscheck size(vordiv_to_uv1) == size(vor) || throw(BoundsError) - - lm = 0 - @inbounds for m in 1:mmax # 1-based l, m, skip last column - - # DIAGONAL (separated to avoid access to l-1, m which is above the diagonal) - lm += 1 - - # div, vor contribution to meridional gradient - ∂ζθ = vordiv_to_uv2[lm]*vor[lm+1] # lm-1 term is zero - ∂Dθ = -vordiv_to_uv2[lm]*div[lm+1] # lm-1 term is zero - - # the following is moved into the muladd - # ∂Dλ = im*vordiv_to_uv_x[lm]*div[lm] # divergence contribution to zonal gradient - # ∂ζλ = im*vordiv_to_uv_x[lm]*vor[lm] # vorticity contribution to zonal gradient + @boundscheck ismatching(S, U) || throw(DimensionMismatch(S, U)) - z = im*vordiv_to_uv_x[lm] - U[lm] = muladd(z, div[lm], ∂ζθ) # = ∂Dλ + ∂ζθ - V[lm] = muladd(z, vor[lm], ∂Dθ) # = ∂ζλ + ∂Dθ + # maximum degree l, order m of spherical harmonics (1-based) + lmax, mmax = size(U, OneBased, as=Matrix) - # BELOW DIAGONAL (all terms) - for l in m+1:lmax # skip last row (lmax+2) + for k in eachmatrix(U, V, vor, div) # also checks size compatibility + lm = 0 + @inbounds for m in 1:mmax-1 # 1-based l, m, skip last column + + # DIAGONAL (separated to avoid access to l-1, m which is above the diagonal) lm += 1 # div, vor contribution to meridional gradient - # ∂ζθ = vordiv_to_uv2[lm]*vor[lm+1] - vordiv_to_uv1[lm]*vor[lm-1] - # ∂Dθ = vordiv_to_uv1[lm]*div[lm-1] - vordiv_to_uv2[lm]*div[lm+1] - ∂ζθ = muladd(vordiv_to_uv2[lm], vor[lm+1], -vordiv_to_uv1[lm]*vor[lm-1]) - ∂Dθ = muladd(vordiv_to_uv1[lm], div[lm-1], -vordiv_to_uv2[lm]*div[lm+1]) - - # The following is moved into the muladd - # ∂Dλ = im*vordiv_to_uv_x[lm]*div[lm] # divergence contribution to zonal gradient - # ∂ζλ = im*vordiv_to_uv_x[lm]*vor[lm] # vorticity contribution to zonal gradient + ∂ζθ = vordiv_to_uv2[lm]*vor[lm+1, k] # lm-1 term is zero + ∂Dθ = -vordiv_to_uv2[lm]*div[lm+1, k] # lm-1 term is zero + + # the following is moved into the muladd + # ∂Dλ = im*vordiv_to_uv_x[lm]*div[lm] # divergence contribution to zonal gradient + # ∂ζλ = im*vordiv_to_uv_x[lm]*vor[lm] # vorticity contribution to zonal gradient z = im*vordiv_to_uv_x[lm] - U[lm] = muladd(z, div[lm], ∂ζθ) # = ∂Dλ + ∂ζθ - V[lm] = muladd(z, vor[lm], ∂Dθ) # = ∂ζλ + ∂Dθ + U[lm, k] = muladd(z, div[lm, k], ∂ζθ) # = ∂Dλ + ∂ζθ + V[lm, k] = muladd(z, vor[lm, k], ∂Dθ) # = ∂ζλ + ∂Dθ + + # BELOW DIAGONAL (all terms) + for l in m+1:lmax-2 # skip last two rows (lmax-1, lmax) + lm += 1 + + # div, vor contribution to meridional gradient + # ∂ζθ = vordiv_to_uv2[lm]*vor[lm+1] - vordiv_to_uv1[lm]*vor[lm-1] + # ∂Dθ = vordiv_to_uv1[lm]*div[lm-1] - vordiv_to_uv2[lm]*div[lm+1] + ∂ζθ = muladd(vordiv_to_uv2[lm], vor[lm+1, k], -vordiv_to_uv1[lm]*vor[lm-1, k]) + ∂Dθ = muladd(vordiv_to_uv1[lm], div[lm-1, k], -vordiv_to_uv2[lm]*div[lm+1, k]) + + # The following is moved into the muladd + # ∂Dλ = im*vordiv_to_uv_x[lm]*div[lm] # divergence contribution to zonal gradient + # ∂ζλ = im*vordiv_to_uv_x[lm]*vor[lm] # vorticity contribution to zonal gradient + + z = im*vordiv_to_uv_x[lm] + U[lm, k] = muladd(z, div[lm, k], ∂ζθ) # = ∂Dλ + ∂ζθ + V[lm, k] = muladd(z, vor[lm, k], ∂Dθ) # = ∂ζλ + ∂Dθ + end + + # SECOND LAST ROW (separated to imply that vor, div are zero in last row) + lm += 1 + U[lm, k] = im*vordiv_to_uv_x[lm]*div[lm, k] - vordiv_to_uv1[lm]*vor[lm-1, k] + V[lm, k] = im*vordiv_to_uv_x[lm]*vor[lm, k] + vordiv_to_uv1[lm]*div[lm-1, k] + + # LAST ROW (separated to avoid out-of-bounds access to lmax+1) + lm += 1 + U[lm, k] = -vordiv_to_uv1[lm]*vor[lm-1, k] # only last term from 2nd last row + V[lm, k] = vordiv_to_uv1[lm]*div[lm-1, k] # only last term from 2nd last row end - # SECOND LAST ROW (separated to imply that vor, div are zero in last row) - lm += 1 - U[lm] = im*vordiv_to_uv_x[lm]*div[lm] - vordiv_to_uv1[lm]*vor[lm-1] - V[lm] = im*vordiv_to_uv_x[lm]*vor[lm] + vordiv_to_uv1[lm]*div[lm-1] + # LAST COLUMN + @inbounds begin + lm += 1 # second last row + U[lm, k] = im*vordiv_to_uv_x[lm]*div[lm, k] # other terms are zero + V[lm, k] = im*vordiv_to_uv_x[lm]*vor[lm, k] # other terms are zero - # LAST ROW (separated to avoid out-of-bounds access to lmax+3 - lm += 1 - U[lm] = -vordiv_to_uv1[lm]*vor[lm-1] # only last term from 2nd last row - V[lm] = vordiv_to_uv1[lm]*div[lm-1] # only last term from 2nd last row + lm += 1 # last row + U[lm, k] = -vordiv_to_uv1[lm]*vor[lm-1, k] # other terms are zero + V[lm, k] = vordiv_to_uv1[lm]*div[lm-1, k] # other terms are zero + end end - # LAST COLUMN - @inbounds begin - lm += 1 # second last row - U[lm] = im*vordiv_to_uv_x[lm]*div[lm] # other terms are zero - V[lm] = im*vordiv_to_uv_x[lm]*vor[lm] # other terms are zero - - lm += 1 # last row - U[lm] = -vordiv_to_uv1[lm]*vor[lm-1] # other terms are zero - V[lm] = vordiv_to_uv1[lm]*div[lm-1] # other terms are zero + # *radius scaling if not unit sphere (*radius² for ∇⁻², then /radius to get from stream function to velocity) + if radius != 1 + U .*= radius + V .*= radius end + + return U, V end """ @@ -365,7 +381,8 @@ $(TYPEDSIGNATURES) Laplace operator ∇² applied to the spectral coefficients `alms` in spherical coordinates. The eigenvalues which are precomputed in `S`. ∇²! is the in-place version which directly stores the output in the first argument `∇²alms`. -Acts on the unit sphere, i.e. it omits any radius scaling as all inplace gradient operators. +Acts on the unit sphere, i.e. it omits any radius scaling as all inplace gradient operators, +unless the `radius` keyword argument is provided. Keyword arguments ================= @@ -375,32 +392,42 @@ Keyword arguments - `inverse=true` computes ∇⁻²(alms) instead Default is `add=false`, `flipsign=false`, `inverse=false`. These options can be combined.""" -function ∇²!( ∇²alms::LowerTriangularMatrix{Complex{NF}}, # Output: (inverse) Laplacian of alms - alms::LowerTriangularMatrix{Complex{NF}}, # Input: spectral coefficients - S::SpectralTransform{NF}; # precomputed eigenvalues - add::Bool=false, # add to output array or overwrite - flipsign::Bool=false, # -∇² or ∇² - inverse::Bool=false, # ∇⁻² or ∇² - ) where {NF<:AbstractFloat} - - @boundscheck size(alms) == size(∇²alms) || throw(BoundsError) - lmax, mmax = size(alms; as=Matrix) .- (1, 1) # 0-based degree l, order m of the Legendre polynomials - +function ∇²!( + ∇²alms::LowerTriangularArray, # Output: (inverse) Laplacian of alms + alms::LowerTriangularArray, # Input: spectral coefficients + S::SpectralTransform; # precomputed eigenvalues + add::Bool=false, # add to output array or overwrite + flipsign::Bool=false, # -∇² or ∇² + inverse::Bool=false, # ∇⁻² or ∇² + radius = DEFAULT_RADIUS, # scale with radius if provided, otherwise unit sphere +) + @boundscheck ismatching(S, ∇²alms) || throw(DimensionMismatch(S, ∇²alms)) + # use eigenvalues⁻¹/eigenvalues for ∇⁻²/∇² based but name both eigenvalues eigenvalues = inverse ? S.eigenvalues⁻¹ : S.eigenvalues - @boundscheck length(eigenvalues) >= lmax+1 || throw(BoundsError) - + @inline kernel(o, a) = flipsign ? (add ? (o-a) : -a) : (add ? (o+a) : a) - - lm = 0 - @inbounds for m in 1:mmax+1 # order m = 0:mmax but 1-based - for l in m:lmax+1 # degree l = m:lmax but 1-based - lm += 1 - ∇²alms[lm] = kernel(∇²alms[lm], alms[lm]*eigenvalues[l]) + + # maximum degree l, order m of spherical harmonics (1-based) + lmax, mmax = size(alms, OneBased, as=Matrix) + + for k in eachmatrix(∇²alms, alms) + lm = 0 + @inbounds for m in 1:mmax + for l in m:lmax + lm += 1 + ∇²alms[lm, k] = kernel(∇²alms[lm, k], alms[lm, k]*eigenvalues[l]) + end end end + # /radius² or *radius² scaling if not unit sphere + if radius != 1 + R_plusminus_squared = inverse ? radius^2 : inv(radius^2) + ∇²alms .*= R_plusminus_squared + end + return ∇²alms end @@ -410,17 +437,12 @@ Laplace operator ∇² applied to input `alms`, using precomputed eigenvalues fr Acts on the unit sphere, i.e. it omits 1/radius^2 scaling unless `radius` keyword argument is provided.""" function ∇²( - alms::LowerTriangularMatrix, # Input: spectral coefficients + alms::LowerTriangularArray, # Input: spectral coefficients S::SpectralTransform; # precomputed eigenvalues - radius = DEFAULT_RADIUS, + kwargs..., ) ∇²alms = similar(alms) - ∇²!(∇²alms, alms, S, add=false, flipsign=false, inverse=false) - - if radius != 1 - ∇²alms .*= inv(radius^2) - end - + ∇²!(∇²alms, alms, S; add=false, flipsign=false, inverse=false, kwargs...) return ∇²alms end @@ -429,7 +451,7 @@ $(TYPEDSIGNATURES) Returns the Laplace operator ∇² applied to input `alms`. Acts on the unit sphere, i.e. it omits 1/radius^2 scaling unless `radius` keyword argument is provided.""" -∇²(alms::LowerTriangularMatrix; kwargs...) = ∇²(alms, SpectralTransform(alms); kwargs...) +∇²(alms::LowerTriangularArray; kwargs...) = ∇²(alms, SpectralTransform(alms); kwargs...) """ $(TYPEDSIGNATURES) @@ -437,18 +459,12 @@ InverseLaplace operator ∇⁻² applied to input `alms`, using precomputed eigenvalues from `S`. Acts on the unit sphere, i.e. it omits radius^2 scaling unless `radius` keyword argument is provided.""" function ∇⁻²( - ∇²alms::LowerTriangularMatrix, # Input: spectral coefficients + ∇²alms::LowerTriangularArray, # Input: spectral coefficients S::SpectralTransform; # precomputed eigenvalues - radius = DEFAULT_RADIUS, + kwargs..., ) - alms = similar(∇²alms) - ∇⁻²!(alms, ∇²alms, S, add=false, flipsign=false) - - if radius != 1 - ∇²alms .*= radius^2 - end - + ∇⁻²!(alms, ∇²alms, S; add=false, flipsign=false, kwargs...) return alms end @@ -457,67 +473,77 @@ $(TYPEDSIGNATURES) Returns the inverse Laplace operator ∇⁻² applied to input `alms`. Acts on the unit sphere, i.e. it omits radius^2 scaling unless `radius` keyword argument is provided.""" -∇⁻²(∇²alms::LowerTriangularMatrix; kwargs...) = ∇⁻²(∇²alms, SpectralTransform(∇²alms); kwargs...) +∇⁻²(∇²alms::LowerTriangularArray; kwargs...) = ∇⁻²(∇²alms, SpectralTransform(∇²alms); kwargs...) """$(TYPEDSIGNATURES) Calls `∇²!(∇⁻²alms, alms, S; add, flipsign, inverse=true)`.""" -function ∇⁻²!( ∇⁻²alms::LowerTriangularMatrix{Complex{NF}}, # Output: inverse Laplacian of alms - alms::LowerTriangularMatrix{Complex{NF}}, # Input: spectral coefficients - S::SpectralTransform{NF}; # precomputed eigenvalues - add::Bool=false, # add to output array or overwrite - flipsign::Bool=false, # -∇⁻² or ∇⁻² - ) where {NF<:AbstractFloat} - +function ∇⁻²!( + ∇⁻²alms::LowerTriangularArray, # Output: inverse Laplacian of alms + alms::LowerTriangularArray, # Input: spectral coefficients + S::SpectralTransform; # precomputed eigenvalues + add::Bool = false, # add to output array or overwrite + flipsign::Bool = false, # -∇⁻² or ∇⁻² + kwargs..., +) inverse = true - return ∇²!(∇⁻²alms, alms, S; add, flipsign, inverse) + return ∇²!(∇⁻²alms, alms, S; add, flipsign, inverse, kwargs...) end """$(TYPEDSIGNATURES) Applies the gradient operator ∇ applied to input `p` and stores the result in `dpdx` (zonal derivative) and `dpdy` (meridional derivative). The gradient operator acts -on the unit sphere and therefore omits the 1/radius scaling""" -function ∇!(dpdx::LowerTriangularMatrix{Complex{NF}}, # Output: zonal gradient - dpdy::LowerTriangularMatrix{Complex{NF}}, # Output: meridional gradient - p::LowerTriangularMatrix{Complex{NF}}, # Input: spectral coefficients - S::SpectralTransform{NF} # includes precomputed arrays - ) where {NF<:AbstractFloat} - - lmax, mmax = size(p, as=Matrix) .- (1, 1) # 0-based, include last row - @boundscheck size(p) == size(dpdx) || throw(BoundsError) - @boundscheck size(p) == size(dpdy) || throw(BoundsError) - - (; grad_y1, grad_y2 ) = S - - lm = 0 - @inbounds for m in 1:mmax # 1-based l, m, skip last column - - # DIAGONAL (separated to avoid access to l-1, m which is above the diagonal) - lm += 1 - - dpdx[lm] = (m-1)*im*p[lm] # zonal gradient: d/dlon = *i*m - dpdy[lm] = grad_y2[lm]*p[lm+1] # meridional gradient: p[lm-1]=0 on diagonal - - # BELOW DIAGONAL (all terms) - for l in m+1:lmax # skip last row +on the unit sphere and therefore omits the 1/radius scaling unless `radius` keyword argument is provided.""" +function ∇!( + dpdx::LowerTriangularArray, # Output: zonal gradient + dpdy::LowerTriangularArray, # Output: meridional gradient + p::LowerTriangularArray, # Input: spectral coefficients + S::SpectralTransform; # includes precomputed arrays + radius = DEFAULT_RADIUS, # scale with radius if provided, otherwise unit sphere +) + (; grad_y1, grad_y2) = S + @boundscheck ismatching(S, p) || throw(DimensionMismatch(S, p)) + + # maximum degree l, order m of spherical harmonics (1-based) + lmax, mmax = size(p, OneBased, as=Matrix) + + for k in eachmatrix(dpdx, dpdy, p) # also performs size checks + lm = 0 + @inbounds for m in 1:mmax-1 # 1-based l, m, skip last column + + # DIAGONAL (separated to avoid access to l-1, m which is above the diagonal) lm += 1 - - dpdx[lm] = (m-1)*im*p[lm] - dpdy[lm] = grad_y1[lm]*p[lm-1] + grad_y2[lm]*p[lm+1] + + dpdx[lm, k] = (m-1)*im*p[lm, k] # zonal gradient: d/dlon = *i*m + dpdy[lm, k] = grad_y2[lm]*p[lm+1, k] # meridional gradient: p[lm-1]=0 on diagonal + + # BELOW DIAGONAL (all terms) + for l in m+1:lmax-1 # skip last row + lm += 1 + dpdx[lm, k] = (m-1)*im*p[lm, k] + dpdy[lm, k] = grad_y1[lm]*p[lm-1, k] + grad_y2[lm]*p[lm+1, k] + end + + # LAST ROW (separated to avoid out-of-bounds access to lmax+1 + lm += 1 + dpdx[lm, k] = (m-1)*im*p[lm, k] + dpdy[lm, k] = grad_y1[lm]*p[lm-1, k] # only first term from 2nd last row end - # LAST ROW (separated to avoid out-of-bounds access to lmax+2 - lm += 1 - dpdx[lm] = (m-1)*im*p[lm] - dpdy[lm] = grad_y1[lm]*p[lm-1] # only first term from 2nd last row - end + # LAST COLUMN + @inbounds begin + lm += 1 # second last row + dpdx[lm, k] = (mmax-1)*im*p[lm, k] + dpdy[lm, k] = grad_y2[lm]*p[lm+1, k] # only 2nd term - # LAST COLUMN - @inbounds begin - lm += 1 - dpdx[lm] = mmax*im*p[lm] - dpdy[lm] = grad_y2[lm]*p[lm+1] # only 2nd term + lm += 1 # last row + dpdx[lm, k] = (mmax-1)*im*p[lm, k] + dpdy[lm, k] = grad_y1[lm]*p[lm-1, k] # only 1st term + end + end - lm += 1 - dpdx[lm] = mmax*im*p[lm] - dpdy[lm] = grad_y1[lm]*p[lm-1] # only 1st term + # 1/radius factor if not unit sphere + if radius != 1 + R⁻¹ = inv(radius) + dpdx .*= R⁻¹ + dpdy .*= R⁻¹ end return dpdx, dpdy @@ -526,23 +552,17 @@ end """$(TYPEDSIGNATURES) The zonal and meridional gradient of `p` using an existing `SpectralTransform` `S`. Acts on the unit sphere, i.e. it omits 1/radius scaling unless `radius` keyword argument is provided.""" -function ∇(p::LowerTriangularMatrix, S::SpectralTransform; radius = DEFAULT_RADIUS) +function ∇(p::LowerTriangularArray, S::SpectralTransform; kwargs...) dpdx = similar(p) dpdy = similar(p) - ∇!(dpdx, dpdy, p, S) - - if radius != 1 - dpdx .*= inv(radius) - dpdy .*= inv(radius) - end - + ∇!(dpdx, dpdy, p, S; kwargs...) return dpdx, dpdy end """$(TYPEDSIGNATURES) The zonal and meridional gradient of `p`. Precomputes a `SpectralTransform` `S`. Acts on the unit-sphere, i.e. it omits 1/radius scaling unless `radius` keyword argument is provided.""" -function ∇(p::LowerTriangularMatrix; kwargs...) +function ∇(p::LowerTriangularArray; kwargs...) S = SpectralTransform(p, one_more_degree=true) return ∇(p, S; kwargs...) end @@ -551,11 +571,11 @@ end Transform to spectral space, takes the gradient and unscales the 1/coslat scaling in the gradient. Acts on the unit-sphere, i.e. it omits 1/radius scaling unless `radius` keyword argument is provided. Makes use of an existing spectral transform `S`.""" -function ∇(grid::AbstractGrid, S::SpectralTransform; kwargs...) - p = spectral(grid, S) +function ∇(grid::AbstractGridArray, S::SpectralTransform; kwargs...) + p = transform(grid, S) dpdx, dpdy = ∇(p, S; kwargs...) - dpdx_grid = gridded(dpdx, S, unscale_coslat=true) - dpdy_grid = gridded(dpdy, S, unscale_coslat=true) + dpdx_grid = transform(dpdx, S, unscale_coslat=true) + dpdy_grid = transform(dpdy, S, unscale_coslat=true) return dpdx_grid, dpdy_grid end @@ -563,7 +583,7 @@ end Transform to spectral space, takes the gradient and unscales the 1/coslat scaling in the gradient. Acts on the unit-sphere, i.e. it omits 1/radius scaling unless `radius` keyword argument is provided.""" -function ∇(grid::AbstractGrid; kwargs...) +function ∇(grid::AbstractGridArray; kwargs...) S = SpectralTransform(grid, one_more_degree=true) return ∇(grid, S; kwargs...) end \ No newline at end of file diff --git a/src/SpeedyTransforms/spectral_transform.jl b/src/SpeedyTransforms/spectral_transform.jl index d3ab8145b..1d83a281d 100644 --- a/src/SpeedyTransforms/spectral_transform.jl +++ b/src/SpeedyTransforms/spectral_transform.jl @@ -1,8 +1,7 @@ """ - S = SpectralTransform{NF<:AbstractFloat}(...) - -SpectralTransform struct that contains all parameters and preallocated arrays -for the spectral transform.""" +SpectralTransform struct that contains all parameters and precomputed arrays +to perform a spectral transform. Fields are +$(TYPEDFIELDS)""" struct SpectralTransform{NF<:AbstractFloat} # GRID @@ -71,14 +70,15 @@ Generator function for a SpectralTransform struct. With `NF` the number format, `Grid` the grid type `<:AbstractGrid` and spectral truncation `lmax, mmax` 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{<:AbstractGridArray}, # type of spatial grid used - lmax::Int, # Spectral truncation: degrees - mmax::Int; # Spectral truncation: orders - 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 +function SpectralTransform( + ::Type{NF}, # Number format NF + Grid::Type{<:AbstractGridArray}, # type of spatial grid used + lmax::Int, # Spectral truncation: degrees + mmax::Int; # Spectral truncation: orders + 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 Grid = RingGrids.nonparametric_type(Grid) # always use nonparametric super type @@ -124,7 +124,8 @@ function SpectralTransform( ::Type{NF}, # Number format lon_offsets = [cispi(m*lon1/π) for m in 0:mmax, lon1 in lon1s] # PREALLOCATE LEGENDRE POLYNOMIALS, +1 for 1-based indexing - Λ = AssociatedLegendrePolArray{NF,2,1,Vector{NF}}(zeros(LowerTriangularMatrix{NF}, lmax+1, mmax+1)) # Legendre polynomials for one latitude + # Legendre polynomials for one latitude + Λ = AssociatedLegendrePolArray{NF, 2, 1, Vector{NF}}(zeros(LowerTriangularMatrix{NF}, lmax+1, mmax+1)) # allocate memory in Λs for polynomials at all latitudes or allocate dummy array if precomputed # Λs is of size (lmax+1) x (mmax+1) x nlat_half unless recomputed @@ -216,40 +217,55 @@ end $(TYPEDSIGNATURES) 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, - ) where NF # number format NF - - lmax, mmax = size(alms) .- 1 # -1 for 0-based degree l, order m - return SpectralTransform(NF, Grid, lmax, mmax; recompute_legendre) +function SpectralTransform( + alms::LowerTriangularArray{NF}; # spectral coefficients + recompute_legendre::Bool = true, # saves memory + Grid::Type{<:AbstractGrid} = DEFAULT_GRID, + dealiasing::Real = DEFAULT_DEALIASING, +) where NF # number format NF (can be complex) + lmax, mmax = size(alms, ZeroBased, as=Matrix) # 0-based degree l, order m + return SpectralTransform(real(NF), Grid, lmax, mmax; recompute_legendre, dealiasing) end """ $(TYPEDSIGNATURES) -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::LowerTriangularMatrix{Complex{NF}}; # spectral coefficients - recompute_legendre::Bool = true, # saves memory - Grid::Type{<:AbstractGrid} = DEFAULT_GRID, - ) where NF # number format NF - - lmax, mmax = size(alms, as=Matrix) .- 1 # -1 for 0-based degree l, order m - return SpectralTransform(NF, Grid, lmax, mmax; recompute_legendre) +Generator function for a `SpectralTransform` struct based on the size and grid type of +gridded field `grids`. Recomputes the Legendre polynomials by default.""" +function SpectralTransform( + grids::AbstractGridArray{NF}; # gridded field + recompute_legendre::Bool=true, # saves memory as transform is garbage collected anyway + one_more_degree::Bool=false, # returns a square LowerTriangularMatrix by default + dealiasing::Real = DEFAULT_DEALIASING, +) where NF # number format NF + Grid = RingGrids.nonparametric_type(typeof(grids)) + trunc = get_truncation(grids, dealiasing) + return SpectralTransform(NF, Grid, trunc+one_more_degree, trunc; recompute_legendre, dealiasing) end -""" -$(TYPEDSIGNATURES) -Generator function for a `SpectralTransform` struct based on the size and grid type of -gridded field `map`. Recomputes the Legendre polynomials by default.""" -function SpectralTransform( map::AbstractGrid{NF}; # gridded field - recompute_legendre::Bool=true, # saves memory - one_more_degree::Bool=false, - ) where NF # number format NF - - Grid = RingGrids.nonparametric_type(typeof(map)) - trunc = get_truncation(map) - return SpectralTransform(NF, Grid, trunc+one_more_degree, trunc; recompute_legendre) +# CHECK MATCHING SIZES +function ismatching(S::SpectralTransform, L::LowerTriangularArray) + return (S.lmax, S.mmax) == size(L, ZeroBased, as=Matrix)[1:2] +end + +function ismatching(S::SpectralTransform, grid::AbstractGridArray) + match = S.Grid == RingGrids.nonparametric_type(typeof(grid)) && S.nlat_half == grid.nlat_half + return match +end + +# make `ismatching` commutative +ismatching(L::LowerTriangularArray, S::SpectralTransform) = ismatching(S, L) +ismatching(G::AbstractGridArray, S::SpectralTransform) = ismatching(S, G) + +function Base.DimensionMismatch(S::SpectralTransform, L::LowerTriangularArray) + s = "SpectralTransform(lmax=$(S.lmax), mmax=$(S.mmax)) and $(LowerTriangularMatrices.size2x_string(size(L, as=Matrix))) "* + "LowerTriangularArray do not match." + return DimensionMismatch(s) +end + +function Base.DimensionMismatch(S::SpectralTransform{NF1}, G::AbstractGridArray{NF2}) where {NF1, NF2} + s = "SpectralTransform{$NF1}($(S.Grid), nlat_half=$(S.nlat_half)) and "* + "$(RingGrids.nonparametric_type(G)){$NF2} with nlat_half=$(G.nlat_half) do not match." + return DimensionMismatch(s) end """ @@ -287,269 +303,278 @@ end # if number format not provided use Float64 get_recursion_factors(lmax::Int, mmax::Int) = get_recursion_factors(Float64, lmax, mmax) -""" -$(TYPEDSIGNATURES) -Spectral transform (spectral to grid) of the spherical harmonic coefficients `alms` to a gridded field -`map`. The spectral transform is number format-flexible as long as the parametric types of `map`, `alms`, `S` -are identical. The spectral transform is grid-flexible as long as the `typeof(map)<:AbstractGrid`. -Uses the precalculated arrays, FFT plans and other constants in the SpectralTransform struct `S`.""" -function gridded!( map::AbstractGrid{NF}, # gridded output - alms::LowerTriangularMatrix{Complex{NF}}, # spectral coefficients input - S::SpectralTransform{NF}; # precomputed parameters struct - unscale_coslat::Bool=false # unscale with cos(lat) on the fly? - ) where {NF<:AbstractFloat} # number format NF +"""$(TYPEDSIGNATURES) +Spectral transform (spectral to grid space) from n-dimensional array `specs` of spherical harmonic +coefficients to an n-dimensional array `grids` of ring grids. Uses FFT in the zonal direction, +and a Legendre Transform in the meridional direction exploiting symmetries. The spectral transform is +number format-flexible but `grids` and the spectral transform `S` have to have the same number format. +Uses the precalculated arrays, FFT plans and other constants in the SpectralTransform struct `S`. +The spectral transform is grid-flexible as long as the `typeof(grids)<:AbstractGridArray` and `S.Grid` +matches.""" +function transform!( # SPECTRAL TO GRID + grids::AbstractGridArray, # gridded output + specs::LowerTriangularArray, # spectral coefficients input + S::SpectralTransform{NF}; # precomputed transform + unscale_coslat::Bool=false, # unscale with cos(lat) on the fly? +) where NF # number format NF (; nlat, nlons, nlat_half, nfreq_max ) = S (; cos_colat, sin_colat, lon_offsets ) = S (; recompute_legendre, Λ, Λs, m_truncs ) = S (; brfft_plans ) = S - recompute_legendre && @boundscheck size(alms; as=Matrix) == size(Λ) || throw(BoundsError) - recompute_legendre || @boundscheck size(alms) == size(Λs[1]) || throw(BoundsError) - lmax, mmax = size(alms; as=Matrix) .- 1 # maximum degree l, order m of spherical harmonics + @boundscheck ismatching(S, grids) || throw(DimensionMismatch(S, grids)) + @boundscheck ismatching(S, specs) || throw(DimensionMismatch(S, specs)) - @boundscheck maximum(m_truncs) <= nfreq_max || throw(BoundsError) - @boundscheck nlat == length(cos_colat) || throw(BoundsError) - @boundscheck typeof(map) <: S.Grid || throw(BoundsError) - @boundscheck get_nlat_half(map) == S.nlat_half || throw(BoundsError) + lmax = specs.m - 1 # 0-based maximum degree l of spherical harmonics + mmax = specs.n - 1 # 0-based maximum order m of spherical harmonics # preallocate work arrays gn = zeros(Complex{NF}, nfreq_max) # phase factors for northern latitudes gs = zeros(Complex{NF}, nfreq_max) # phase factors for southern latitudes - rings = eachring(map) # precompute ring indices + rings = eachring(grids) # precomputed ring indices Λw = Legendre.Work(Legendre.λlm!, Λ, Legendre.Scalar(zero(Float64))) - @inbounds for j_north in 1:nlat_half # symmetry: loop over northern latitudes only - j_south = nlat - j_north + 1 # southern latitude index - nlon = nlons[j_north] # number of longitudes on this ring - nfreq = nlon÷2 + 1 # linear max Fourier frequency wrt to nlon - m_trunc = m_truncs[j_north] # (lin/quad/cub) max frequency to shorten loop over m - not_equator = j_north != j_south # is the latitude ring not on equator? - - # Recalculate or use precomputed Legendre polynomials Λ - recompute_legendre && Legendre.unsafe_legendre!(Λw, Λ, lmax, mmax, Float64(cos_colat[j_north])) - Λj = recompute_legendre ? Λ.data : Λs[j_north] - - # inverse Legendre transform by looping over wavenumbers l, m - lm = 1 # single index for non-zero l, m indices - for m in 1:m_trunc # Σ_{m=0}^{mmax}, but 1-based index, shortened to m_trunc - acc_odd = zero(Complex{NF}) # accumulator for isodd(l+m) - acc_even = zero(Complex{NF}) # accumulator for iseven(l+m) - - # integration over l = m:lmax+1 - lm_end = lm + lmax-m+1 # first index lm plus lmax-m+1 (length of column -1) - even_degrees = iseven(lm+lm_end) # is there an even number of degrees in column m? - - # anti-symmetry: sign change of odd harmonics on southern hemisphere - # but put both into one loop for contiguous memory access - for lm_even in lm:2:lm_end-even_degrees - # split into even, i.e. iseven(l+m) - # acc_even += alms[lm_even] * Λj[lm_even], but written with muladd - acc_even = muladd(alms[lm_even], Λj[lm_even], acc_even) - - # and odd (isodd(l+m)) harmonics - # acc_odd += alms[lm_odd] * Λj[lm_odd], but written with muladd - acc_odd = muladd(alms[lm_even+1], Λj[lm_even+1], acc_odd) - end - - # for even number of degrees, one acc_even iteration is skipped, do now - acc_even = even_degrees ? muladd(alms[lm_end], Λj[lm_end], acc_even) : acc_even - - acc_n = (acc_even + acc_odd) # accumulators for northern - acc_s = (acc_even - acc_odd) # and southern hemisphere - - # CORRECT FOR LONGITUDE OFFSETTS - o = lon_offsets[m, j_north] # longitude offset rotation + # loop over all specs/grids (e.g. vertical dimension) + # k, k_grid only differ when specs/grids have a singleton dimension + @inbounds for (k, k_grid) in zip(eachmatrix(specs), eachgrid(grids)) + for j_north in 1:nlat_half # symmetry: loop over northern latitudes only + j_south = nlat - j_north + 1 # southern latitude index + nlon = nlons[j_north] # number of longitudes on this ring + nfreq = nlon÷2 + 1 # linear max Fourier frequency wrt to nlon + m_trunc = m_truncs[j_north] # (lin/quad/cub) max frequency to shorten loop over m + not_equator = j_north != j_south # is the latitude ring not on equator? + + # Recalculate or use precomputed Legendre polynomials Λ + recompute_legendre && Legendre.unsafe_legendre!(Λw, Λ, lmax, mmax, Float64(cos_colat[j_north])) + Λj = recompute_legendre ? Λ : Λs[j_north] + + # INVERSE LEGENDRE TRANSFORM by looping over wavenumbers l, m + lm = 1 # single index for non-zero l, m indices + for m in 1:m_trunc # Σ_{m=0}^{mmax}, but 1-based index, shortened to m_trunc + acc_odd = zero(Complex{NF}) # accumulator for isodd(l+m) + acc_even = zero(Complex{NF}) # accumulator for iseven(l+m) + + # integration over l = m:lmax+1 + lm_end = lm + lmax-m+1 # first index lm plus lmax-m+1 (length of column -1) + even_degrees = iseven(lm+lm_end) # is there an even number of degrees in column m? + + # anti-symmetry: sign change of odd harmonics on southern hemisphere + # but put both into one loop for contiguous memory access + for lm_even in lm:2:lm_end-even_degrees + # split into even, i.e. iseven(l+m) + # acc_even += specs[lm_even] * Λj[lm_even], but written with muladd + acc_even = muladd(specs[lm_even, k], Λj[lm_even], acc_even) + + # and odd (isodd(l+m)) harmonics + # acc_odd += specs[lm_odd] * Λj[lm_odd], but written with muladd + acc_odd = muladd(specs[lm_even+1, k], Λj[lm_even+1], acc_odd) + end + + # for even number of degrees, one acc_even iteration is skipped, do now + acc_even = even_degrees ? muladd(specs[lm_end, k], Λj[lm_end], acc_even) : acc_even + + acc_n = (acc_even + acc_odd) # accumulators for northern + acc_s = (acc_even - acc_odd) # and southern hemisphere + + # CORRECT FOR LONGITUDE OFFSETTS + o = lon_offsets[m, j_north] # longitude offset rotation - gn[m] = muladd(acc_n, o, gn[m]) # accumulate in phase factors for northern - gs[m] = muladd(acc_s, o, gs[m]) # and southern hemisphere + gn[m] = muladd(acc_n, o, gn[m]) # accumulate in phase factors for northern + gs[m] = muladd(acc_s, o, gs[m]) # and southern hemisphere - lm = lm_end + 1 # first index of next m column - end + lm = lm_end + 1 # first index of next m column + end - if unscale_coslat - @inbounds cosθ = sin_colat[j_north] # sin(colat) = cos(lat) - gn ./= cosθ # scale in place - gs ./= cosθ - end + if unscale_coslat + @inbounds cosθ = sin_colat[j_north] # sin(colat) = cos(lat) + gn ./= cosθ # scale in place + gs ./= cosθ + end - # INVERSE FOURIER TRANSFORM in zonal direction - brfft_plan = brfft_plans[j_north] # FFT planned wrt nlon on ring - ilons = rings[j_north] # in-ring indices northern ring - LinearAlgebra.mul!(view(map.data, ilons), brfft_plan, view(gn, 1:nfreq)) # perform FFT + # INVERSE FOURIER TRANSFORM in zonal direction + brfft_plan = brfft_plans[j_north] # FFT planned wrt nlon on ring + ilons = rings[j_north] # in-ring indices northern ring + LinearAlgebra.mul!(view(grids.data, ilons, k_grid), brfft_plan, view(gn, 1:nfreq)) # perform FFT - # southern latitude, don't call redundant 2nd fft if ring is on equator - ilons = rings[j_south] # in-ring indices southern ring - not_equator && LinearAlgebra.mul!(view(map.data, ilons), brfft_plan, view(gs, 1:nfreq)) # perform FFT + # southern latitude, don't call redundant 2nd fft if ring is on equator + ilons = rings[j_south] # in-ring indices southern ring + not_equator && LinearAlgebra.mul!(view(grids.data, ilons, k_grid), brfft_plan, view(gs, 1:nfreq)) # perform FFT - fill!(gn, zero(Complex{NF})) # set phase factors back to zero - fill!(gs, zero(Complex{NF})) + fill!(gn, zero(Complex{NF})) # set phase factors back to zero + fill!(gs, zero(Complex{NF})) + end end - return map + return grids end -""" -$(TYPEDSIGNATURES) -Spectral transform (grid to spectral space) from the gridded field `map` on a `grid<:AbstractGrid` to -a `LowerTriangularMatrix` of spherical harmonic coefficients `alms`. Uses FFT in the zonal direction, +"""$(TYPEDSIGNATURES) +Spectral transform (grid to spectral space) from n-dimensional array of `grids` to an n-dimensional +array `specs` of spherical harmonic coefficients. Uses FFT in the zonal direction, and a Legendre Transform in the meridional direction exploiting symmetries. The spectral transform is -number format-flexible as long as the parametric types of `map`, `alms`, `S` are identical. -The spectral transform is grid-flexible as long as the `typeof(map)<:AbstractGrid`. -Uses the precalculated arrays, FFT plans and other constants in the SpectralTransform struct `S`.""" -function spectral!( alms::LowerTriangularMatrix{Complex{NF}}, # output: spectral coefficients - map::AbstractGrid{NF}, # input: gridded values - S::SpectralTransform{NF} - ) where {NF<:AbstractFloat} +number format-flexible but `grids` and the spectral transform `S` have to have the same number format. +Uses the precalculated arrays, FFT plans and other constants in the SpectralTransform struct `S`. +The spectral transform is grid-flexible as long as the `typeof(grids)<:AbstractGridArray` and `S.Grid` +matches.""" +function transform!( # grid -> spectral + specs::LowerTriangularArray, # output: spectral coefficients + grids::AbstractGridArray{NF}, # input: gridded values + S::SpectralTransform{NF} # precomputed spectral transform +) where NF # number format (; nlat, nlat_half, nlons, nfreq_max, cos_colat ) = S (; recompute_legendre, Λ, Λs, solid_angles ) = S (; rfft_plans, lon_offsets, m_truncs ) = S - recompute_legendre && @boundscheck size(alms; as=Matrix) == size(Λ) || throw(BoundsError) - recompute_legendre || @boundscheck size(alms) == size(Λs[1]) || throw(BoundsError) - lmax, mmax = size(alms; as=Matrix) .- 1 # maximum degree l, order m of spherical harmonics + @boundscheck ismatching(S, grids) || throw(DimensionMismatch(S, grids)) + @boundscheck ismatching(S, specs) || throw(DimensionMismatch(S, specs)) - @boundscheck maximum(m_truncs) <= nfreq_max || throw(BoundsError) - @boundscheck nlat == length(cos_colat) || throw(BoundsError) - @boundscheck typeof(map) <: S.Grid || throw(BoundsError) - @boundscheck get_nlat_half(map) == S.nlat_half || throw(BoundsError) + lmax = specs.m - 1 # 0-based maximum degree l of spherical harmonics + mmax = specs.n - 1 # 0-based maximum order m of spherical harmonics - # preallocate work warrays + # preallocate work arrays fn = zeros(Complex{NF}, nfreq_max) # Fourier-transformed northern latitude fs = zeros(Complex{NF}, nfreq_max) # Fourier-transformed southern latitude - rings = eachring(map) # precompute ring indices + rings = eachring(grids) # precompute ring indices - # partial sums are accumulated in alms, force zeros initially. - fill!(alms, 0) + # partial sums are accumulated in specs, force zeros initially. + fill!(specs, 0) Λw = Legendre.Work(Legendre.λlm!, Λ, Legendre.Scalar(zero(Float64))) - @inbounds for j_north in 1:nlat_half # symmetry: loop over northern latitudes only - j_south = nlat - j_north + 1 # corresponding southern latitude index - nlon = nlons[j_north] # number of longitudes on this ring - nfreq = nlon÷2 + 1 # linear max Fourier frequency wrt to nlon - m_trunc = m_truncs[j_north] # (lin/quad/cub) max frequency to shorten loop over m - not_equator = j_north != j_south # is the latitude ring not on equator? - - # FOURIER TRANSFORM in zonal direction - rfft_plan = rfft_plans[j_north] # FFT planned wrt nlon on ring - ilons = rings[j_north] # in-ring indices northern ring - LinearAlgebra.mul!(view(fn, 1:nfreq), rfft_plan, view(map.data, ilons)) # Northern latitude - - ilons = rings[j_south] # in-ring indices southern ring - # Southern latitude (don't call FFT on Equator) - # then fill fs with zeros and no changes needed further down - not_equator ? LinearAlgebra.mul!(view(fs, 1:nfreq), rfft_plan, view(map.data, ilons)) : fill!(fs, 0) - - # LEGENDRE TRANSFORM in meridional direction - # Recalculate or use precomputed Legendre polynomials Λ - recompute_legendre && Legendre.unsafe_legendre!(Λw, Λ, lmax, mmax, Float64(cos_colat[j_north])) - Λj = recompute_legendre ? Λ.data : Λs[j_north] - - # SOLID ANGLES including quadrature weights (sinθ Δθ) and azimuth (Δϕ) on ring j - ΔΩ = solid_angles[j_north] # = sinθ Δθ Δϕ, solid angle for a grid point - - lm = 1 # single index for spherical harmonics - for m in 1:m_trunc # Σ_{m=0}^{mmax}, but 1-based index - - an, as = fn[m], fs[m] - - # SOLID ANGLE QUADRATURE WEIGHTS and LONGITUDE OFFSET - o = lon_offsets[m, j_north] # longitude offset rotation - ΔΩ_rotated = ΔΩ*conj(o) # complex conjugate for rotation back to prime meridian - - # LEGENDRE TRANSFORM - a_even = (an + as)*ΔΩ_rotated # sign flip due to anti-symmetry with - a_odd = (an - as)*ΔΩ_rotated # odd polynomials - - # integration over l = m:lmax+1 - lm_end = lm + lmax-m+1 # first index lm plus lmax-m+1 (length of column -1) - even_degrees = iseven(lm+lm_end) # is there an even number of degrees in column m? + # loop over all specs/grids (e.g. vertical dimension) + # k, k_grid only differ when specs/grids have a singleton dimension + @inbounds for (k, k_grid) in zip(eachmatrix(specs), eachgrid(grids)) + for j_north in 1:nlat_half # symmetry: loop over northern latitudes only + j_south = nlat - j_north + 1 # corresponding southern latitude index + nlon = nlons[j_north] # number of longitudes on this ring + nfreq = nlon÷2 + 1 # linear max Fourier frequency wrt to nlon + m_trunc = m_truncs[j_north] # (lin/quad/cub) max frequency to shorten loop over m + not_equator = j_north != j_south # is the latitude ring not on equator? + + # FOURIER TRANSFORM in zonal direction + rfft_plan = rfft_plans[j_north] # FFT planned wrt nlon on ring + ilons = rings[j_north] # in-ring indices northern ring + LinearAlgebra.mul!(view(fn, 1:nfreq), rfft_plan, view(grids.data, ilons, k_grid)) # Northern latitude + + ilons = rings[j_south] # in-ring indices southern ring + # Southern latitude (don't call FFT on Equator) + # then fill fs with zeros and no changes needed further down + not_equator ? LinearAlgebra.mul!(view(fs, 1:nfreq), rfft_plan, view(grids.data, ilons, k_grid)) : fill!(fs, 0) + + # LEGENDRE TRANSFORM in meridional direction + # Recalculate or use precomputed Legendre polynomials Λ + recompute_legendre && Legendre.unsafe_legendre!(Λw, Λ, lmax, mmax, Float64(cos_colat[j_north])) + Λj = recompute_legendre ? Λ : Λs[j_north] - # anti-symmetry: sign change of odd harmonics on southern hemisphere - # but put both into one loop for contiguous memory access - for lm_even in lm:2:lm_end-even_degrees - # lm_odd = lm_even+1 - # split into even, i.e. iseven(l+m) - # alms[lm_even] += a_even * Λj[lm_even]#, but written with muladd - alms[lm_even] = muladd(a_even, Λj[lm_even], alms[lm_even]) - - # and odd (isodd(l+m)) harmonics - # alms[lm_odd] += a_odd * Λj[lm_odd]#, but written with muladd - alms[lm_even+1] = muladd(a_odd, Λj[lm_even+1], alms[lm_even+1]) - end + # SOLID ANGLES including quadrature weights (sinθ Δθ) and azimuth (Δϕ) on ring j + ΔΩ = solid_angles[j_north] # = sinθ Δθ Δϕ, solid angle for a grid point + + lm = 1 # single index for spherical harmonics + for m in 1:m_trunc # Σ_{m=0}^{mmax}, but 1-based index + + an, as = fn[m], fs[m] - # for even number of degrees, one even iteration is skipped, do now - alms[lm_end] = even_degrees ? muladd(a_even, Λj[lm_end], alms[lm_end]) : alms[lm_end] + # SOLID ANGLE QUADRATURE WEIGHTS and LONGITUDE OFFSET + o = lon_offsets[m, j_north] # longitude offset rotation + ΔΩ_rotated = ΔΩ*conj(o) # complex conjugate for rotation back to prime meridian - lm = lm_end + 1 # first index of next m column + # LEGENDRE TRANSFORM + a_even = (an + as)*ΔΩ_rotated # sign flip due to anti-symmetry with + a_odd = (an - as)*ΔΩ_rotated # odd polynomials + + # integration over l = m:lmax+1 + lm_end = lm + lmax-m+1 # first index lm plus lmax-m+1 (length of column -1) + even_degrees = iseven(lm+lm_end) # is there an even number of degrees in column m? + + # anti-symmetry: sign change of odd harmonics on southern hemisphere + # but put both into one loop for contiguous memory access + for lm_even in lm:2:lm_end-even_degrees + # lm_odd = lm_even+1 + # split into even, i.e. iseven(l+m) + # specs[lm_even] += a_even * Λj[lm_even]#, but written with muladd + specs[lm_even, k] = muladd(a_even, Λj[lm_even], specs[lm_even, k]) + + # and odd (isodd(l+m)) harmonics + # specs[lm_odd] += a_odd * Λj[lm_odd]#, but written with muladd + specs[lm_even+1, k] = muladd(a_odd, Λj[lm_even+1], specs[lm_even+1, k]) + end + + # for even number of degrees, one even iteration is skipped, do now + specs[lm_end, k] = even_degrees ? muladd(a_even, Λj[lm_end], specs[lm_end, k]) : specs[lm_end, k] + + lm = lm_end + 1 # first index of next m column + end end end - return alms + return specs end -""" -$(TYPEDSIGNATURES) -Spectral transform (spectral to grid space) from spherical coefficients `alms` to a newly allocated gridded -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::LowerTriangularMatrix{T}; # spectral coefficients - recompute_legendre::Bool = true, # saves memory - Grid::Type{<:AbstractGrid} = DEFAULT_GRID, - kwargs... - ) where {NF, T<:Complex{NF}} # number format NF - - lmax, mmax = size(alms; as=Matrix) .- 1 # -1 for 0-based degree l, order m - S = SpectralTransform(NF, Grid, lmax, mmax; recompute_legendre) - return gridded(alms, S; kwargs...) +# CONVENIENCE/ALLOCATING VERSIONS + +"""$(TYPEDSIGNATURES) +Spherical harmonic transform from `grids` to a newly allocated `specs::LowerTriangularArray` +using the precomputed spectral transform `S`.""" +function transform( # GRID TO SPECTRAL + grids::AbstractGridArray, # input grid + S::SpectralTransform{NF}, # precomputed spectral transform +) where NF + ks = size(grids)[2:end] # the non-horizontal dimensions + specs = zeros(LowerTriangularArray{Complex{NF}}, S.lmax+1, S.mmax+1, ks...) + transform!(specs, grids, S) + return specs end -""" -$(TYPEDSIGNATURES) -Spectral transform (spectral to grid space) from spherical coefficients `alms` to a newly allocated gridded -field `map` with precalculated properties based on the SpectralTransform struct `S`. `alms` is converted to -a `LowerTriangularMatrix` to execute the in-place `gridded!`.""" -function gridded( alms::LowerTriangularMatrix, # spectral coefficients - S::SpectralTransform{NF}; # struct for spectral transform parameters - kwargs... - ) where NF # number format NF - - map = zeros(S.Grid{NF}, S.nlat_half) # preallocate output - almsᴸ = zeros(LowerTriangularMatrix{Complex{NF}}, S.lmax+1, S.mmax+1) - copyto!(almsᴸ, alms) # drop the upper triangle and convert to NF - gridded!(map, almsᴸ, S; kwargs...) # now execute the in-place version - return map +"""$(TYPEDSIGNATURES) +Spherical harmonic transform from `specs` to a newly allocated `grids::AbstractGridArray` +using the precomputed spectral transform `S`.""" +function transform( # SPECTRAL TO GRID + specs::LowerTriangularArray, # input spectral coefficients + S::SpectralTransform{NF}; # precomputed spectral transform + kwargs... # pass on unscale_coslat=true/false(default) +) where NF + ks = size(specs)[2:end] # the non-horizontal dimensions + grids = zeros(S.Grid{NF}, S.nlat_half, ks...) + transform!(grids, specs, S; kwargs...) + return grids end """ $(TYPEDSIGNATURES) -Converts `map` to `Grid(map)` to execute `spectral(map::AbstractGrid; kwargs...)`.""" -function spectral( map::AbstractGrid{NF}; # gridded field - recompute_legendre::Bool = true, # saves memory - one_more_degree::Bool = false, # for lmax+2 x mmax+1 output size - ) where NF # number format NF - - Grid = typeof(map) - trunc = get_truncation(map.nlat_half) - S = SpectralTransform(NF, Grid, trunc+one_more_degree, trunc; recompute_legendre) - return spectral(map, S) +Spectral transform (spectral to grid space) from spherical coefficients `alms` to a newly allocated gridded +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 `transform(alms, S)`.""" +function transform( # SPECTRAL TO GRID + specs::LowerTriangularArray{NF}; # spectral coefficients input + recompute_legendre::Bool = true, # saves memory + Grid::Type{<:AbstractGrid} = DEFAULT_GRID, + dealiasing::Real = DEFAULT_DEALIASING, + kwargs... # pass on unscale_coslat=true/false(default) +) where NF # number format NF + lmax, mmax = size(specs, ZeroBased, as=Matrix) + S = SpectralTransform(real(NF), Grid, lmax, mmax; recompute_legendre, dealiasing) + return transform(specs, S; kwargs...) end """ $(TYPEDSIGNATURES) -Spectral transform (grid to spectral) `map` to `grid(map)` to execute `spectral(map::AbstractGrid; kwargs...)`.""" -function spectral( map::AbstractGrid, # gridded field - S::SpectralTransform{NF}, # spectral transform struct - ) where NF # number format NF - - map_NF = similar(map, NF) # convert map to NF - copyto!(map_NF, map) - - alms = LowerTriangularMatrix{Complex{NF}}(undef, S.lmax+1, S.mmax+1) - return spectral!(alms, map_NF, S) # in-place version -end +Spectral transform (grid to spectral space) from `grids` to a newly allocated `LowerTriangularArray`. +Based on the size of `grids` and the keyword `dealiasing` the spectral resolution trunc is +retrieved. SpectralTransform struct `S` is allocated to execute `transform(grids, S)`.""" +function transform( + grids::AbstractGridArray{NF}; # gridded fields + recompute_legendre::Bool = true, # saves memory + one_more_degree::Bool = false, # for lmax+2 x mmax+1 output size + dealiasing::Real = DEFAULT_DEALIASING, +) where NF # number format NF + + Grid = RingGrids.nonparametric_type(typeof(grids)) + trunc = get_truncation(grids.nlat_half, dealiasing) + S = SpectralTransform(NF, Grid, trunc+one_more_degree, trunc; recompute_legendre, dealiasing) + return transform(grids, S) +end \ No newline at end of file diff --git a/src/SpeedyTransforms/spectral_truncation.jl b/src/SpeedyTransforms/spectral_truncation.jl index 1dd5a41dd..a1ef5f8f4 100644 --- a/src/SpeedyTransforms/spectral_truncation.jl +++ b/src/SpeedyTransforms/spectral_truncation.jl @@ -45,7 +45,6 @@ function spectral_truncation!(A::AbstractMatrix) return A end - """ $(TYPEDSIGNATURES) Triangular truncation of `alms` to degree and order `trunc` in-place.""" @@ -132,8 +131,7 @@ function zero_imaginary_zonal_modes!( return alms end -""" -$(TYPEDSIGNATURES) +"""$(TYPEDSIGNATURES) Smooth the spectral field `A` following A_smooth = (1-c*∇²ⁿ)A with power n of a normalised Laplacian so that the highest degree lmax is dampened by multiplication with c. Anti-diffusion for c<0.""" function spectral_smoothing(A::LowerTriangularArray, c::Real; power::Real=1) @@ -142,8 +140,7 @@ function spectral_smoothing(A::LowerTriangularArray, c::Real; power::Real=1) return A_smooth end -""" -$(TYPEDSIGNATURES) +"""$(TYPEDSIGNATURES) Smooth the spectral field `A` following A *= (1-(1-c)*∇²ⁿ) with power n of a normalised Laplacian so that the highest degree lmax is dampened by multiplication with c. Anti-diffusion for c>1.""" function spectral_smoothing!( L::LowerTriangularArray, @@ -152,6 +149,7 @@ function spectral_smoothing!( L::LowerTriangularArray, truncation::Int=-1) # smoothing wrt wavenumber (0 = largest) lmax, mmax = size(L; as=Matrix) + # normalize by largest eigenvalue by default, or wrt to given truncation eigenvalue_norm = truncation == -1 ? -mmax*(mmax+1) : -truncation*(truncation+1) diff --git a/src/SpeedyTransforms/spectrum.jl b/src/SpeedyTransforms/spectrum.jl index e15ad7ea2..cc48f2652 100644 --- a/src/SpeedyTransforms/spectrum.jl +++ b/src/SpeedyTransforms/spectrum.jl @@ -1,9 +1,9 @@ function power_spectrum(alms::LowerTriangularMatrix{Complex{NF}}; normalize::Bool=true) where NF - lmax, mmax = size(alms; as=Matrix) # 1-based max degree l, order m + lmax, mmax = size(alms, OneBased, as=Matrix) # 1-based max degree l, order m trunc = min(lmax, mmax) # consider only the triangle - # ignore higher degrees if lmax > mmax + # ignore higher degrees if lmax > mmax spectrum = zeros(NF, trunc) # zonal modes m = 0, *1 as not mirrored at -m diff --git a/src/SpeedyWeather.jl b/src/SpeedyWeather.jl index f0a4f4925..f40a5c451 100644 --- a/src/SpeedyWeather.jl +++ b/src/SpeedyWeather.jl @@ -11,14 +11,13 @@ import LinearAlgebra: LinearAlgebra, Diagonal # GPU, PARALLEL import Base.Threads: Threads, @threads -import FLoops: FLoops, @floop import KernelAbstractions -import CUDA: CUDA, CUDAKernels +import CUDA: CUDA, CUDAKernels, CuArray import Adapt: Adapt, adapt, adapt_structure # INPUT OUTPUT import TOML -import Dates: Dates, DateTime, Period, Millisecond, Second, Minute, Hour, Day +import Dates: Dates, DateTime, Period, Millisecond, Second, Minute, Hour, Day, Week import Printf: Printf, @sprintf import Random: randstring import NCDatasets: NCDatasets, NCDataset, defDim, defVar @@ -29,7 +28,7 @@ import UnicodePlots import ProgressMeter # to avoid a `using Dates` to pass on DateTime arguments -export DateTime, Second, Minute, Hour, Day +export DateTime, Second, Minute, Hour, Day, Week # export functions that have many cross-component methods export initialize!, finish! @@ -49,7 +48,7 @@ include("LowerTriangularMatrices/LowerTriangularMatrices.jl") using .LowerTriangularMatrices # RingGrids -export RingGrids +export RingGrids export AbstractGrid, AbstractGridArray, AbstractFullGridarray, AbstractReducedGridArray export FullClenshawGrid, FullClenshawArray, @@ -61,13 +60,14 @@ export FullClenshawGrid, FullClenshawArray, HEALPixGrid, HEALPixArray, OctaHEALPixGrid, OctaHEALPixArray, eachring, eachgrid, plot +export AnvilInterpolator include("RingGrids/RingGrids.jl") using .RingGrids # SpeedyTransforms export SpeedyTransforms, SpectralTransform -export spectral, gridded, spectral!, gridded! +export transform, transform! export spectral_truncation, spectral_truncation! export curl, divergence, curl!, divergence! export ∇, ∇², ∇⁻², ∇!, ∇²!, ∇⁻²! @@ -82,7 +82,6 @@ include("gpu.jl") # abstract types include("models/abstract_models.jl") include("dynamics/abstract_types.jl") -include("output/abstract_types.jl") include("physics/abstract_types.jl") # GEOMETRY CONSTANTS ETC @@ -139,8 +138,9 @@ include("physics/land.jl") # OUTPUT include("output/schedule.jl") -include("output/output.jl") include("output/feedback.jl") +include("output/netcdf_output.jl") +include("output/restart_file.jl") include("output/plot.jl") include("output/callbacks.jl") include("output/particle_tracker.jl") diff --git a/src/dynamics/adiabatic_conversion.jl b/src/dynamics/adiabatic_conversion.jl index dc5950e78..f142d89be 100644 --- a/src/dynamics/adiabatic_conversion.jl +++ b/src/dynamics/adiabatic_conversion.jl @@ -2,16 +2,16 @@ abstract type AbstractAdiabaticConversion <: AbstractModelComponent end export AdiabaticConversion Base.@kwdef struct AdiabaticConversion{NF} <: AbstractAdiabaticConversion - nlev::Int + nlayers::Int "σ-related factor A needed for adiabatic conversion term" - σ_lnp_A::Vector{NF} = zeros(NF, nlev) + σ_lnp_A::Vector{NF} = zeros(NF, nlayers) "σ-related factor B needed for adiabatic conversion term" - σ_lnp_B::Vector{NF} = zeros(NF, nlev) + σ_lnp_B::Vector{NF} = zeros(NF, nlayers) end -AdiabaticConversion(SG::SpectralGrid; kwargs...) = AdiabaticConversion{SG.NF}(; nlev=SG.nlev, kwargs...) +AdiabaticConversion(SG::SpectralGrid; kwargs...) = AdiabaticConversion{SG.NF}(; nlayers=SG.nlayers, kwargs...) function initialize!( adiabatic::AdiabaticConversion, diff --git a/src/dynamics/coriolis.jl b/src/dynamics/coriolis.jl index d99e0ab85..87465d276 100644 --- a/src/dynamics/coriolis.jl +++ b/src/dynamics/coriolis.jl @@ -11,7 +11,7 @@ end Coriolis(SG::SpectralGrid; kwargs...) = Coriolis{SG.NF}(nlat=SG.nlat; kwargs...) -function initialize!(coriolis::Coriolis, model::ModelSetup) +function initialize!(coriolis::Coriolis, model::AbstractModel) (; rotation) = model.planet (; sinlat, radius) = model.geometry @@ -28,17 +28,18 @@ Return the Coriolis parameter `f` on the grid `Grid` of resolution `nlat_half` on a planet of `ratation` [1/s]. Default rotation of Earth.""" function coriolis( ::Type{Grid}, - nlat_half::Integer; + nlat_half::Integer, # resolution parameter + ks::Integer...; # non-horizontal dimensions rotation = DEFAULT_ROTATION -) where {Grid<:AbstractGrid} +) where {Grid<:AbstractGridArray} - f = zeros(Grid, nlat_half) # preallocate - lat = get_lat(Grid, nlat_half) # in radians [-π/2, π/2] + f = zeros(Grid, nlat_half, ks...) # preallocate + lat = get_lat(Grid, nlat_half) # in radians [-π/2, π/2] for (j, ring) in enumerate(eachring(f)) fⱼ = 2rotation*sin(lat[j]) for ij in ring - f[ij] = fⱼ + f[ij, :] .= fⱼ # setindex across all ks dimensions end end return f @@ -51,6 +52,6 @@ on a planet of `ratation` [1/s]. Default rotation of Earth.""" function coriolis( grid::Grid; kwargs... -) where {Grid<:AbstractGrid} - return coriolis(Grid, grid.nlat_half; kwargs...) +) where {Grid<:AbstractGridArray} + return coriolis(Grid, grid.nlat_half, size(grid)[2:end]...; kwargs...) end \ No newline at end of file diff --git a/src/dynamics/diagnostic_variables.jl b/src/dynamics/diagnostic_variables.jl index 4f4e7d162..b079120ae 100644 --- a/src/dynamics/diagnostic_variables.jl +++ b/src/dynamics/diagnostic_variables.jl @@ -1,239 +1,464 @@ -# const LTM = LowerTriangularMatrix # already defined in prognostic_variables +function Base.show(io::IO, A::AbstractDiagnosticVariables) + println(io, "$(typeof(A).name.wrapper)") + keys = propertynames(A) + keys_filtered = filter(key -> ~(getfield(A, key) isa Integer), keys) + n = length(keys_filtered) + for (i, key) in enumerate(keys_filtered) + last = i == n + val = getfield(A, key) + T = typeof(val) + + if T <: AbstractGridArray + NF = first_parameter(T) + nlayers = size(val, 2) + nlat = RingGrids.get_nlat(val) + Grid = RingGrids.nonparametric_type(T) + s = "$nlat-ring, $nlayers-layer $Grid{$NF}" + elseif T <: LowerTriangularArray + NF = first_parameter(T) + nlayers = size(val, 2) + trunc = val.n - 1 + nlayers = size(val, 3) + s = "T$trunc, $nlayers-layer LowerTriangularArray{$NF}" + else + s = "$T" + end + + ~last ? println(io, "├ $key: $s") : + print(io, "└ $key: $s") + end +end -""" -Tendencies of the prognostic spectral variables for a given layer. +first_parameter(::Type{<:AbstractArray{T}}) where T = T + +export Tendencies + +"""Tendencies of the prognostic variables in spectral and grid-point space $(TYPEDFIELDS)""" -Base.@kwdef struct Tendencies{NF<:AbstractFloat, Grid<:AbstractGrid{NF}} - nlat_half::Int - trunc::Int - vor_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) # Vorticity of horizontal wind field [1/s] - div_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) # Divergence of horizontal wind field [1/s] - temp_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) # Absolute temperature [K] - humid_tend::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) # Specific humidity [g/kg] - u_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) # zonal velocity (spectral) - v_tend ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) # meridional velocity (spectral) - u_tend_grid ::Grid = zeros(Grid, nlat_half) # zonal velocity (grid) - v_tend_grid ::Grid = zeros(Grid, nlat_half) # meridinoal velocity (grid) - temp_tend_grid ::Grid = zeros(Grid, nlat_half) # temperature - humid_tend_grid ::Grid = zeros(Grid, nlat_half) # specific humidity +@kwdef struct Tendencies{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + SpectralVariable2D, # <: LowerTriangularArray + SpectralVariable3D, # <: LowerTriangularArray + GridVariable2D, # <: AbstractGridArray + GridVariable3D, # <: AbstractGridArray +} <: AbstractDiagnosticVariables + + trunc::Int # spectral resolution: maximum degree and order of spherical harmonics + nlat_half::Int # grid resolution: number of latitude rings on one hemisphere (Eq. incl.) + nlayers::Int # number of vertical layers + + # SPECTRAL TENDENCIES + "Vorticity of horizontal wind field [1/s]" + vor_tend ::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + "Divergence of horizontal wind field [1/s]" + div_tend ::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + "Absolute temperature [K]" + temp_tend ::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + "Specific humidity [kg/kg]" + humid_tend::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + "Zonal velocity [m/s]" + u_tend ::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + "Meridional velocity [m/s]" + v_tend ::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + "Logarithm of surface pressure [Pa]" + pres_tend ::SpectralVariable2D = zeros(SpectralVariable2D, trunc+2, trunc+1) + + "Zonal velocity [m/s], grid" + u_tend_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Meridinoal velocity [m/s], grid" + v_tend_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Absolute temperature [K], grid" + temp_tend_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Specific humidity [kg/kg], grid" + humid_tend_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Logarith of surface pressure [Pa], grid" + pres_tend_grid ::GridVariable2D = zeros(GridVariable2D, nlat_half) end -# generator function based on a SpectralGrid -function Base.zeros(::Type{Tendencies}, - SG::SpectralGrid) - (; NF, trunc, Grid, nlat_half) = SG - return Tendencies{NF, Grid{NF}}(; nlat_half, trunc) +"""$(TYPEDSIGNATURES) +Generator function.""" +function Tendencies(SG::SpectralGrid) + (; trunc, nlat_half, nlayers, NF, ArrayType) = SG + (; SpectralVariable2D, SpectralVariable3D) = SG + (; GridVariable2D, GridVariable3D) = SG + + return Tendencies{NF, ArrayType, SpectralVariable2D, SpectralVariable3D, + GridVariable2D, GridVariable3D}(; + trunc, nlat_half, nlayers, + ) end -""" -Transformed prognostic variables (plus a few others) into grid-point space. +export GridVariables + +"""Transformed prognostic variables (and u, v, temp_virt) into grid-point space. $TYPEDFIELDS.""" -Base.@kwdef struct GridVariables{NF<:AbstractFloat, Grid<:AbstractGrid{NF}} - nlat_half::Int # resolution parameter for any grid - vor_grid ::Grid = zeros(Grid, nlat_half) # vorticity - div_grid ::Grid = zeros(Grid, nlat_half) # divergence - temp_grid ::Grid = zeros(Grid, nlat_half) # absolute temperature [K] - temp_grid_prev ::Grid = zeros(Grid, nlat_half) # absolute temperature of previous time step [K] - temp_virt_grid ::Grid = zeros(Grid, nlat_half) # virtual tempereature [K] - humid_grid ::Grid = zeros(Grid, nlat_half) # specific_humidity [kg/kg] - humid_grid_prev ::Grid = zeros(Grid, nlat_half) # specific_humidity at previous time step - u_grid ::Grid = zeros(Grid, nlat_half) # zonal velocity *coslat [m/s] - v_grid ::Grid = zeros(Grid, nlat_half) # meridional velocity *coslat [m/s] - u_grid_prev ::Grid = zeros(Grid, nlat_half) # zonal velocity *coslat of previous time step [m/s] - v_grid_prev ::Grid = zeros(Grid, nlat_half) # meridional velocity *coslat of previous time step [m/s] -end - -# generator function based on a SpectralGrid -function Base.zeros(::Type{GridVariables}, SG::SpectralGrid) - (; NF, Grid, nlat_half) = SG - return GridVariables{NF, Grid{NF}}(; nlat_half) +@kwdef struct GridVariables{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + GridVariable2D, # <: AbstractGridArray + GridVariable3D, # <: AbstractGridArray +} <: AbstractDiagnosticVariables + + nlat_half::Int # grid resolution: number of latitude rings on one hemisphere (Eq. incl.) + nlayers::Int # number of vertical layers + + "Relative vorticity of the horizontal wind [1/s]" + vor_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Divergence of the horizontal wind [1/s]" + div_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Absolute temperature [K]" + temp_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Virtual tempereature [K]" + temp_virt_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Specific_humidity [kg/kg]" + humid_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Zonal velocity [m/s]" + u_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Meridional velocity [m/s]" + v_grid ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Logarithm of surface pressure [Pa]" + pres_grid ::GridVariable2D = zeros(GridVariable2D, nlat_half) + + # PREVIOUS TIME STEP + "Absolute temperature [K] at previous time step" + temp_grid_prev ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Specific humidity [kg/kg] at previous time step" + humid_grid_prev ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Zonal velocity [m/s] at previous time step" + u_grid_prev ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Meridional velocity [m/s] at previous time step" + v_grid_prev ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + "Logarithm of surface pressure [Pa] at previous time step" + pres_grid_prev ::GridVariable2D = zeros(GridVariable2D, nlat_half) end -""" -Intermediate quantities for the dynamics of a given layer. -$(TYPEDFIELDS)""" -Base.@kwdef struct DynamicsVariables{NF<:AbstractFloat, Grid<:AbstractGrid{NF}} +"""$(TYPEDSIGNATURES) +Generator function.""" +function GridVariables(SG::SpectralGrid) + (; nlat_half, nlayers, NF, ArrayType) = SG + (; GridVariable2D, GridVariable3D) = SG - nlat_half::Int - trunc::Int + return GridVariables{NF, ArrayType, GridVariable2D, GridVariable3D}(; + nlat_half, nlayers, + ) +end - # MULTI-PURPOSE VECTOR (a, b), work array to be reused in various places, examples: - # uω_coslat⁻¹, vω_coslat⁻¹ = a, b (all models) - # uω_coslat⁻¹_grid, vω_coslat⁻¹_grid = a_grid, b_grid (all models) - # uh_coslat⁻¹, vh_coslat⁻¹ = a, b (ShallowWaterModel) - # uh_coslat⁻¹_grid, vh_coslat⁻¹_grid = a_grid, b_grid (ShallowWaterModel) - # Bernoulli potential: 1/2*(u^2+v^2) + Φ = a, a_grid (ShallowWater + PrimitiveEquation) - a::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) - b::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) - a_grid::Grid = zeros(Grid, nlat_half) - b_grid::Grid = zeros(Grid, nlat_half) +export DynamicsVariables + +"""Intermediate quantities for the dynamics of a given layer. +$(TYPEDFIELDS)""" +@kwdef struct DynamicsVariables{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + SpectralVariable2D, # <: LowerTriangularArray + SpectralVariable3D, # <: LowerTriangularArray + GridVariable2D, # <: AbstractGridArray + GridVariable3D, # <: AbstractGridArray +} <: AbstractDiagnosticVariables - # VERTICAL INTEGRATION - uv∇lnp ::Grid = zeros(Grid, nlat_half) # = (uₖ, vₖ)⋅∇ln(pₛ), pressure flux - uv∇lnp_sum_above::Grid = zeros(Grid, nlat_half) # sum of Δσₖ-weighted uv∇lnp above - div_sum_above ::Grid = zeros(Grid, nlat_half) # sum of div_weighted from top to k + trunc::Int # spectral resolution: maximum degree and order of spherical harmonics + nlat_half::Int # grid resolution: number of latitude rings on one hemisphere (Eq. incl.) + nlayers::Int # number of vertical layers - # virtual temperature spectral for geopot, geopotential on full layers - temp_virt ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) - geopot ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) + "Multi-purpose a, 3D work array to be reused in various places" + a::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + + "Multi-purpose b, 3D work array to be reused in various places" + b::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) + + "Multi-purpose a, 3D work array to be reused in various places" + a_grid::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + + "Multi-purpose b, 3D work array to be reused in various places" + b_grid::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + + "Multi-purpose a, work array to be reused in various places" + a_2D::SpectralVariable2D = zeros(SpectralVariable2D, trunc+2, trunc+1) + + "Multi-purpose b, work array to be reused in various places" + b_2D::SpectralVariable2D = zeros(SpectralVariable2D, trunc+2, trunc+1) + + "Multi-purpose a, work array to be reused in various places" + a_2D_grid::GridVariable2D = zeros(GridVariable2D, nlat_half) + + "Multi-purpose b, work array to be reused in various places" + b_2D_grid::GridVariable2D = zeros(GridVariable2D, nlat_half) - # VERTICAL VELOCITY (̇̇dσ/dt) - σ_tend::Grid = zeros(Grid, nlat_half) # = dσ/dt, on half levels below, at k+1/2 -end + "Pressure flux (uₖ, vₖ)⋅∇ln(pₛ)" + uv∇lnp::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + + "Sum of Δσₖ-weighted uv∇lnp above" + uv∇lnp_sum_above::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + + "Sum of div_weighted from top to k" + div_sum_above::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) + + "Virtual temperature [K], spectral for geopotential" + temp_virt::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) -# generator function based on a SpectralGrid -function Base.zeros(::Type{DynamicsVariables}, - SG::SpectralGrid) - (; NF, trunc, Grid, nlat_half) = SG - return DynamicsVariables{NF, Grid{NF}}(; nlat_half, trunc) -end + "Geopotential [m²/s²] on full layers" + geopot::SpectralVariable3D = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers) -export DiagnosticVariablesLayer + "Vertical velocity (dσ/dt), on half levels k+1/2 below, pointing to the surface (σ=1)" + σ_tend::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers) -""" -All diagnostic variables for a given layer: tendencies, prognostic varibles on the grid, -and intermediate dynamics variables. -$(TYPEDFIELDS)""" -struct DiagnosticVariablesLayer{NF<:AbstractFloat, Grid<:AbstractGrid{NF}} - npoints ::Int # number of grid points - k ::Int # which vertical model level? + "Zonal gradient of log surf pressure" + ∇lnp_x::GridVariable2D = zeros(GridVariable2D, nlat_half) - tendencies ::Tendencies{NF, Grid} - grid_variables ::GridVariables{NF, Grid} - dynamics_variables ::DynamicsVariables{NF, Grid} - temp_average ::Base.RefValue{NF} # average temperature for this level + "Meridional gradient of log surf pressure" + ∇lnp_y::GridVariable2D = zeros(GridVariable2D, nlat_half) + + "Vertical average of zonal velocity [m/s]" + u_mean_grid::GridVariable2D = zeros(GridVariable2D, nlat_half) + + "Vertical average of meridional velocity [m/s]" + v_mean_grid::GridVariable2D = zeros(GridVariable2D, nlat_half) + + "Vertical average of divergence [1/s], grid" + div_mean_grid::GridVariable2D = zeros(GridVariable2D, nlat_half) + + "Vertical average of divergence [1/s], spectral" + div_mean::SpectralVariable2D = zeros(SpectralVariable2D, trunc+2, trunc+1) end -# generator function based on a SpectralGrid -function Base.zeros(::Type{DiagnosticVariablesLayer}, - SG::SpectralGrid, - k::Integer=0) # use k=0 (i.e. unspecified) as default - (; npoints) = SG - tendencies = zeros(Tendencies, SG) - grid_variables = zeros(GridVariables, SG) - dynamics_variables = zeros(DynamicsVariables, SG) - temp_average = Ref(zero(SG.NF)) - return DiagnosticVariablesLayer(npoints, k, tendencies, grid_variables, dynamics_variables, temp_average) +"""$(TYPEDSIGNATURES) +Generator function.""" +function DynamicsVariables(SG::SpectralGrid) + (; trunc, nlat_half, nlayers, NF, ArrayType) = SG + (; SpectralVariable2D, SpectralVariable3D) = SG + (; GridVariable2D, GridVariable3D) = SG + + return DynamicsVariables{NF, ArrayType, SpectralVariable2D, SpectralVariable3D, + GridVariable2D, GridVariable3D}(; + trunc, nlat_half, nlayers, + ) end +export PhysicsVariables + """ -Diagnostic variables for the surface layer. +Diagnostic variables of the physical parameterizations. $(TYPEDFIELDS)""" -Base.@kwdef struct SurfaceVariables{NF<:AbstractFloat, Grid<:AbstractGrid{NF}} +@kwdef struct PhysicsVariables{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + GridVariable2D, # <: AbstractGridArray +} <: AbstractDiagnosticVariables nlat_half::Int - trunc::Int - npoints::Int - # log surface pressure, tendency of it and gridded tendency - pres_grid::Grid = zeros(Grid, nlat_half) - pres_tend::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) - pres_tend_grid::Grid = zeros(Grid, nlat_half) + "Accumualted large-scale precipitation [m]" + precip_large_scale::GridVariable2D = zeros(GridVariable2D, nlat_half) - ∇lnp_x::Grid = zeros(Grid, nlat_half) # zonal gradient of log surf pressure - ∇lnp_y::Grid = zeros(Grid, nlat_half) # meridional gradient of log surf pres + "Accumualted large-scale precipitation [m]" + precip_convection::GridVariable2D = zeros(GridVariable2D, nlat_half) - u_mean_grid::Grid = zeros(Grid, nlat_half) # vertical average of: zonal velocity *coslat - v_mean_grid::Grid = zeros(Grid, nlat_half) # meridional velocity *coslat - div_mean_grid::Grid = zeros(Grid, nlat_half) # divergence - div_mean::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) # divergence (in spectral though) + "Cloud top [m]" + cloud_top::GridVariable2D = zeros(GridVariable2D, nlat_half) - precip_large_scale::Grid = zeros(Grid, nlat_half) # large scale precipitation (for output) - precip_convection::Grid = zeros(Grid, nlat_half) # convective precipitation (for output) - cloud_top::Grid = zeros(Grid, nlat_half) # cloud top [hPa] - soil_moisture_availability::Grid = zeros(Grid, nlat_half) - cos_zenith::Grid = zeros(Grid, nlat_half) # cosine of solar zenith angle + "Availability of soil moisture to evaporation [1]" + soil_moisture_availability::GridVariable2D = zeros(GridVariable2D, nlat_half) + + "Cosine of solar zenith angle [1]" + cos_zenith::GridVariable2D = zeros(GridVariable2D, nlat_half) end -# generator function based on a SpectralGrid -function Base.zeros(::Type{SurfaceVariables}, - SG::SpectralGrid) +"""$(TYPEDSIGNATURES) +Generator function.""" +function PhysicsVariables(SG::SpectralGrid) + (; nlat_half, NF, ArrayType) = SG + (; GridVariable2D) = SG - (; NF, trunc, Grid, nlat_half, npoints) = SG - return SurfaceVariables{NF, Grid{NF}}(; nlat_half, trunc, npoints) + return PhysicsVariables{NF, ArrayType, GridVariable2D}(; nlat_half) end -Base.@kwdef struct ParticleVariables{NF<:AbstractFloat,Grid<:AbstractGrid} +export ParticleVariables + +"""Diagnostic variables for the particle advection +$(TYPEDFIELDS)""" +@kwdef struct ParticleVariables{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + ParticleVector, # <: AbstractGridArray + VectorNF, # Vector{NF} or CuVector{NF} + Grid, # <:AbstractGridArray +} <: AbstractDiagnosticVariables "Number of particles" - n_particles::Int + nparticles::Int "Number of latitudes on one hemisphere (Eq. incld.), resolution parameter of Grid" nlat_half::Int "Work array: particle locations" - locations::Vector{Particle{NF}} = zeros(Particle{NF}, n_particles) + locations::ParticleVector = zeros(ParticleVector, nparticles) "Work array: velocity u" - u::Vector{NF} = zeros(NF,n_particles) + u::VectorNF = VectorNF(zeros(nparticles)) "Work array: velocity v" - v::Vector{NF} = zeros(NF,n_particles) + v::VectorNF = VectorNF(zeros(nparticles)) "Work array: velocity w = dσ/dt" - σ_tend::Vector{NF} = zeros(NF,n_particles) + σ_tend::VectorNF = VectorNF(zeros(nparticles)) "Interpolator to interpolate velocity fields onto particle positions" - interpolator::AnvilInterpolator{NF,Grid} = AnvilInterpolator(NF, Grid, nlat_half, n_particles) + interpolator::AnvilInterpolator{NF, Grid} = AnvilInterpolator(NF, Grid, nlat_half, nparticles) end -function Base.zeros(::Type{ParticleVariables}, SG::SpectralGrid) - (; n_particles, nlat_half) = SG - ParticleVariables{SG.NF, SG.Grid}(; n_particles, nlat_half) +"""$(TYPEDSIGNATURES) +Generator function.""" +function ParticleVariables(SG::SpectralGrid) + (; nlat_half, nparticles, NF, ArrayType, Grid) = SG + (; ParticleVector) = SG + VectorNF = ArrayType{NF, 1} + return ParticleVariables{NF, ArrayType, ParticleVector, VectorNF, Grid}(; nlat_half, nparticles) end +# to be removed +struct DiagnosticVariablesLayer{NF} end +struct SurfaceVariables{NF} end + export DiagnosticVariables -""" -All diagnostic variables. +"""All diagnostic variables. $(TYPEDFIELDS)""" struct DiagnosticVariables{ - NF<:AbstractFloat, - Grid<:AbstractGrid{NF}, - Model<:ModelSetup + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + Grid, # <:AbstractGridArray + SpectralVariable2D, # <: LowerTriangularArray + SpectralVariable3D, # <: LowerTriangularArray + GridVariable2D, # <: AbstractGridArray + GridVariable3D, # <: AbstractGridArray + ParticleVector, # <: AbstractGridArray + VectorNF, # Vector{NF} or CuVector{NF} } <: AbstractDiagnosticVariables - layers ::Vector{DiagnosticVariablesLayer{NF, Grid}} - surface ::SurfaceVariables{NF, Grid} - columns ::Vector{ColumnVariables{NF}} - particles::ParticleVariables{NF} + # DIMENSIONS + "Spectral resolution: Max degree of spherical harmonics (0-based)" + trunc::Int + + "Grid resoltion: Number of latitude rings on one hemisphere (Equator incl.)" + nlat_half::Int + + "Number of vertical layers" + nlayers::Int + + "Number of particles for particle advection" + nparticles::Int + + "Tendencies (spectral and grid) of the prognostic variables" + tendencies::Tendencies{NF, ArrayType, SpectralVariable2D, SpectralVariable3D, GridVariable2D, GridVariable3D} + + "Gridded prognostic variables" + grid::GridVariables{NF, ArrayType, GridVariable2D, GridVariable3D} + + "Intermediate variables for the dynamical core" + dynamics::DynamicsVariables{NF, ArrayType, SpectralVariable2D, SpectralVariable3D, GridVariable2D, GridVariable3D} + + "Global fields returned from physics parameterizations" + physics::PhysicsVariables{NF, ArrayType, GridVariable2D} + + "Intermediate variables for the particle advection" + particles::ParticleVariables{NF, ArrayType, ParticleVector, VectorNF, Grid} + + "Vertical column for the physics parameterizations" + columns::Vector{ColumnVariables{NF}} - nlat_half::Int # resolution parameter of any Grid - nlev ::Int # number of vertical levels - npoints ::Int # number of grid points + "Average temperature of every horizontal layer [K]" + temp_average::VectorNF - scale::Base.RefValue{NF} # vorticity and divergence are scaled by radius + "Scale applied to vorticity and divergence" + scale::Base.RefValue{NF} end -# generator function based on a SpectralGrid -function Base.zeros( - ::Type{DiagnosticVariables}, - SG::SpectralGrid, - Model::Type{<:ModelSetup} -) +"""$(TYPEDSIGNATURES) +Generator function.""" +function DiagnosticVariables(SG::SpectralGrid) - (; NF, Grid, nlat_half, nlev, npoints) = SG - layers = [zeros(DiagnosticVariablesLayer, SG, k) for k in 1:nlev] - surface = zeros(SurfaceVariables, SG) + (; trunc, nlat_half, nparticles, NF, nlayers) = SG + + tendencies = Tendencies(SG) + grid = GridVariables(SG) + dynamics = DynamicsVariables(SG) + physics = PhysicsVariables(SG) + particles = ParticleVariables(SG) # create one column variable per thread to avoid race conditions nthreads = Threads.nthreads() - columns = [ColumnVariables{NF}(; nlev) for _ in 1:nthreads] + columns = [ColumnVariables{NF}(; nlayers) for _ in 1:nthreads] - # particle work arrays - particles = zeros(ParticleVariables, SG) + temp_average = SG.ArrayType{NF, 1}(undef, nlayers) scale = Ref(one(NF)) - return DiagnosticVariables{NF, Grid{NF}, Model}( - layers, surface, columns, particles, - nlat_half, nlev, npoints, scale) + return DiagnosticVariables( + trunc, nlat_half, nlayers, nparticles, + tendencies, grid, dynamics, physics, particles, + columns, temp_average, scale, + ) +end + +function Base.show( + io::IO, + diagn::DiagnosticVariables{NF, ArrayType, Grid}, +) where {NF, ArrayType, Grid} + println(io, "DiagnosticVariables{$NF, $ArrayType, $Grid}") + + (; trunc, nlat_half, nlayers, nparticles) = diagn + nlat = RingGrids.get_nlat(Grid, nlat_half) + println(io, "├ resolution: T$trunc, $nlat rings, $nlayers layers, $nparticles particles") + println(io, "├ tendencies::Tendencies") + println(io, "├ grid::GridVariables") + println(io, "├ dynamics::DynamicsVariables") + println(io, "├ physics::PhysicsVariables") + println(io, "├ particles::ParticleVariables") + println(io, "├ columns::Vector{ColumnVariables}") + println(io, "├ temp_average::$(typeof(diagn.temp_average))") + print(io, "└ scale: $(diagn.scale[])") +end + +"""$(TYPEDSIGNATURES) +Set the tendencies for the barotropic model to `x`.""" +function Base.fill!(tendencies::Tendencies, x, ::Type{<:Barotropic}) + fill!(tendencies.u_tend_grid, x) + fill!(tendencies.v_tend_grid, x) + fill!(tendencies.vor_tend, x) + return tendencies +end + +"""$(TYPEDSIGNATURES) +Set the tendencies for the shallow-water model to `x`.""" +function Base.fill!(tendencies::Tendencies, x, ::Type{<:ShallowWater}) + fill!(tendencies, x, Barotropic) # all tendencies also in Barotropic + fill!(tendencies.div_tend, x) # plus divergence and pressure + fill!(tendencies.pres_tend_grid, x) + fill!(tendencies.pres_tend, x) + return tendencies +end + +"""$(TYPEDSIGNATURES) +Set the tendencies for the primitive dry model to `x`.""" +function Base.fill!(tendencies::Tendencies, x, ::Type{<:PrimitiveDry}) + fill!(tendencies, x, ShallowWater) # all tendencies also in ShallowWater + fill!(tendencies.temp_tend, x) # plus temperature + fill!(tendencies.temp_tend_grid, x) + return tendencies +end + +""" +$(TYPEDSIGNATURES) +Set the tendencies for the primitive wet model to `x`.""" +function Base.fill!(tendencies::Tendencies, x, ::Type{<:PrimitiveWet}) + fill!(tendencies, x, PrimitiveDry) # all tendencies also in PrimitiveDry + fill!(tendencies.humid_tend, x) # plus humidity + fill!(tendencies.humid_tend_grid, x) + return tendencies end -DiagnosticVariables(SG::SpectralGrid) = zeros(DiagnosticVariables, SG, DEFAULT_MODEL) -DiagnosticVariables(SG::SpectralGrid, Model::Type{<:ModelSetup}) = zeros(DiagnosticVariables, SG, Model) -DiagnosticVariables(SG::SpectralGrid, model::ModelSetup) = zeros(DiagnosticVariables, SG, model_class(model)) +# fallback to primitive wet +Base.fill!(tendencies::Tendencies, x) = Base.fill!(tendencies, x, PrimitiveWet) +Base.fill!(tendencies::Tendencies, x, model::AbstractModel) = Base.fill!(tendencies, x, typeof(model)) -# LOOP OVER ALL GRID POINTS (extend from RingGrids module) -RingGrids.eachgridpoint(diagn::DiagnosticVariables) = Base.OneTo(diagn.npoints) -RingGrids.eachgridpoint(layer::DiagnosticVariablesLayer) = Base.OneTo(layer.npoints) -RingGrids.eachgridpoint(surface::SurfaceVariables) = Base.OneTo(surface.npoints) \ No newline at end of file +RingGrids.eachgridpoint(diagn::DiagnosticVariables) = eachgridpoint(diagn.grid.vor_grid) \ No newline at end of file diff --git a/src/dynamics/drag.jl b/src/dynamics/drag.jl index cea7be939..71d4696a8 100644 --- a/src/dynamics/drag.jl +++ b/src/dynamics/drag.jl @@ -1,20 +1,20 @@ abstract type AbstractDrag <: AbstractModelComponent end -## NO DRAG +## NO DRAG export NoDrag struct NoDrag <: AbstractDrag end NoDrag(SG::SpectralGrid) = NoDrag() -initialize!(::NoDrag, ::ModelSetup) = nothing +initialize!(::NoDrag, ::AbstractModel) = nothing -function drag!( diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, +function drag!( diagn::DiagnosticVariables, + progn::PrognosticVariables, drag::NoDrag, - time::DateTime, - model::ModelSetup) + model::AbstractModel, + lf::Integer) return nothing end -# Quadratic drag +# Quadratic drag export QuadraticDrag Base.@kwdef mutable struct QuadraticDrag{NF} <: AbstractDrag "[OPTION] drag coefficient [1]" @@ -27,17 +27,19 @@ end QuadraticDrag(SG::SpectralGrid; kwargs...) = QuadraticDrag{SG.NF}(; kwargs...) function initialize!( drag::QuadraticDrag, - model::ModelSetup) - # c = c_D / H * R + model::AbstractModel) + # c = c_D / H * R drag.c[] = drag.c_D / model.atmosphere.layer_thickness * model.geometry.radius end # function barrier -function drag!( diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - drag::QuadraticDrag, - time::DateTime, - model::ModelSetup) +function drag!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + drag::QuadraticDrag, + model::AbstractModel, + lf::Integer, +) drag!(diagn, drag) end @@ -48,15 +50,14 @@ Quadratic drag for the momentum equations. F = -c_D/H*|(u, v)|*(u, v) with c_D the non-dimensional drag coefficient as defined in `drag::QuadraticDrag`. -c_D and layer thickness `H` are precomputed in initialize!(::QuadraticDrag, ::ModelSetup) +c_D and layer thickness `H` are precomputed in initialize!(::QuadraticDrag, ::AbstractModel) and scaled by the radius as are the momentum equations.""" function drag!( - diagn::DiagnosticVariablesLayer, - drag::QuadraticDrag{NF}, -) where NF - - u = diagn.grid_variables.u_grid - v = diagn.grid_variables.v_grid + diagn::DiagnosticVariables, + drag::QuadraticDrag, +) + u = diagn.grid.u_grid + v = diagn.grid.v_grid Fu = diagn.tendencies.u_tend_grid Fv = diagn.tendencies.v_tend_grid @@ -64,9 +65,10 @@ function drag!( # total drag coefficient with radius scaling and /layer_thickness c = drag.c[] + k = diagn.nlayers # only apply to surface layer @inbounds for ij in eachgridpoint(u, v, Fu, Fv) - speed = sqrt(u[ij]^2 + v[ij]^2) - Fu[ij] -= c*speed*u[ij] # -= as the tendencies already contain forcing - Fv[ij] -= c*speed*v[ij] + speed = sqrt(u[ij, k]^2 + v[ij, k]^2) + Fu[ij, k] -= c*speed*u[ij, k] # -= as the tendencies already contain forcing + Fv[ij, k] -= c*speed*v[ij, k] end end \ No newline at end of file diff --git a/src/dynamics/forcing.jl b/src/dynamics/forcing.jl index 8550105b4..aa24b57ab 100644 --- a/src/dynamics/forcing.jl +++ b/src/dynamics/forcing.jl @@ -1,20 +1,20 @@ abstract type AbstractForcing <: AbstractModelComponent end -## NO FORCING = dummy forcing +## NO FORCING = dummy forcing export NoForcing struct NoForcing <: AbstractForcing end NoForcing(SG::SpectralGrid) = NoForcing() -initialize!(::NoForcing, ::ModelSetup) = nothing +initialize!(::NoForcing, ::AbstractModel) = nothing -function forcing!( diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, +function forcing!( diagn::DiagnosticVariables, + progn::PrognosticVariables, forcing::NoForcing, - time::DateTime, - model::ModelSetup) + model::AbstractModel, + lf::Integer) return nothing end -# JET STREAM FORCING +# JET STREAM FORCING export JetStreamForcing """ @@ -24,16 +24,22 @@ Galewsky, 2004, but mirrored for both hemispheres. $(TYPEDFIELDS) """ -Base.@kwdef mutable struct JetStreamForcing{NF} <: AbstractForcing +@kwdef mutable struct JetStreamForcing{NF} <: AbstractForcing "Number of latitude rings" nlat::Int = 0 + "Number of vertical layers" + nlayers::Int = 0 + "jet latitude [˚N]" latitude::NF = 45 "jet width [˚], default ≈ 19.29˚" width::NF = (1/4-1/7)*180 + "sigma level [1], vertical location of jet" + sigma::NF = 0.2 + "jet speed scale [m/s]" speed::NF = 85 @@ -42,23 +48,26 @@ Base.@kwdef mutable struct JetStreamForcing{NF} <: AbstractForcing "precomputed amplitude vector [m/s²]" amplitude::Vector{NF} = zeros(NF, nlat) + + "precomputed vertical tapering" + tapering::Vector{NF} = zeros(NF, nlayers) end JetStreamForcing(SG::SpectralGrid; kwargs...) = JetStreamForcing{SG.NF}( - ; nlat=SG.nlat, kwargs...) + ; nlat=SG.nlat, nlayers=SG.nlayers, kwargs...) function initialize!( forcing::JetStreamForcing, - model::ModelSetup) + model::AbstractModel) (; latitude, width, speed, time_scale, amplitude) = forcing (; radius) = model.spectral_grid - # Some constants similar to Galewsky 2004 + # Some constants similar to Galewsky 2004 θ₀ = (latitude-width)/360*2π # southern boundary of jet [radians] θ₁ = (latitude+width)/360*2π # northern boundary of jet eₙ = exp(-4/(θ₁-θ₀)^2) # normalisation, so that speed is at max A₀ = speed/eₙ/time_scale.value # amplitude [m/s²] without lat dependency - A₀ *= radius # scale by radius as are the momentum equations + A₀ *= radius # scale by radius as are the momentum equations (; nlat, colat) = model.geometry @@ -73,33 +82,45 @@ function initialize!( forcing::JetStreamForcing, end end + # vertical tapering + (; nlayers, sigma, tapering) = forcing + (; σ_levels_full) = model.geometry + + for k in 1:nlayers + tapering[k] = 1 - abs(sigma - σ_levels_full[k]) + end + return nothing end # function barrier -function forcing!( diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - forcing::JetStreamForcing, - time::DateTime, - model::ModelSetup) +function forcing!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + forcing::JetStreamForcing, + model::AbstractModel, + lf::Integer, +) forcing!(diagn, forcing) end -""" -$(TYPEDSIGNATURES) - +"""$(TYPEDSIGNATURES) Set for every latitude ring the tendency to the precomputed forcing in the momentum equations following the JetStreamForcing. -The forcing is precomputed in `initialize!(::JetStreamForcing, ::ModelSetup)`.""" -function forcing!( diagn::DiagnosticVariablesLayer, - forcing::JetStreamForcing) - Fu = diagn.tendencies.u_tend_grid - (; amplitude) = forcing +The forcing is precomputed in `initialize!(::JetStreamForcing, ::AbstractModel)`.""" +function forcing!( + diagn::DiagnosticVariables, + forcing::JetStreamForcing) - @inbounds for (j, ring) in enumerate(eachring(Fu)) - F = amplitude[j] - for ij in ring - Fu[ij] = F + Fu = diagn.tendencies.u_tend_grid + (; amplitude, tapering) = forcing + + @inbounds for k in eachgrid(Fu) + for (j, ring) in enumerate(eachring(Fu)) + F = amplitude[j] + for ij in ring + Fu[ij] = tapering[k]*F + end end end end \ No newline at end of file diff --git a/src/dynamics/geometry.jl b/src/dynamics/geometry.jl index 2d75475f7..f6c685689 100644 --- a/src/dynamics/geometry.jl +++ b/src/dynamics/geometry.jl @@ -7,7 +7,7 @@ Construct Geometry struct containing parameters and arrays describing an iso-lat and the vertical levels. Pass on `SpectralGrid` to calculate the following fields $(TYPEDFIELDS) """ -Base.@kwdef struct Geometry{NF<:AbstractFloat} <: AbstractGeometry +@kwdef struct Geometry{NF<:AbstractFloat} <: AbstractGeometry "SpectralGrid that defines spectral and grid resolution" spectral_grid::SpectralGrid @@ -28,9 +28,9 @@ Base.@kwdef struct Geometry{NF<:AbstractFloat} <: AbstractGeometry "number of latitude rings" nlat::Int = spectral_grid.nlat - + "number of vertical levels" - nlev::Int = spectral_grid.nlev + nlayers::Int = spectral_grid.nlayers "total number of grid points" npoints::Int = spectral_grid.npoints @@ -100,6 +100,9 @@ end $(TYPEDSIGNATURES) Generator function for `Geometry` struct based on `spectral_grid`.""" function Geometry(spectral_grid::SpectralGrid) + error_message = "nlayers=$(spectral_grid.nlayers) does not match length nlayers="* + "$(spectral_grid.vertical_coordinates.nlayers) in spectral_grid.vertical_coordinates." + @assert spectral_grid.nlayers == spectral_grid.vertical_coordinates.nlayers error_message return Geometry{spectral_grid.NF}(; spectral_grid) end @@ -116,16 +119,16 @@ function σ_interpolation_weights( σ_levels_half::AbstractVector) weights = zero(σ_levels_full) - nlev = length(weights) - nlev == 1 && return weights # escape early for 1 layer to avoid out-of-bounds access + nlayers = length(weights) + nlayers == 1 && return weights # escape early for 1 layer to avoid out-of-bounds access - for k in 1:nlev-1 + for k in 1:nlayers-1 weights[k] = (log(σ_levels_half[k+1]) - log(σ_levels_full[k])) / (log(σ_levels_full[k+1]) - log(σ_levels_full[k])) end # was log(0.99) in Fortran SPEEDY code but doesn't make sense to me - weights[end] = (log(σ_levels_half[nlev+1]) - log(σ_levels_full[nlev])) / - (log(σ_levels_full[nlev]) - log(σ_levels_full[nlev-1])) + weights[end] = (log(σ_levels_half[nlayers+1]) - log(σ_levels_full[nlayers])) / + (log(σ_levels_full[nlayers]) - log(σ_levels_full[nlayers-1])) return weights end @@ -134,30 +137,30 @@ end """ $(TYPEDSIGNATURES) Given a vector in column defined at full levels, do a linear interpolation in -log(σ) to calculate its values at half-levels, skipping top (k=1/2), extrapolating to bottom (k=NLEV+1/2). +log(σ) to calculate its values at half-levels, skipping top (k=1/2), extrapolating to bottom (k=nlayers+1/2). """ function vertical_interpolate!( A_half::Vector, # quantity A on half levels (excl top) A_full::Vector, # quantity A on full levels G::Geometry, ) - nlev = length(A_half) + nlayers = length(A_half) weights = G.full_to_half_interpolation # full levels contain one more for surface # TODO this is currently confusing because the surface fluxes use full[end] # as surface value which is technically on half levels though! - @boundscheck nlev <= length(A_full) || throw(BoundsError) - @boundscheck nlev <= length(weights) || throw(BoundsError) + @boundscheck nlayers <= length(A_full) || throw(BoundsError) + @boundscheck nlayers <= length(weights) || throw(BoundsError) # 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. Fortran SPEEDY documentation eq. (1) - for k = 1:nlev-1 + for k = 1:nlayers-1 A_half[k] = A_full[k] + weights[k]*(A_full[k+1] - A_full[k]) end # Compute the values at the surface separately - A_half[nlev] = A_full[nlev] + weights[nlev]*(A_full[nlev] - A_full[nlev-1]) + A_half[nlayers] = A_full[nlayers] + weights[nlayers]*(A_full[nlayers] - A_full[nlayers-1]) return nothing end \ No newline at end of file diff --git a/src/dynamics/geopotential.jl b/src/dynamics/geopotential.jl index ed964a317..b88a2750f 100644 --- a/src/dynamics/geopotential.jl +++ b/src/dynamics/geopotential.jl @@ -2,12 +2,12 @@ abstract type AbstractGeopotential <: AbstractModelComponent end export Geopotential Base.@kwdef struct Geopotential{NF} <: AbstractGeopotential - nlev::Int - Δp_geopot_half::Vector{NF} = zeros(NF, nlev) - Δp_geopot_full::Vector{NF} = zeros(NF, nlev) + nlayers::Int + Δp_geopot_half::Vector{NF} = zeros(NF, nlayers) + Δp_geopot_full::Vector{NF} = zeros(NF, nlayers) end -Geopotential(SG::SpectralGrid) = Geopotential{SG.NF}(; nlev=SG.nlev) +Geopotential(SG::SpectralGrid) = Geopotential{SG.NF}(; nlayers=SG.nlayers) """ $(TYPEDSIGNATURES) @@ -21,18 +21,18 @@ function initialize!( geopotential::Geopotential, model::PrimitiveEquation ) - (; Δp_geopot_half, Δp_geopot_full, nlev) = geopotential + (; Δp_geopot_half, Δp_geopot_full, nlayers) = geopotential (; R_dry) = model.atmosphere (; σ_levels_full, σ_levels_half) = model.geometry # 1. integration onto half levels - for k in 1:nlev-1 # k is full level index, 1=top, nlev=bottom + for k in 1:nlayers-1 # k is full level index, 1=top, nlayers=bottom # used for: Φ_{k+1/2} = Φ_{k+1} + R*T_{k+1}*(ln(p_{k+1}) - ln(p_{k+1/2})) Δp_geopot_half[k+1] = R_dry*log(σ_levels_full[k+1]/σ_levels_half[k+1]) end # 2. integration onto full levels (same formula but k -> k-1/2) - for k in 1:nlev + for k in 1:nlayers # used for: Φ_k = Φ_{k+1/2} + R*T_k*(ln(p_{k+1/2}) - ln(p_k)) Δp_geopot_full[k] = R_dry*log(σ_levels_half[k+1]/σ_levels_full[k]) end @@ -47,34 +47,27 @@ function geopotential!( geopotential::Geopotential, orography::AbstractOrography, ) - + (; geopot, temp_virt) = diagn.dynamics (; geopot_surf) = orography # = orography*gravity (; Δp_geopot_half, Δp_geopot_full) = geopotential # = R*Δlnp either on half or full levels - (; nlev) = diagn # number of vertical levels + (; nlayers) = diagn # number of vertical levels - @boundscheck nlev == length(Δp_geopot_full) || throw(BoundsError) + @boundscheck nlayers == length(Δp_geopot_full) || throw(BoundsError) # for PrimitiveDry virtual temperature = absolute temperature here # note these are not anomalies here as they are only in grid-point fields # BOTTOM FULL LAYER - temp = diagn.layers[end].dynamics_variables.temp_virt - geopot = diagn.layers[end].dynamics_variables.geopot - - @inbounds for lm in eachharmonic(geopot, geopot_surf, temp) - geopot[lm] = geopot_surf[lm] + temp[lm]*Δp_geopot_full[end] + local k::Int = nlayers + for lm in eachharmonic(geopot, geopot_surf, temp_virt) + geopot[lm, k] = geopot_surf[lm] + temp_virt[lm, k]*Δp_geopot_full[k] end # OTHER FULL LAYERS, integrate two half-layers from bottom to top - @inbounds for k in nlev-1:-1:1 - temp_k = diagn.layers[k].dynamics_variables.temp_virt - temp_k1 = diagn.layers[k+1].dynamics_variables.temp_virt - geopot_k = diagn.layers[k].dynamics_variables.geopot - geopot_k1 = diagn.layers[k+1].dynamics_variables.geopot - - for lm in eachharmonic(temp_k, temp_k1, geopot_k, geopot_k1) - geopot_k½ = geopot_k1[lm] + temp_k1[lm]*Δp_geopot_half[k+1] # 1st half layer integration - geopot_k[lm] = geopot_k½ + temp_k[lm]*Δp_geopot_full[k] # 2nd onto full layer + for k in nlayers-1:-1:1 + for lm in eachharmonic(geopot, temp_virt) + geopot_k½ = geopot[lm, k+1] + temp_virt[lm, k+1]*Δp_geopot_half[k+1] # 1st half layer integration + geopot[lm, k] = geopot_k½ + temp_virt[lm, k]*Δp_geopot_full[k] # 2nd onto full layer end end end @@ -91,18 +84,18 @@ function geopotential!( geopot_surf::Real = 0 ) - nlev = length(geopot) + nlayers = length(geopot) (; Δp_geopot_half, Δp_geopot_full) = G # = R*Δlnp either on half or full levels - @boundscheck length(temp) >= nlev || throw(BoundsError) - @boundscheck length(Δp_geopot_full) >= nlev || throw(BoundsError) - @boundscheck length(Δp_geopot_half) >= nlev || throw(BoundsError) + @boundscheck length(temp) >= nlayers || throw(BoundsError) + @boundscheck length(Δp_geopot_full) >= nlayers || throw(BoundsError) + @boundscheck length(Δp_geopot_half) >= nlayers || throw(BoundsError) # bottom layer - geopot[nlev] = geopot_surf + temp[nlev]*Δp_geopot_full[end] + geopot[nlayers] = geopot_surf + temp[nlayers]*Δp_geopot_full[end] # OTHER FULL LAYERS, integrate two half-layers from bottom to top - @inbounds for k in nlev-1:-1:1 + @inbounds for k in nlayers-1:-1:1 geopot[k] = geopot[k+1] + temp[k+1]*Δp_geopot_half[k+1] + temp[k]*Δp_geopot_full[k] end end @@ -118,9 +111,14 @@ 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, +function geopotential!( diagn::DiagnosticVariables, + pres::LowerTriangularArray, planet::AbstractPlanet) - (; geopot) = diagn.dynamics_variables - geopot .= pres * planet.gravity -end \ No newline at end of file + (; geopot) = diagn.dynamics + + # don't use broadcasting as geopot will have size Nxnlayers but pres N + # [lm] indexing bypasses the incompatible sizes (necessary for primitive models) + for lm in eachindex(geopot, pres) + geopot[lm] = pres[lm] * planet.gravity + end +end \ No newline at end of file diff --git a/src/dynamics/hole_filling.jl b/src/dynamics/hole_filling.jl index 6790e6677..4be10cb00 100644 --- a/src/dynamics/hole_filling.jl +++ b/src/dynamics/hole_filling.jl @@ -13,15 +13,13 @@ initialize!(::ClipNegatives, ::PrimitiveWet) = nothing # function barrier function hole_filling!( - A::AbstractGrid, + q::AbstractGridArray, H::ClipNegatives, model::PrimitiveWet ) - hole_filling!(A, H) + hole_filling!(q, H) end -function hole_filling!(A::AbstractGrid,::ClipNegatives) - @inbounds for ij in eachgridpoint(A) - A[ij] = max(A[ij], 0) - end +function hole_filling!(q::AbstractGridArray,::ClipNegatives) + @. q = max(q, 0) end \ No newline at end of file diff --git a/src/dynamics/horizontal_diffusion.jl b/src/dynamics/horizontal_diffusion.jl index 31f5ea927..56befca85 100644 --- a/src/dynamics/horizontal_diffusion.jl +++ b/src/dynamics/horizontal_diffusion.jl @@ -19,7 +19,7 @@ $(TYPEDFIELDS)""" trunc::Int "number of vertical levels" - nlev::Int + nlayers::Int # PARAMETERS "[OPTION] power of Laplacian" @@ -42,20 +42,20 @@ $(TYPEDFIELDS)""" tapering_σ::Float64 = 0.2 # ARRAYS, precalculated for each spherical harmonics degree and vertical layer - ∇²ⁿ::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 + ∇²ⁿ::Vector{Vector{NF}} = [zeros(NF, trunc+2) for _ in 1:nlayers] # explicit part + ∇²ⁿ_implicit::Vector{Vector{NF}} = [ones(NF, trunc+2) for _ in 1:nlayers] # implicit part # ARRAYS but no scaling or tapering and using time_scale_temp_humid - ∇²ⁿc::Vector{Vector{NF}} = [zeros(NF, trunc+2) for _ in 1:nlev] # explicit part - ∇²ⁿc_implicit::Vector{Vector{NF}} = [ones(NF, trunc+2) for _ in 1:nlev] # implicit part + ∇²ⁿc::Vector{Vector{NF}} = [zeros(NF, trunc+2) for _ in 1:nlayers] # explicit part + ∇²ⁿc_implicit::Vector{Vector{NF}} = [ones(NF, trunc+2) for _ in 1:nlayers] # 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...) + (; NF, trunc, nlayers) = spectral_grid # take resolution parameters from spectral_grid + return HyperDiffusion{NF}(; trunc, nlayers, kwargs...) end """$(TYPEDSIGNATURES) @@ -64,7 +64,7 @@ model time step, and possibly with a changing strength/power in the vertical. """ function initialize!( scheme::HyperDiffusion, - model::ModelSetup) + model::AbstractModel) initialize!(scheme, model.geometry, model.time_stepping) end @@ -76,7 +76,7 @@ function initialize!( G::AbstractGeometry, L::AbstractTimeStepper, ) - (; trunc, nlev, resolution_scaling) = scheme + (; trunc, nlayers, resolution_scaling) = scheme (; ∇²ⁿ, ∇²ⁿ_implicit, ∇²ⁿc, ∇²ⁿc_implicit) = scheme (; power, power_stratosphere, tapering_σ) = scheme (; Δt, radius) = L @@ -93,7 +93,7 @@ function initialize!( # (=more scale-selective for smaller wavenumbers) largest_eigenvalue = -trunc*(trunc+1) - for k in 1:nlev + for k in 1:nlayers # VERTICAL TAPERING for the stratosphere # go from 1 to 0 between σ=0 and tapering_σ σ = G.σ_levels_full[k] @@ -126,52 +126,58 @@ Apply horizontal diffusion to a 2D field `A` in spectral space by updating its t 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 +function horizontal_diffusion!( + tendency::LowerTriangularArray, # tendency of a + A::LowerTriangularArray, # spectral horizontal field + ∇²ⁿ_expl::AbstractVector, # explicit spectral damping (vector of k vectors of lmax length) + ∇²ⁿ_impl::AbstractVector # implicit spectral damping (vector of k vectors of lmax length) +) lmax, mmax = size(tendency; as=Matrix) # 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] + @boundscheck lmax <= length(∇²ⁿ_expl[1]) == length(∇²ⁿ_impl[1]) || throw(BoundsError) + + for k in eachmatrix(tendency, A) + lm = 0 + 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, k] = (tendency[lm, k] + ∇²ⁿ_expl[k][l]*A[lm, k]) * ∇²ⁿ_impl[k][l] + end + lm += 1 # skip last row for scalar quantities end - lm += 1 # skip last row for scalar quantities end end """$(TYPEDSIGNATURES) Apply horizontal diffusion to vorticity in the BarotropicModel.""" -function horizontal_diffusion!( diagn::DiagnosticVariablesLayer, - progn::PrognosticLayerTimesteps, - model::Barotropic, - lf::Int=1) # leapfrog index used (2 is unstable) - ∇²ⁿ = model.horizontal_diffusion.∇²ⁿ[diagn.k] - ∇²ⁿ_implicit = model.horizontal_diffusion.∇²ⁿ_implicit[diagn.k] +function horizontal_diffusion!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + model::Barotropic, + lf::Integer = 1, # leapfrog index used (2 is unstable) +) + (; ∇²ⁿ, ∇²ⁿ_implicit) = model.horizontal_diffusion # Barotropic model diffuses vorticity (only variable) - (; vor) = progn.timesteps[lf] + vor = progn.vor[lf] (; vor_tend) = diagn.tendencies horizontal_diffusion!(vor_tend, vor, ∇²ⁿ, ∇²ⁿ_implicit) end """$(TYPEDSIGNATURES) Apply horizontal diffusion to vorticity and divergence in the ShallowWaterModel.""" -function horizontal_diffusion!( progn::PrognosticLayerTimesteps, - diagn::DiagnosticVariablesLayer, - model::ShallowWater, - lf::Int=1) # leapfrog index used (2 is unstable) - ∇²ⁿ = model.horizontal_diffusion.∇²ⁿ[diagn.k] - ∇²ⁿ_implicit = model.horizontal_diffusion.∇²ⁿ_implicit[diagn.k] +function horizontal_diffusion!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + model::ShallowWater, + lf::Integer = 1, # leapfrog index used (2 is unstable) +) + (; ∇²ⁿ, ∇²ⁿ_implicit) = model.horizontal_diffusion # ShallowWater model diffuses vorticity and divergence - (; vor, div) = progn.timesteps[lf] + vor = progn.vor[lf] + div = progn.div[lf] (; vor_tend, div_tend) = diagn.tendencies horizontal_diffusion!(vor_tend, vor, ∇²ⁿ, ∇²ⁿ_implicit) horizontal_diffusion!(div_tend, div, ∇²ⁿ, ∇²ⁿ_implicit) @@ -180,20 +186,23 @@ end """$(TYPEDSIGNATURES) Apply horizontal diffusion applied to vorticity, divergence, temperature, and humidity (PrimitiveWet only) in the PrimitiveEquation models.""" -function horizontal_diffusion!( progn::PrognosticLayerTimesteps, - diagn::DiagnosticVariablesLayer, - model::PrimitiveEquation, - lf::Int=1) # leapfrog index used (2 is unstable) +function horizontal_diffusion!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + model::PrimitiveEquation, + lf::Integer = 1, # leapfrog index used (2 is unstable) +) # use weaker diffusion operators that don't taper or scale with resolution for temperature and humidity - ∇²ⁿc = model.horizontal_diffusion.∇²ⁿc[diagn.k] - ∇²ⁿc_implicit = model.horizontal_diffusion.∇²ⁿc_implicit[diagn.k] + (; ∇²ⁿc, ∇²ⁿc_implicit) = model.horizontal_diffusion # and the ones that do for vorticity and divergence - ∇²ⁿ = model.horizontal_diffusion.∇²ⁿ[diagn.k] - ∇²ⁿ_implicit = model.horizontal_diffusion.∇²ⁿ_implicit[diagn.k] + (; ∇²ⁿ, ∇²ⁿ_implicit) = model.horizontal_diffusion # Primitive equation models diffuse vor, divergence, temp (and humidity for wet core) - (; vor, div, temp, humid) = progn.timesteps[lf] + vor = progn.vor[lf] + div = progn.div[lf] + temp = progn.temp[lf] + humid = progn.humid[lf] (; vor_tend, div_tend, temp_tend, humid_tend) = diagn.tendencies horizontal_diffusion!(vor_tend, vor, ∇²ⁿ, ∇²ⁿ_implicit) horizontal_diffusion!(div_tend, div, ∇²ⁿ, ∇²ⁿ_implicit) diff --git a/src/dynamics/implicit.jl b/src/dynamics/implicit.jl index e27cc13a3..694fb9f23 100644 --- a/src/dynamics/implicit.jl +++ b/src/dynamics/implicit.jl @@ -1,6 +1,6 @@ abstract type AbstractImplicit <: AbstractModelComponent end -# BAROTROPIC MODEL (no implicit needed) +# BAROTROPIC MODEL (no implicit needed) export NoImplicit struct NoImplicit <: AbstractImplicit end NoImplicit(SG::SpectralGrid) = NoImplicit() @@ -14,20 +14,20 @@ export ImplicitShallowWater Struct that holds various precomputed arrays for the semi-implicit correction to prevent gravity waves from amplifying in the shallow water model. $(TYPEDFIELDS)""" -Base.@kwdef struct ImplicitShallowWater{NF<:AbstractFloat} <: AbstractImplicit +@kwdef struct ImplicitShallowWater{NF<:AbstractFloat} <: AbstractImplicit # DIMENSIONS trunc::Int - "coefficient for semi-implicit computations to filter gravity waves" - α::Float64 = 1 + "[OPTION] coefficient for semi-implicit computations to filter gravity waves, 0.5 <= α <= 1" + α::NF = 1 # PRECOMPUTED ARRAYS, to be initiliased with initialize! - H::Base.RefValue{NF} = Ref(zero(NF)) # layer_thickness - ξ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 + H::Base.RefValue{NF} = Ref(zero(NF)) # layer_thickness + ξ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 """ @@ -62,7 +62,7 @@ function initialize!( # α = 0.5 evaluates at i+1 and i-1 (centered implicit) # α = 1 evaluates at i+1 (backward implicit) # α ∈ [0.5, 1] are also possible which controls the strength of the gravity wave dampening. - # α = 0.5 slows gravity waves and prevents them from amplifying + # α = 0.5 slows gravity waves and prevents them from amplifying # α > 0.5 will dampen the gravity waves within days to a few timesteps (α=1) ξ = α*dt # new implicit timestep ξ = α*dt = 2αΔt (for leapfrog) from input dt @@ -71,7 +71,7 @@ function initialize!( # loop over degree l of the harmonics (implicit terms are independent of order m) @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) + eigenvalue = -l*(l-1) # =∇², with without 1/radius², 1-based -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[]*ξg∇²[l]) # update 1/(1-ξ²gH∇²) with new ξ @@ -83,45 +83,46 @@ $(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}, - implicit::ImplicitShallowWater) where NF - - (; div_tend) = diagn.tendencies # divergence tendency - div_old = progn.timesteps[1].div # divergence at t - div_new = progn.timesteps[2].div # divergence at t+dt - pres_old = progn_surface.timesteps[1].pres # pressure/η at t - pres_new = progn_surface.timesteps[2].pres # pressure/η at t+dt - (; pres_tend) = diagn_surface # tendency of pressure/η +function implicit_correction!( diagn::DiagnosticVariables, + progn::PrognosticVariables, + implicit::ImplicitShallowWater) + + (; div_tend, pres_tend) = diagn.tendencies # tendency of divergence and pressure/η + div_old = progn.div[1] # divergence at t + div_new = progn.div[2] # divergence at t+dt + pres_old = progn.pres[1] # pressure/η at t + pres_new = progn.pres[2] # pressure/η at t+dt (; 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 - lmax, mmax = size(div_tend; as=Matrix) .- (2, 1) + lmax, mmax = size(div_tend, ZeroBased, as=Matrix) + lmax -= 1 + @boundscheck length(S⁻¹) == lmax+2 || throw(BoundsError) @boundscheck length(ξg∇²) == lmax+2 || throw(BoundsError) @boundscheck length(g∇²) == lmax+2 || throw(BoundsError) - lm = 0 - @inbounds for m in 1:mmax+1 - for l in m:lmax+1 - lm += 1 # single index lm corresponding to harmonic l, m with a LowerTriangularMatrix - - # calculate the G = N(Vⁱ) + NI(Vⁱ⁻¹ - Vⁱ) term. - # Vⁱ is a prognostic variable at time step i - # N is the right hand side of ∂V\∂t = N(V) - # NI is the part of N that's calculated semi-implicitily: N = NE + NI - G_div = div_tend[lm] - g∇²[l]*(pres_old[lm] - pres_new[lm]) - G_η = pres_tend[lm] - H*(div_old[lm] - div_new[lm]) - - # using the Gs correct the tendencies for semi-implicit time stepping - div_tend[lm] = S⁻¹[l]*(G_div - ξg∇²[l]*G_η) - pres_tend[lm] = G_η - ξH*div_tend[lm] + for k in eachmatrix(div_tend) + lm = 0 + for m in 1:mmax+1 + for l in m:lmax+1 + lm += 1 # single index lm corresponding to harmonic l, m with a LowerTriangularMatrix + + # calculate the G = N(Vⁱ) + NI(Vⁱ⁻¹ - Vⁱ) term. + # Vⁱ is a prognostic variable at time step i + # N is the right hand side of ∂V\∂t = N(V) + # NI is the part of N that's calculated semi-implicitily: N = NE + NI + G_div = div_tend[lm, k] - g∇²[l]*(pres_old[lm] - pres_new[lm]) + G_η = pres_tend[lm] - H*(div_old[lm, k] - div_new[lm, k]) + + # using the Gs correct the tendencies for semi-implicit time stepping + div_tend[lm, k] = S⁻¹[l]*(G_div - ξg∇²[l]*G_η) + pres_tend[lm] = G_η - ξH*div_tend[lm, k] + end + lm += 1 # loop skips last row end - lm += 1 # loop skips last row end end @@ -131,65 +132,65 @@ export ImplicitPrimitiveEquation Struct that holds various precomputed arrays for the semi-implicit correction to prevent gravity waves from amplifying in the primitive equation model. $(TYPEDFIELDS)""" -Base.@kwdef struct ImplicitPrimitiveEquation{NF<:AbstractFloat} <: AbstractImplicit +@kwdef struct ImplicitPrimitiveEquation{NF<:AbstractFloat} <: AbstractImplicit # DIMENSIONS "spectral resolution" trunc::Int - "number of vertical levels" - nlev::Int + "number of vertical layers" + nlayers::Int - # PARAMETERS + # PARAMETERS "time-step coefficient: 0=explicit, 0.5=centred implicit, 1=backward implicit" - α::Float64 = 1 + α::NF = 1 # PRECOMPUTED ARRAYS, to be initiliased with initialize! - "vertical temperature profile, obtained from diagn" - temp_profile::Vector{NF} = zeros(NF, nlev) + "vertical temperature profile, obtained from diagn on first time step" + temp_profile::Vector{NF} = zeros(NF, nlayers) "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) + R::Matrix{NF} = zeros(NF, nlayers, nlayers) "divergence: the -RdTₖ∇² term excl the eigenvalues from ∇² for divergence" - U::Vector{NF} = zeros(NF, nlev) + U::Vector{NF} = zeros(NF, nlayers) "temperature: operator for the TₖD + κTₖDlnps/Dt term" - L::Matrix{NF} = zeros(NF, nlev, nlev) + L::Matrix{NF} = zeros(NF, nlayers, nlayers) "pressure: vertical averaging of the -D̄ term in the log surface pres equation" - W::Vector{NF} = zeros(NF, nlev) + W::Vector{NF} = zeros(NF, nlayers) "components to construct L, 1/ 2Δσ" - L0::Vector{NF} = zeros(NF, nlev) + L0::Vector{NF} = zeros(NF, nlayers) "vert advection term in the temperature equation (below+above)" - L1::Matrix{NF} = zeros(NF, nlev, nlev) + L1::Matrix{NF} = zeros(NF, nlayers, nlayers) "factor in front of the div_sum_above term" - L2::Vector{NF} = zeros(NF, nlev) + L2::Vector{NF} = zeros(NF, nlayers) "_sum_above operator itself" - L3::Matrix{NF} = zeros(NF, nlev, nlev) + L3::Matrix{NF} = zeros(NF, nlayers, nlayers) "factor in front of div term in Dlnps/Dt" - L4::Vector{NF} = zeros(NF, nlev) + L4::Vector{NF} = zeros(NF, nlayers) "for every l the matrix to be inverted" - S::Matrix{NF} = zeros(NF, nlev, nlev) + S::Matrix{NF} = zeros(NF, nlayers, nlayers) "combined inverted operator: S = 1 - ξ²(RL + UW)" - S⁻¹::Array{NF, 3} = zeros(NF, trunc+1, nlev, nlev) + S⁻¹::Array{NF, 3} = zeros(NF, trunc+1, nlayers, nlayers) end """$(TYPEDSIGNATURES) Generator using the resolution from SpectralGrid.""" function ImplicitPrimitiveEquation(spectral_grid::SpectralGrid, kwargs...) - (; NF, trunc, nlev) = spectral_grid - return ImplicitPrimitiveEquation{NF}(; trunc, nlev, kwargs...) + (; NF, trunc, nlayers) = spectral_grid + return ImplicitPrimitiveEquation{NF}(; trunc, nlayers, kwargs...) end # function barrier to unpack the constants struct for primitive eq models @@ -215,29 +216,27 @@ function initialize!( adiabatic_conversion::AbstractAdiabaticConversion, ) - (; trunc, nlev, α, temp_profile, S, S⁻¹, L, R, U, W, L0, L1, L2, L3, L4) = implicit + (; trunc, nlayers, α, temp_profile, S, S⁻¹, L, R, U, W, L0, L1, L2, L3, L4) = implicit (; σ_levels_full, σ_levels_thick) = geometry (; R_dry, κ) = atmosphere (; Δp_geopot_half, Δp_geopot_full) = geopotential (; σ_lnp_A, σ_lnp_B) = adiabatic_conversion - for k in 1:nlev - # use current vertical temperature profile - temp_profile[k] = diagn.layers[k].temp_average[] + # use current vertical temperature profile + temp_profile .= diagn.temp_average - # return immediately if temp_profile contains NaRs, model blew up in that case - if !isfinite(temp_profile[k]) return nothing end - end + # return immediately if temp_profile contains NaRs, model blew up in that case + all(isfinite.(temp_profile)) || return nothing # set up R, U, L, W operators from # δD = G_D + ξ(RδT + Uδlnps) divergence D correction # δT = G_T + ξLδD temperature T correction # δlnps = G_lnps + ξWδD log surface pressure lnps correction - #  + # # G_X is the uncorrected explicit tendency calculated as RHS_expl(Xⁱ) + RHS_impl(Xⁱ⁻¹) # with RHS_expl being the nonlinear terms calculated from the centered time step i # and RHS_impl are the linear terms that are supposed to be calcualted semi-implicitly - # however, they have sofar only been evaluated explicitly at time step i-1 + # however, they have sofar only been evaluated explicitly at time step i-1 # and are subject to be corrected to δX following the equations above # 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 @@ -246,9 +245,9 @@ function initialize!( 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 - R[1:k, k] .= -Δp_geopot_full[k] # otherwise equivalent to geopotential! with zero orography - R[1:k-1, k] .+= -Δp_geopot_half[k] # incl the minus but excluding the eigenvalues as with U + @inbounds for k in 1:nlayers # vertical geopotential integration as matrix operator + R[1:k, k] .= -Δp_geopot_full[k] # otherwise equivalent to geopotential! with zero orography + R[1:k-1, k] .+= -Δp_geopot_half[k] # incl the minus but excluding the eigenvalues as with U end U .= -R_dry*temp_profile # the R_d*Tₖ∇² term excl the eigenvalues from ∇² for divergence @@ -257,16 +256,16 @@ function initialize!( L2 .= κ*temp_profile.*σ_lnp_A # factor in front of the div_sum_above term L4 .= κ*temp_profile.*σ_lnp_B # factor in front of div term in Dlnps/Dt - @inbounds for k in 1:nlev + @inbounds for k in 1:nlayers Tₖ = temp_profile[k] # average temperature at k - k_above = max(1, k-1) # layer index above - k_below = min(k+1, nlev) # layer index below + k_above = max(1, k-1) # layer index above + k_below = min(k+1, nlayers) # layer index below ΔT_above = Tₖ - temp_profile[k_above] # temperature difference to layer above ΔT_below = temp_profile[k_below] - Tₖ # and to layer below σₖ = σ_levels_full[k] # should be Σ_r=1^k Δσᵣ for model top at >0hPa σₖ_above = σ_levels_full[k_above] - for r in 1:nlev + for r in 1:nlayers L1[k, r] = ΔT_below*σ_levels_thick[r]*σₖ # vert advection operator below L1[k, r] -= k>=r ? σ_levels_thick[r] : 0 @@ -285,19 +284,19 @@ function initialize!( W .= -σ_levels_thick # the -D̄ term in the log surface pres equation # solving the equations above for δD yields - # δD = SG, with G = G_D + ξRG_T + ξUG_lnps and the operator S + # δ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) + I = LinearAlgebra.I(nlayers) @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) - Sinv = L1 # reuse L1 matrix to store inv(S) - Sinv .= I # use ldiv! so last arg needs to be unity matrix - LinearAlgebra.ldiv!(luS, Sinv) # now do S\I = S⁻¹ via LU decomposition - S⁻¹[l, :, :] .= Sinv # store in array + Sinv = L1 # reuse L1 matrix to store inv(S) + Sinv .= I # use ldiv! so last arg needs to be unity matrix + LinearAlgebra.ldiv!(luS, Sinv) # now do S\I = S⁻¹ via LU decomposition + S⁻¹[l, :, :] .= Sinv # store in array end end @@ -308,87 +307,69 @@ function implicit_correction!( implicit::ImplicitPrimitiveEquation, progn::PrognosticVariables, ) - # escape immediately if explicit implicit.α == 0 && return nothing - # (; Δp_geopot_half, Δp_geopot_full) = model.geometry # = R*Δlnp on half or full levels - (; nlev, trunc, S⁻¹, R, U, L, W) = implicit + (; nlayers, trunc) = implicit + (; 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 # so is the -D̄ term for surface pressure in tendencies! - @floop for k in 1:nlev - (; temp_tend) = diagn.layers[k].tendencies # unpack temp_tend - for r in 1:nlev - div_old = progn.layers[r].timesteps[1].div # divergence at i-1 - div_new = progn.layers[r].timesteps[2].div # divergence at i - - # RHS_expl(Vⁱ) + RHS_impl(Vⁱ⁻¹) = RHS(Vⁱ) + RHS_impl(Vⁱ⁻¹ - Vⁱ) - # for temperature tendency do the latter as its cheaper. - @. temp_tend += L[k, r]*(div_old-div_new) - # @. temp_tend += L[k, r]*div_old # for the former + (; temp_tend) = diagn.tendencies + div_old, div_new = progn.div # divergence at i-1 (old), i (new, i.e. current) + + for k in eachmatrix(temp_tend, div_old, div_new) + for r in eachmatrix(temp_tend, div_old, div_new) + for lm in eachharmonic(temp_tend, div_old, div_new) + # RHS_expl(Vⁱ) + RHS_impl(Vⁱ⁻¹) = RHS(Vⁱ) + RHS_impl(Vⁱ⁻¹ - Vⁱ) + # for temperature tendency do the latter as its cheaper. + temp_tend[lm, k] += L[k, r] * (div_old[lm, r] - div_new[lm, r]) + # temp_tend[lm, k] += L[k, r] * div_old[lm, r] # for the former + end end - end + end # SEMI IMPLICIT CORRECTIONS FOR DIVERGENCE - (; pres_tend) = diagn.surface - @floop for k in 1:nlev # loop from bottom layer to top for geopotential calculation - - # calculate the combined tendency G = G_D + ξRG_T + ξUG_lnps to solve for divergence δD - G = diagn.layers[k].dynamics_variables.a # reuse work arrays, used for combined tendency G - geopot = diagn.layers[k].dynamics_variables.b # used for geopotential - (; div_tend) = diagn.layers[k].tendencies # unpack div_tend - - # 1. the ξ*R*G_T term, vertical integration of geopotential (excl ξ, this is done in 2.) - @inbounds for r in k:nlev # skip 1:k-1 as integration is surface to k - (; temp_tend) = diagn.layers[r].tendencies # unpack temp_tend - @. geopot += R[k, r]*temp_tend + # calculate the combined tendency G = G_D + ξRG_T + ξUG_lnps to solve for divergence δD + (; pres_tend, div_tend) = diagn.tendencies + G = diagn.dynamics.a # reuse work arrays, used for combined tendency G + geopot = diagn.dynamics.b # used for geopotential + + for k in 1:nlayers + for r in k:nlayers # skip 1:k-1 as integration is surface to k + for lm in eachharmonic(temp_tend, div_old, div_new) + # 1. the ξ*R*G_T term, vertical integration of geopotential (excl ξ, this is done in 2.) + geopot[lm, k] += R[k, r]*temp_tend[lm, r] + end end - # alternative way to calculate the geopotential (not thread-safe because geopot from below is reused) - # R is not used here as it's cheaper to reuse the geopotential from k+1 than - # to multiply with the entire upper triangular matrix R which recalculates - # the geopotential for k from all lower levels k...nlev - # TODO swap sign here and not in + geopot[lm] further down - # if k == nlev - # @. geopot = Δp_geopot_full[k]*temp_tend # surface geopotential without orography - # else - # temp_tend_k1 = diagn.layers[k+1].tendencies.temp_tend # temp tendency from layer below - # geopot_k1 = diagn.layers[k+1].dynamics_variables.b # geopotential from layer below - # @. geopot = geopot_k1 + Δp_geopot_half[k+1]*temp_tend_k1 + Δp_geopot_full[k]*temp_tend - # end - # 2. the G = G_D + ξRG_T + ξUG_lnps terms using geopot from above lm = 0 - @inbounds for m in 1:trunc+1 # loops over all columns/order m + 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 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]) + G[lm, k] = div_tend[lm, k] + ξ*eigenvalue*(U[k]*pres_tend[lm] + geopot[lm, k]) + + # div_tend is now in G, fill with zeros here so that it can be used as an accumulator + # in the δD = S⁻¹G calculation below + div_tend[lm, k] = 0 end lm += 1 # skip last row, LowerTriangularMatrices are of size lmax+2 x mmax+1 end - - # div_tend is now in G, fill with zeros here so that it can be used as an accumulator - # in the δD = S⁻¹G calculation below - fill!(div_tend, 0) end # NOW SOLVE THE δD = S⁻¹G to correct divergence tendency - @floop for k in 1:nlev - (; div_tend) = diagn.layers[k].tendencies # unpack div_tend - - @inbounds for r in 1:nlev - G = diagn.layers[r].dynamics_variables.a # reuse work arrays - + for k in eachmatrix(div_tend, G) + for r in eachmatrix(div_tend, G) lm = 0 - 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) + 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] + div_tend[lm, k] += S⁻¹[l, k, r]*G[lm, r] end lm += 1 # skip last row, LowerTriMatrices are of size lmax+2 x mmax+1 end @@ -396,16 +377,17 @@ function implicit_correction!( end # SEMI IMPLICIT CORRECTIONS FOR PRESSURE AND TEMPERATURE, insert δD to get δT, δlnpₛ - @floop for k in 1:nlev - (; temp_tend) = diagn.layers[k].tendencies # unpack temp_tend - @inbounds for r in 1:nlev - (; div_tend) = diagn.layers[r].tendencies # unpack div_tend - @. temp_tend += ξ*L[k, r]*div_tend # δT = G_T + ξLδD + for k in eachmatrix(div_tend, temp_tend) + for r in eachmatrix(div_tend, temp_tend) + for lm in eachharmonic(div_tend, temp_tend) + # δT = G_T + ξLδD + temp_tend[lm, k] += ξ*L[k, r]*div_tend[lm, r] + end end - end - for k in 1:nlev # not thread safe - (; div_tend) = diagn.layers[k].tendencies # unpack div_tend - @. pres_tend += ξ*W[k]*div_tend # δlnpₛ = G_lnpₛ + ξWδD + for lm in eachharmonic(div_tend, temp_tend) + # δlnpₛ = G_lnpₛ + ξWδD + pres_tend[lm] += ξ*W[k]*div_tend[lm, k] + end end end \ No newline at end of file diff --git a/src/dynamics/initial_conditions.jl b/src/dynamics/initial_conditions.jl index 442892b92..c43e26eeb 100644 --- a/src/dynamics/initial_conditions.jl +++ b/src/dynamics/initial_conditions.jl @@ -1,7 +1,7 @@ abstract type AbstractInitialConditions <: AbstractModelComponent end export InitialConditions -Base.@kwdef struct InitialConditions{V,P,T,H} <: AbstractInitialConditions +@kwdef struct InitialConditions{V,P,T,H} <: AbstractInitialConditions vordiv::V = ZeroInitially() pres::P = ZeroInitially() temp::T = ZeroInitially() @@ -11,7 +11,7 @@ end function initialize!( progn::PrognosticVariables, IC::InitialConditions, - model::ModelSetup + model::AbstractModel ) has(model, :vor) && initialize!(progn, IC.vordiv, model) has(model, :pres) && initialize!(progn, IC.pres, model) @@ -38,18 +38,18 @@ end export ZeroInitially struct ZeroInitially <: AbstractInitialConditions end -initialize!(::PrognosticVariables,::ZeroInitially,::ModelSetup) = nothing +initialize!(::PrognosticVariables,::ZeroInitially,::AbstractModel) = nothing # to avoid a breaking change, like ZeroInitially export StartFromRest struct StartFromRest <: AbstractInitialConditions end -initialize!(::PrognosticVariables,::StartFromRest,::ModelSetup) = nothing +initialize!(::PrognosticVariables,::StartFromRest,::AbstractModel) = nothing export StartWithRandomVorticity """Start with random vorticity as initial conditions $(TYPEDFIELDS)""" -Base.@kwdef mutable struct StartWithRandomVorticity <: AbstractInitialConditions +@kwdef mutable struct StartWithRandomVorticity <: AbstractInitialConditions "Power of the spectral distribution k^power" power::Float64 = -3 @@ -62,30 +62,31 @@ $(TYPEDSIGNATURES) Start with random vorticity as initial conditions""" function initialize!( progn::PrognosticVariables{NF}, initial_conditions::StartWithRandomVorticity, - model::ModelSetup) where NF + model::AbstractModel) where NF - lmax = progn.trunc+1 + lmax = progn.trunc + 1 power = initial_conditions.power + 1 # +1 as power is summed of orders m - ξ = randn(Complex{NF}, lmax, lmax)*convert(NF, initial_conditions.amplitude) + + ξ = randn(LowerTriangularArray{Complex{NF}}, lmax, lmax, 1)*convert(NF, initial_conditions.amplitude) + ξ[1] = 0 # don't perturb l=m=0 mode to have zero mean - for progn_layer in progn.layers - for m in 1:lmax - for l in m:lmax - progn_layer.timesteps[1].vor[l, m] = ξ[l, m]*l^power - end + for m in 1:lmax + for l in m:lmax + ξ[l, m, 1] *= l^power end - # don't perturb l=m=0 mode to have zero mean - progn_layer.timesteps[1].vor[1] = 0 end + + ξ = repeat(ξ, 1, model.spectral_grid.nlayers) # repeat for each level to have the same IC across all vertical layers + + set!(progn, model; vor=ξ, lf=1) end export ZonalJet -""" -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. +"""A struct that contains all parameters for the Galewsky et al, 2004 zonal jet +intitial conditions for the ShallowWaterModel. Default values as in Galewsky. $(TYPEDFIELDS)""" -Base.@kwdef mutable struct ZonalJet <: AbstractInitialConditions +@kwdef mutable struct ZonalJet <: AbstractInitialConditions "jet latitude [˚N]" latitude::Float64 = 45 @@ -116,7 +117,7 @@ $(TYPEDSIGNATURES) Initial conditions from Galewsky, 2004, Tellus""" function initialize!( progn::PrognosticVariables, initial_conditions::ZonalJet, - model::ModelSetup) + model::AbstractModel) (; latitude, width, umax) = initial_conditions # for jet (; perturb_lat, perturb_lon, perturb_xwidth, # for perturbation @@ -132,15 +133,15 @@ function initialize!( progn::PrognosticVariables, λ = perturb_lon*2π/360 # perturbation longitude [radians] (; rotation, gravity) = model.planet - (; Grid, NF, radius, nlat_half) = model.spectral_grid - (; coslat⁻¹) = model.geometry + (; Grid, NF, nlat_half) = model.spectral_grid + (; coslat⁻¹, radius) = model.geometry - u_grid = zeros(Grid{NF}, nlat_half) + u_grid = zeros(Grid{NF}, nlat_half, 1) η_perturb_grid = zeros(Grid{NF}, nlat_half) lat = RingGrids.get_lat(Grid, nlat_half) _, lons = RingGrids.get_colatlons(Grid, nlat_half) - for (j, ring) in enumerate(eachring(u_grid, η_perturb_grid)) + for (j, ring) in enumerate(eachring(u_grid)) θ = lat[j] # latitude in radians # velocity per latitude @@ -167,40 +168,40 @@ function initialize!( progn::PrognosticVariables, # 0 = -∇⋅((ζ+f)*(-v, u)) - ∇²((u^2 + v^2)/2), i.e. # invert the Laplacian for # ∇²(gη) = -∇⋅(0, (ζ+f)*u) - ∇²(u^2/2) - u = spectral(u_grid, model.spectral_transform) + u = transform(u_grid, model.spectral_transform) # get vorticity initial conditions from curl of u, v v = zero(u) # meridional velocity zero for these initial conditions - (; vor) = progn.layers[end].timesteps[1] + vor = progn.vor[1] curl!(vor, u, v, model.spectral_transform) - # compute the div = -∇⋅(0,(ζ+f)*u) = -∇×((ζ+f)*u, 0) term, v=0 - vor_grid = gridded(vor, model.spectral_transform) + # compute the div = -∇⋅(0,(ζ+f)*u) = ∇×((ζ+f)*u, 0) term, v=0 + vor_grid = transform(vor, model.spectral_transform) f = coriolis(vor_grid; rotation) # includes 1/coslat/radius from above for curl! # but *radius^2 for the ∇⁻²! operation below! vor_flux_grid = @. (vor_grid + f) * u_grid * radius^2 - vor_flux = spectral(vor_flux_grid, model.spectral_transform) - div = zero(v) - curl!(div, vor_flux, v, model.spectral_transform) + vor_flux = transform(vor_flux_grid, model.spectral_transform) + div = curl(vor_flux, v, model.spectral_transform) # compute the -∇²(u^2/2) term, add to div, divide by gravity RingGrids.scale_coslat!(u_grid) # remove coslat scaling u_grid .*= radius # no radius scaling as we'll apply ∇⁻²(∇²) (would cancel) @. u_grid = convert(NF,1/2) * u_grid^2 - u²_half = spectral!(u, u_grid, model.spectral_transform) + u²_half = transform!(u, u_grid, model.spectral_transform) ∇²!(div, u²_half, model.spectral_transform, flipsign=true, add=true) div .*= inv(gravity) # invert Laplacian to obtain η - (; pres) = progn.surface.timesteps[1] + pres = progn.pres[1] ∇⁻²!(pres, div, model.spectral_transform) - # add perturbation - η_perturb = spectral!(u, η_perturb_grid, model.spectral_transform) + # add perturbation (reuse u array) + η_perturb = transform(η_perturb_grid, model.spectral_transform) pres .+= η_perturb spectral_truncation!(pres) + return nothing end export ZonalWind @@ -210,7 +211,7 @@ $(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. $(TYPEDFIELDS)""" -Base.@kwdef struct ZonalWind <: AbstractInitialConditions +@kwdef struct ZonalWind <: AbstractInitialConditions "conversion from σ to Jablonowski's ηᵥ-coordinates" η₀::Float64 = 0.252 @@ -240,17 +241,17 @@ function initialize!( progn::PrognosticVariables{NF}, (; u₀, η₀) = initial_conditions (; perturb_lat, perturb_lon, perturb_uₚ, perturb_radius) = initial_conditions - (; radius, Grid, nlat_half) = model.spectral_grid + (; radius, Grid, nlat_half, nlayers) = model.spectral_grid (; σ_levels_full) = model.geometry φ, λ = model.geometry.latds, model.geometry.londs S = model.spectral_transform # VORTICITY - ζ = zeros(Grid{NF}, nlat_half) # relative vorticity - D = zeros(Grid{NF}, nlat_half) # divergence (perturbation only) + vor_grid = zeros(Grid{NF}, nlat_half, nlayers) # relative vorticity + div_grid = zeros(Grid{NF}, nlat_half, nlayers) # divergence (perturbation only) - for (k, layer) in enumerate(progn.layers) + for k in eachgrid(vor_grid, div_grid) η = σ_levels_full[k] # Jablonowski and Williamson use η for σ coordinates ηᵥ = (η - η₀)*π/2 # auxiliary variable for vertical coordinate @@ -264,7 +265,7 @@ function initialize!( progn::PrognosticVariables{NF}, tanφ = tand(φij) # Jablonowski and Williamson, eq. (3) - ζ[ij] = -4u₀/radius*cos_ηᵥ*sinφ*cosφ*(2 - 5sinφ^2) + vor_grid[ij, k] = -4u₀/radius*cos_ηᵥ*sinφ*cosφ*(2 - 5sinφ^2) # PERTURBATION sinφc = sind(perturb_lat) # location of centre @@ -279,19 +280,17 @@ function initialize!( progn::PrognosticVariables{NF}, exp_decay = exp(-(r/R)^2) # Jablonowski and Williamson, eq. (12) - ζ[ij] += perturb_uₚ/radius*exp_decay* + vor_grid[ij, k] += perturb_uₚ/radius*exp_decay* (tanφ - 2*(radius/R)^2*acos(X)*X_norm*(sinφc*cosφ - cosφc*sinφ*cosd(λij-λc))) # Jablonowski and Williamson, eq. (13) - D[ij] = -2perturb_uₚ*radius/R^2 * exp_decay * acos(X) * X_norm * cosφc*sind(λij-λc) + div_grid[ij, k] = -2perturb_uₚ*radius/R^2 * exp_decay * acos(X) * X_norm * cosφc*sind(λij-λc) end - - (; vor, div) = layer.timesteps[1] - spectral!(vor, ζ, S) - spectral!(div, D, S) - spectral_truncation!(vor) - spectral_truncation!(div) end + + set!(progn, model; vor=vor_grid, div=vor_grid, lf=1) + + return nothing end export JablonowskiTemperature @@ -301,7 +300,7 @@ $(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. $(TYPEDFIELDS)""" -Base.@kwdef struct JablonowskiTemperature <: AbstractInitialConditions +@kwdef struct JablonowskiTemperature <: AbstractInitialConditions "conversion from σ to Jablonowski's ηᵥ-coordinates" η₀::Float64 = 0.252 @@ -323,13 +322,13 @@ $(TYPEDSIGNATURES) Initial conditions from Jablonowski and Williamson, 2006, QJR Meteorol. Soc""" function initialize!( progn::PrognosticVariables{NF}, initial_conditions::JablonowskiTemperature, - model::ModelSetup) where NF + model::AbstractModel) where NF (;u₀, η₀, ΔT, Tmin) = initial_conditions (;σ_tropopause) = initial_conditions lapse_rate = model.atmosphere.moist_lapse_rate (;temp_ref, R_dry) = model.atmosphere - (;radius, Grid, nlat_half, nlev) = model.spectral_grid + (;radius, Grid, nlat_half, nlayers) = model.spectral_grid (;rotation, gravity) = model.planet (;σ_levels_full) = model.geometry @@ -338,7 +337,7 @@ function initialize!( progn::PrognosticVariables{NF}, # vertical profile Tη = zero(σ_levels_full) - for k in 1:nlev + for k in 1:nlayers σ = σ_levels_full[k] Tη[k] = temp_ref*σ^(R_dry*lapse_rate/gravity) # Jablonowski and Williamson eq. 4 @@ -348,12 +347,10 @@ function initialize!( progn::PrognosticVariables{NF}, end Tη .= max.(Tη, Tmin) - - T = zeros(Grid{NF}, nlat_half) # temperature + temp_grid = zeros(Grid{NF}, nlat_half, nlayers) # temperature aΩ = radius*rotation - for (k, layer) in enumerate(progn.layers) - + for k in eachgrid(temp_grid) η = σ_levels_full[k] # Jablonowski and Williamson use η for σ coordinates ηᵥ = (η - η₀)*π/2 # auxiliary variable for vertical coordinate @@ -366,13 +363,13 @@ function initialize!( progn::PrognosticVariables{NF}, cosφ = cosd(φij) # Jablonowski and Williamson, eq. (6) - T[ij] = Tη[k] + A1*((-2sinφ^6*(cosφ^2 + 1/3) + 10/63)*A2 + (8/5*cosφ^3*(sinφ^2 + 2/3) - π/4)*aΩ) + temp_grid[ij, k] = Tη[k] + A1*((-2sinφ^6*(cosφ^2 + 1/3) + 10/63)*A2 + (8/5*cosφ^3*(sinφ^2 + 2/3) - π/4)*aΩ) end - - (; temp) = layer.timesteps[1] - spectral!(temp, T, S) - spectral_truncation!(temp) end + + set!(progn, model; temp=temp_grid, lf=1) + + return nothing end export StartFromFile @@ -382,7 +379,7 @@ Restart from a previous SpeedyWeather.jl simulation via the restart file restart Applies interpolation in the horizontal but not in the vertical. restart.jld2 is identified by $(TYPEDFIELDS)""" -Base.@kwdef struct StartFromFile <: AbstractInitialConditions +@kwdef struct StartFromFile <: AbstractInitialConditions "path for restart file" path::String = pwd() @@ -396,7 +393,7 @@ Restart from a previous SpeedyWeather.jl simulation via the restart file restart Applies interpolation in the horizontal but not in the vertical.""" function initialize!( progn_new::PrognosticVariables, initial_conditions::StartFromFile, - model::ModelSetup) + model::AbstractModel) (; path, id ) = initial_conditions @@ -421,22 +418,23 @@ function homogeneous_temperature!( progn::PrognosticVariables, # R_dry: Specific gas constant for dry air [J/kg/K] (; temp_ref, lapse_rate, R_dry) = model.atmosphere (; gravity) = model.planet - (; nlev, σ_levels_full) = model.geometry - (; norm_sphere) = model.spectral_transform # normalization of the l=m=0 spherical harmonic + (; nlayers, σ_levels_full) = model.geometry + (; norm_sphere) = model.spectral_transform # normalization of the l=m=0 spherical harmonic # Lapse rate scaled by gravity [K/m / (m²/s²)] Γg⁻¹ = lapse_rate/gravity - # SURFACE TEMPERATURE (store in k = nlev, but it's actually surface, i.e. k=nlev+1/2) + # SURFACE TEMPERATURE (store in k = nlayers, but it's actually surface, i.e. k=nlayers+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 = progn.temp[1][:,end] # 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 end # Use lapserate and vertical coordinate σ for profile - for k in 1:nlev # k=nlev overwrites the surface temperature + for k in 1:nlayers # k=nlayers overwrites the surface temperature # with lowermost layer temperature temp = progn.layers[k].timesteps[1].temp σₖᴿ = σ_levels_full[k]^(R_dry*Γg⁻¹) # from hydrostatic equation @@ -480,10 +478,10 @@ function initialize!( progn::PrognosticVariables, for ij in eachgridpoint(lnp_grid, orography) lnp_grid[ij] = lnp₀ + log(1 - ΓT⁻¹*orography[ij])/RΓg⁻¹ end + + set!(progn, model; pres=lnp_grid, lf=1) - lnp = progn.surface.timesteps[1].pres - spectral!(lnp, lnp_grid, model.spectral_transform) - spectral_truncation!(lnp) # set lmax+1 row to zero + return nothing end export ConstantPressure @@ -497,43 +495,43 @@ function initialize!( progn::PrognosticVariables, # logarithm of reference surface pressure [log(Pa)] # set the l=m=0 mode, normalize correctly - progn.surface.timesteps[1].pres[1] = log(pres_ref) * norm_sphere + set!(progn, model; pres=log(pres_ref) * norm_sphere) - # set other modes explicitly to zero - progn.surface.timesteps[1].pres[2:end] .= 0 return nothing end export ConstantRelativeHumidity -Base.@kwdef struct ConstantRelativeHumidity <: AbstractInitialConditions +@kwdef struct ConstantRelativeHumidity <: AbstractInitialConditions relhumid_ref::Float64 = 0.7 end function initialize!( progn::PrognosticVariables, IC::ConstantRelativeHumidity, - model::ModelSetup, + model::AbstractModel, ) (; relhumid_ref) = IC - (; nlev, σ_levels_full) = model.geometry - lnpₛ = progn.surface.timesteps[1].pres - pres_grid = gridded(lnpₛ, model.spectral_transform) + (; nlayers, σ_levels_full) = model.geometry + + # get pressure [Pa] on grid + lnpₛ = progn.pres[1] + pres_grid = transform(lnpₛ, model.spectral_transform) pres_grid .= exp.(pres_grid) - for k in 1:nlev - temp_grid = gridded(progn.layers[k].timesteps[1].temp, model.spectral_transform) - humid_grid = zero(temp_grid) + temp_grid = transform(progn.temp[1], model.spectral_transform) + humid_grid = zero(temp_grid) + for k in eachgrid(temp_grid, humid_grid) for ij in eachgridpoint(humid_grid) pₖ = σ_levels_full[k] * pres_grid[ij] - q_sat = saturation_humidity(temp_grid[ij], pₖ, model.clausius_clapeyron) - humid_grid[ij] = relhumid_ref*q_sat + q_sat = saturation_humidity(temp_grid[ij, k], pₖ, model.clausius_clapeyron) + humid_grid[ij, k] = relhumid_ref*q_sat end - - (;humid) = progn.layers[k].timesteps[1] - spectral!(humid, humid_grid, model.spectral_transform) - spectral_truncation!(humid) end + + set!(progn, model; humid=humid_grid, lf=1) + + return nothing end export RandomWaves @@ -541,13 +539,15 @@ export RandomWaves """Parameters for random initial conditions for the interface displacement η in the shallow water equations. $(TYPEDFIELDS)""" -Base.@kwdef struct RandomWaves <: AbstractInitialConditions - # random interface displacement field - A::Float64 = 2000 # amplitude [m] +@kwdef struct RandomWaves <: AbstractInitialConditions + # random interface displacement field + A::Float64 = 2000 # amplitude [m] lmin::Int64 = 10 # minimum wavenumber lmax::Int64 = 20 # maximum wavenumber end +RandomWaves(S::SpectralGrid; kwargs...) = RandomWaves(;kwargs...) + """ $(TYPEDSIGNATURES) Random initial conditions for the interface displacement η @@ -560,17 +560,18 @@ function initialize!( progn::PrognosticVariables{NF}, (; A, lmin, lmax) = initial_conditions (; trunc) = progn - η = progn.surface.timesteps[1].pres - η .= randn(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) + η = randn(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) # zero out other wavenumbers η[1:min(lmin, trunc+2), :] .= 0 η[min(lmax+2, trunc+2):trunc+2, :] .= 0 # scale to amplitude - η_grid = gridded(η, model.spectral_transform) + η_grid = transform(η, model.spectral_transform) η_min, η_max = extrema(η_grid) η .*= (A/max(abs(η_min), abs(η_max))) + set!(progn, model; pres=η) + return nothing end \ No newline at end of file diff --git a/src/dynamics/orography.jl b/src/dynamics/orography.jl index 3e5692888..297a2ad98 100644 --- a/src/dynamics/orography.jl +++ b/src/dynamics/orography.jl @@ -26,7 +26,7 @@ function (::Type{Orography})( end # no further initialization needed -initialize!(::NoOrography, ::ModelSetup) = nothing +initialize!(::NoOrography, ::AbstractModel) = nothing export ZonalRidge @@ -50,7 +50,7 @@ end # function barrier function initialize!( orog::ZonalRidge, - model::ModelSetup) + model::AbstractModel) initialize!(orog, model.planet, model.spectral_transform, model.geometry) end @@ -83,7 +83,7 @@ function initialize!( orog::ZonalRidge, 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 + transform!(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 return nothing @@ -130,7 +130,7 @@ end # function barrier function initialize!( orog::EarthOrography, - model::ModelSetup) + model::AbstractModel) initialize!(orog, model.planet, model.spectral_transform) end @@ -160,18 +160,20 @@ function initialize!( orog::EarthOrography, # Interpolate/coarsen to desired resolution interpolate!(orography, orography_highres) - orography .*= scale # scale orography (default 1) - spectral!(geopot_surf, orography, S) # no *gravity yet + orography .*= scale # scale orography (default 1) + transform!(geopot_surf, orography, S) # no *gravity yet - if orog.smoothing # smooth orography in spectral space? - trunc = (size(geopot_surf, 1, as=Matrix) - 2) # get trunc=lmax from size of geopot_surf - truncation = round(Int, trunc * (1-orog.smoothing_fraction)) # degree of harmonics to be truncated + if orog.smoothing # smooth orography in spectral space? + # get trunc=lmax from size of geopot_surf + trunc = (size(geopot_surf, 1, as=Matrix) - 2) + # degree of harmonics to be truncated + truncation = round(Int, trunc * (1-orog.smoothing_fraction)) c = orog.smoothing_strength power = orog.smoothing_power SpeedyTransforms.spectral_smoothing!(geopot_surf, c; power, truncation) end - gridded!(orography, geopot_surf, S) # to grid-point space + transform!(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 return nothing diff --git a/src/dynamics/particle_advection.jl b/src/dynamics/particle_advection.jl index fe648a7c5..f79079926 100644 --- a/src/dynamics/particle_advection.jl +++ b/src/dynamics/particle_advection.jl @@ -4,7 +4,7 @@ abstract type AbstractParticleAdvection <: AbstractModelComponent end export NoParticleAdvection struct NoParticleAdvection <: AbstractParticleAdvection end NoParticleAdvection(::SpectralGrid) = NoParticleAdvection() -initialize!(::NoParticleAdvection, ::ModelSetup) = nothing +initialize!(::NoParticleAdvection, ::AbstractModel) = nothing initialize!(particles, progn, diagn, ::NoParticleAdvection) = nothing particle_advection!(progn, diagn, ::NoParticleAdvection) = nothing @@ -21,17 +21,17 @@ Base.@kwdef struct ParticleAdvection2D{NF} <: AbstractParticleAdvection end function ParticleAdvection2D(SG::SpectralGrid; kwargs...) - SG.n_particles == 0 && @warn "ParticleAdvection2D created but n_particles = 0 in spectral grid." + SG.nparticles == 0 && @warn "ParticleAdvection2D created but nparticles = 0 in spectral grid." ParticleAdvection2D{SG.NF}(; kwargs...) end function initialize!( particle_advection::ParticleAdvection2D, - model::ModelSetup, + model::AbstractModel, ) - (; nlev) = model.spectral_grid + (; nlayers) = model.spectral_grid (; layer) = particle_advection - nlev < layer && @warn "Particle advection on layer $layer on spectral grid with nlev=$nlev." + nlayers < layer && @warn "Particle advection on layer $layer on spectral grid with nlayers=$nlayers." (; every_n_timesteps) = particle_advection # Δt [˚*s/m] is scaled by radius to convert more easily from velocity [m/s] @@ -47,8 +47,8 @@ vertical σ coordinates. This uses a cosin-distribution in latitude for an equal-area uniformity.""" function initialize!( particles::Vector{P}, - model::ModelSetup, -) where {P<:Particle} + model::AbstractModel, +) where {P <: Particle} for i in eachindex(particles) # uniform random in lon (360*rand), lat (cos-distribution), σ (rand) particles[i] = rand(P) @@ -69,7 +69,8 @@ function initialize!( length(particles) == 0 && return nothing k = particle_advection.layer - (; u_grid, v_grid) = diagn.layers[k].grid_variables + u_grid = view(diagn.grid.u_grid, :, k) + v_grid = view(diagn.grid.v_grid, :, k) (; interpolator) = diagn.particles # interpolate initial velocity on initial locations @@ -91,7 +92,6 @@ function initialize!( interpolate!(v0, v_grid, interpolator) end - # function barrier function particle_advection!(progn, diagn, adv::ParticleAdvection2D) particle_advection!(progn.particles, diagn, progn.clock, adv) @@ -110,7 +110,7 @@ function particle_advection!( # decide whether to execute on this time step: # execute always on last time step *before* time step is divisible by # `particle_advection.every_n_timesteps`, e.g. 7, 15, 23, ... for n=8 which - # already contains u, v at i=8, 16, 24, etc as executed after `gridded!` + # already contains u, v at i=8, 16, 24, etc as executed after `transform!` # even though the clock hasn't be step forward yet, this means time = time + Δt here # should not be called on the 1st step in first_timesteps, which is excluded @@ -157,7 +157,8 @@ function particle_advection!( # CORRECTOR STEP, use u, v at new location and new time step k = particle_advection.layer - (; u_grid, v_grid) = diagn.layers[k].grid_variables + u_grid = view(diagn.grid.u_grid, :, k) + v_grid = view(diagn.grid.v_grid, :, k) (; interpolator) = diagn.particles RingGrids.update_locator!(interpolator, lats, lons) diff --git a/src/dynamics/particles.jl b/src/dynamics/particles.jl index a9a7dae6e..4a932dff7 100644 --- a/src/dynamics/particles.jl +++ b/src/dynamics/particles.jl @@ -43,6 +43,10 @@ Base.zero(::Type{Particle}) = Particle{DEFAULT_NF,true}(0,0,0) Base.zero(::Type{Particle{NF}}) where NF = Particle{NF,true}(0,0,0) Base.zero(::Type{Particle{NF,isactive}}) where {NF,isactive} = Particle{NF,isactive}(0,0,0) Base.zero(::P) where {P<:Particle} = zero(P) +function Base.zeros(ArrayType::Type{<:AbstractArray{P}}, n::Int...) where {P<:Particle} + z = ArrayType(undef, n...) + fill!(z, zero(P)) +end Base.rand(rng::Random.AbstractRNG, ::Random.Sampler{Particle}) = rand(rng, Particle{DEFAULT_NF,true}) Base.rand(rng::Random.AbstractRNG, ::Random.Sampler{Particle{NF}}) where NF = rand(rng, Particle{NF,true}) diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl index 951b88dc7..9efa12dc4 100644 --- a/src/dynamics/prognostic_variables.jl +++ b/src/dynamics/prognostic_variables.jl @@ -1,411 +1,449 @@ -# how many time steps have to be stored for the time integration? Leapfrog = 2 -const N_STEPS = 2 -const LTM = LowerTriangularMatrix # just because it's shorter here -export PrognosticVariablesLayer - -"""A layer of the prognostic variables in spectral space. -$(TYPEDFIELDS)""" -Base.@kwdef struct PrognosticVariablesLayer{NF<:AbstractFloat} <: AbstractVariables - - "Spectral resolution as max degree of spherical harmonics" - trunc::Int - - "Vorticity of horizontal wind field [1/s]" - vor ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) - - "Divergence of horizontal wind field [1/s]" - div ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) - - "Absolute temperature [K]" - temp ::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) - - "Specific humidity [kg/kg]" - humid::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) -end - -# generator function based on spectral grid -PrognosticVariablesLayer(SG::SpectralGrid) = PrognosticVariablesLayer{SG.NF}(trunc=SG.trunc) - function Base.show(io::IO, A::AbstractVariables) println(io, "$(typeof(A))") keys = propertynames(A) print_fields(io, A, keys) end -"""Collect the n time steps of PrognosticVariablesLayer -of an n-step time integration (leapfrog=2) into a single struct. -$(TYPEDFIELDS).""" -struct PrognosticLayerTimesteps{NF<:AbstractFloat} <: AbstractVariables - timesteps::Vector{PrognosticVariablesLayer{NF}} # N_STEPS-element vector for time steps -end - -# generator function based on spectral grid -function PrognosticLayerTimesteps(SG::SpectralGrid) - return PrognosticLayerTimesteps([PrognosticVariablesLayer(SG) for _ in 1:N_STEPS]) -end - -"""The spectral and gridded prognostic variables at the surface. -$(TYPEDFIELDS)""" -Base.@kwdef struct PrognosticVariablesSurface{NF<:AbstractFloat} <: AbstractVariables - - "Spectral resolution as max degree of spherical harmonics" - trunc::Int - - "log of surface pressure [log(Pa)] for PrimitiveEquation, interface displacement [m] for ShallowWaterModel" - pres::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1) -end - -# generator function based on a SpectralGrid -PrognosticVariablesSurface(SG::SpectralGrid) = PrognosticVariablesSurface{SG.NF}(trunc=SG.trunc) - -Base.@kwdef mutable struct PrognosticVariablesOcean{NF<:AbstractFloat, Grid<:AbstractGrid{NF}} <: AbstractVariables - - "Resolution parameter of grid" - const nlat_half::Int +# to be removed +struct PrognosticLayerTimesteps end +struct PrognosticSurfaceTimesteps end +struct PrognosticVariablesLayer end - "Current time of the ocean variables" - time::DateTime = DEFAULT_DATE +export PrognosticVariablesOcean +@kwdef struct PrognosticVariablesOcean{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + GridVariable2D, # <: AbstractGridArray +} <: AbstractPrognosticVariables + # DIMENSION + "Number of latitude rings on one hemisphere (Equator incl.), resolution parameter of grid" + nlat_half::Int - # SEA + # OCEAN VARIABLES "Sea surface temperature [K]" - const sea_surface_temperature::Grid = zeros(Grid, nlat_half) + sea_surface_temperature::GridVariable2D = zeros(GridVariable2D, nlat_half) "Sea ice concentration [1]" - const sea_ice_concentration::Grid = zeros(Grid, nlat_half) -end - -# generator function based on a SpectralGrid -function PrognosticVariablesOcean(SG::SpectralGrid) - (; nlat_half, Grid, NF) = SG - return PrognosticVariablesOcean{NF, Grid{NF}}(; nlat_half) + sea_ice_concentration::GridVariable2D = zeros(GridVariable2D, nlat_half) end -Base.@kwdef mutable struct PrognosticVariablesLand{NF<:AbstractFloat, Grid<:AbstractGrid{NF}} <: AbstractVariables - - "Resolution parameter of grid" - const nlat_half::Int - - "Current time of the land variables" - time::DateTime = DEFAULT_DATE +export PrognosticVariablesLand +@kwdef struct PrognosticVariablesLand{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + GridVariable2D, # <: AbstractGridArray +} <: AbstractPrognosticVariables + # DIMENSION + "Number of latitude rings on one hemisphere (Equator incl.), resolution parameter of grid" + nlat_half::Int - # LAND + # LAND VARIABLES "Land surface temperature [K]" - const land_surface_temperature::Grid = zeros(Grid, nlat_half) + land_surface_temperature::GridVariable2D = zeros(GridVariable2D, nlat_half) "Snow depth [m]" - const snow_depth::Grid = zeros(Grid, nlat_half) + snow_depth::GridVariable2D = zeros(GridVariable2D, nlat_half) "Soil moisture layer 1, volume fraction [1]" - const soil_moisture_layer1::Grid = zeros(Grid, nlat_half) + soil_moisture_layer1::GridVariable2D = zeros(GridVariable2D, nlat_half) "Soil moisture layer 2, volume fraction [1]" - const soil_moisture_layer2::Grid = zeros(Grid, nlat_half) + soil_moisture_layer2::GridVariable2D = zeros(GridVariable2D, nlat_half) end -# generator function based on a SpectralGrid -function PrognosticVariablesLand(SG::SpectralGrid) - (; nlat_half, Grid, NF) = SG - return PrognosticVariablesLand{NF, Grid{NF}}(; nlat_half) -end +export PrognosticVariables +@kwdef struct PrognosticVariables{ + NF, # <: AbstractFloat + ArrayType, # Array, CuArray, ... + NSTEPS, # number of timesteps + SpectralVariable2D, # <: LowerTriangularArray + SpectralVariable3D, # <: LowerTriangularArray + GridVariable2D, # <: AbstractGridArray + ParticleVector, # <: AbstractVector{Particle{NF}} +} <: AbstractPrognosticVariables -"""Collect the n time steps of PrognosticVariablesSurface -of an n-step time integration (leapfrog=2) into a single struct. -$(TYPEDFIELDS).""" -struct PrognosticSurfaceTimesteps{NF<:AbstractFloat} <: AbstractVariables - timesteps::Vector{PrognosticVariablesSurface{NF}} # N_STEPS-element vector for time steps -end + # DIMENSIONS + "max degree of spherical harmonics (0-based)" + trunc::Int -# generator function based on spectral grid -function PrognosticSurfaceTimesteps(SG::SpectralGrid) - return PrognosticSurfaceTimesteps([PrognosticVariablesSurface(SG) for _ in 1:N_STEPS]) -end + "Number of latitude rings on one hemisphere (Equator excl.), resolution parameter of grids" + nlat_half::Int -export PrognosticVariables -struct PrognosticVariables{ - NF<:AbstractFloat, - Grid<:AbstractGrid{NF}, - M<:ModelSetup -} <: AbstractPrognosticVariables + "number of vertical layers" + nlayers::Int - # dimensions - trunc::Int # max degree of spherical harmonics - nlat_half::Int # resolution parameter of grids - nlev::Int # number of vertical levels - n_steps::Int # N_STEPS time steps that are stored + "Number of particles for particle advection" + nparticles::Int - layers::Vector{PrognosticLayerTimesteps{NF}} # vector of vertical layers - surface::PrognosticSurfaceTimesteps{NF} - ocean::PrognosticVariablesOcean{NF, Grid} - land::PrognosticVariablesLand{NF, Grid} - particles::Vector{Particle{NF}} + # LAYERED VARIABLES + "Vorticity of horizontal wind field [1/s], but scaled by scale (=radius during simulation)" + vor::NTuple{NSTEPS, SpectralVariable3D} = + ntuple(i -> zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers), NSTEPS) - # scaling - scale::Base.RefValue{NF} + "Divergence of horizontal wind field [1/s], but scaled by scale (=radius during simulation)" + div::NTuple{NSTEPS, SpectralVariable3D} = + ntuple(i -> zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers), NSTEPS) - clock::Clock -end + "Absolute temperature [K]" + temp::NTuple{NSTEPS, SpectralVariable3D} = + ntuple(i -> zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers), NSTEPS) + + "Specific humidity [kg/kg]" + humid::NTuple{NSTEPS, SpectralVariable3D} = + ntuple(i -> zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers), NSTEPS) -function PrognosticVariables(SG::SpectralGrid, model::ModelSetup) + "Logarithm of surface pressure [log(Pa)] for PrimitiveEquation, interface displacement [m] for ShallowWaterModel" + pres::NTuple{NSTEPS, SpectralVariable2D} = + ntuple(i -> zeros(SpectralVariable2D, trunc+2, trunc+1), NSTEPS) + + "Ocean variables, sea surface temperature and sea ice concentration" + ocean::PrognosticVariablesOcean{NF, ArrayType, GridVariable2D} = + PrognosticVariablesOcean{NF, ArrayType, GridVariable2D}(; nlat_half) - (; trunc, nlat_half, nlev, Grid, NF) = SG - (; n_particles) = SG + "Land variables, land surface temperature, snow and soil moisture" + land::PrognosticVariablesLand{NF, ArrayType, GridVariable2D} = + PrognosticVariablesLand{NF, ArrayType, GridVariable2D}(; nlat_half) + + "Particles for particle advection" + particles::ParticleVector = zeros(ParticleVector, nparticles) - # data structs - layers = [PrognosticLayerTimesteps(SG) for _ in 1:nlev] # vector of nlev layers - surface = PrognosticSurfaceTimesteps(SG) - ocean = PrognosticVariablesOcean(SG) - land = PrognosticVariablesLand(SG) + "Scaling for vor, div. scale=1 outside simulation, =radius during simulation" + scale::Base.RefValue{NF} = Ref(one(NF)) - # particles advection - particles = zeros(Particle{NF}, n_particles) + "Clock that keeps track of time, number of timesteps to integrate for." + clock::Clock = Clock() +end - scale = Ref(one(NF)) # initialize with scale=1, wrapped in RefValue for mutability - clock = Clock() +"""$(TYPEDSIGNATURES) +Generator function.""" +function PrognosticVariables(SG::SpectralGrid; nsteps=DEFAULT_NSTEPS) + (; trunc, nlat_half, nlayers, nparticles) = SG + (; NF, ArrayType) = SG + (; SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector) = SG + + return PrognosticVariables{NF, ArrayType, nsteps, + SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}(; + trunc, nlat_half, nlayers=nlayers, nparticles, + ) +end - Model = model_class(model) # strip away the parameters - return PrognosticVariables{NF, Grid{NF}, Model}(trunc, nlat_half, nlev, N_STEPS, - layers, surface, ocean, land, particles, - scale, clock) +"""$(TYPEDSIGNATURES) +Generator function.""" +function PrognosticVariables(SG::SpectralGrid, model::AbstractModel) + PrognosticVariables(SG, nsteps = model.time_stepping.nsteps) end -has(::PrognosticVariables{NF, Grid, M}, var_name::Symbol) where {NF, Grid, M} = has(M, var_name) +function Base.show( + io::IO, + progn::PrognosticVariables{NF, ArrayType, NSTEPS}, +) where {NF, ArrayType, NSTEPS} + Grid = typeof(progn.ocean.sea_surface_temperature) + println(io, "PrognosticVariables{$NF, $ArrayType}") + + # variables + (; trunc, nlat_half, nlayers, nparticles) = progn + nlat = RingGrids.get_nlat(Grid, nlat_half) + println(io, "├ vor: T$trunc, $nlayers-layer, $NSTEPS-steps LowerTriangularArray{$NF}") + println(io, "├ div: T$trunc, $nlayers-layer, $NSTEPS-steps LowerTriangularArray{$NF}") + println(io, "├ temp: T$trunc, $nlayers-layer, $NSTEPS-steps LowerTriangularArray{$NF}") + println(io, "├ humid: T$trunc, $nlayers-layer, $NSTEPS-steps LowerTriangularArray{$NF}") + println(io, "├ pres: T$trunc, 1-layer, $NSTEPS-steps LowerTriangularArray{$NF}") + println(io, "├┐ocean: PrognosticVariablesOcean{$NF}") + println(io, "│├ sea_surface_temperature: $nlat-ring $Grid") + println(io, "│└ sea_ice_concentration: $nlat-ring $Grid") + println(io, "├┐land: PrognosticVariablesLand{$NF}") + println(io, "│├ land_surface_temperature: $nlat-ring $Grid") + println(io, "│├ snow_depth: $nlat-ring $Grid") + println(io, "│├ soil_moisture_layer1: $nlat-ring $Grid") + println(io, "│└ soil_moisture_layer2: $nlat-ring $Grid") + println(io, "├ particles: $nparticles-element $(typeof(progn.particles))") + println(io, "├ scale: $(progn.scale[])") + print(io, "└ clock: $(progn.clock.time)") +end -""" - copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariables) +# has(::PrognosticVariables{NF, Grid, M}, var_name::Symbol) where {NF, Grid, M} = has(M, var_name) -Copies entries of `progn_old` into `progn_new`. Only copies those variables that are present -in the model of both `progn_new` and `progn_old`. -""" +"""$(TYPEDSIGNATURES) +Copies entries of `progn_old` into `progn_new`.""" function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariables) - var_names = propertynames(progn_old.layers[1].timesteps[1]) - - for var_name in var_names - if has(progn_new, var_name) - var = get_var(progn_old, var_name) - set_var!(progn_new, var_name, var) - end - end - pres = get_pressure(progn_old) - set_pressure!(progn_new, pres) + for i in eachindex(progn_new.vor) # each leapfrog time step + progn_new.vor[i] .= progn_old.vor[i] + progn_new.div[i] .= progn_old.div[i] + progn_new.temp[i] .= progn_old.temp[i] + progn_new.humid[i] .= progn_old.humid[i] + progn_new.pres[i] .= progn_old.pres[i] + end + + # ocean + progn_new.ocean.sea_surface_temperature .= progn_old.ocean.sea_surface_temperature + progn_new.ocean.sea_ice_concentration .= progn_old.ocean.sea_ice_concentration - # synchronize the clock + # land + progn_new.land.land_surface_temperature .= progn_old.land.land_surface_temperature + progn_new.land.snow_depth .= progn_old.land.snow_depth + progn_new.land.soil_moisture_layer1 .= progn_old.land.soil_moisture_layer1 + progn_new.land.soil_moisture_layer2 .= progn_old.land.soil_moisture_layer2 + + # copy largest subset of particles + if length(progn_new.particles) != length(progn_old.particles) + nnew = length(progn_new.particles) + nold = length(progn_old.particles) + nsub = min(nnew, nold) + @warn "Number of particles changed (origin: $nold, destination: $nnew), copying over only the largest subset ($nsub particles)" + progn_new.particles[1:nsub] .= progn_old.particles[1:nsub] + else + progn_new.particles .= progn_old.particles + end + progn_new.clock.time = progn_old.clock.time + progn_new.scale[] = progn_old.scale[] return progn_new end -# SET_VAR FUNCTIONS TO ASSIGN NEW VALUES TO PrognosticVariables +export set! """ - set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:LowerTriangularMatrix}; - lf::Integer=1) where NF - -Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` -with values given in `var` a vector with all information for all layers in spectral space. +$(TYPEDSIGNATURES) +Sets new values for the keyword arguments (velocities, vorticity, divergence, etc..) into the +prognostic variable struct `progn` at timestep index `lf`. If `add==true` they are added to the +current value instead. If a `SpectralTransform` S is provided, it is used when needed to set +the variable, otherwise it is recomputed. In case `u` and `v` are provied, actually the divergence +and vorticity are set and `coslat_scaling_included` specficies whether or not the 1/cos(lat) +scaling is already included in the arrays or not (default: `false`) + +The input may be: +* A function or callable object `f(lond, latd, σ) -> value` (multilevel variables) +* A function or callable object `f(lond, latd) -> value` (surface level variables) +* An instance of `AbstractGridArray` +* An instance of `LowerTriangularArray` +* A scalar `<: Number` (interpreted as a constant field in grid space) """ -function set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:LowerTriangularMatrix}; - lf::Integer=1) where NF +function set!( + progn::PrognosticVariables, + geometry::Geometry; + u = nothing, + v = nothing, + vor = nothing, + div = nothing, + temp = nothing, + humid = nothing, + pres = nothing, + sea_surface_temperature = nothing, + sea_ice_concentration = nothing, + land_surface_temperature = nothing, + snow_depth = nothing, + soil_moisture_layer1 = nothing, + soil_moisture_layer2 = nothing, + lf::Integer = 1, + add::Bool = false, + S::Union{Nothing, SpectralTransform} = nothing, + coslat_scaling_included::Bool = false, +) + # ATMOSPHERE + isnothing(vor) || set!(progn.vor[lf], vor, geometry, S; add) + isnothing(div) || set!(progn.div[lf], div, geometry, S; add) + isnothing(temp) || set!(progn.temp[lf], temp, geometry, S; add) + isnothing(humid) || set!(progn.humid[lf], humid, geometry, S; add) + isnothing(pres) || set!(progn.pres[lf], pres, geometry, S; add) + + # or provide u, v instead of vor, div + isnothing(u) | isnothing(v) || set_vordiv!(progn.vor[lf], progn.div[lf], u, v, geometry, S; add, coslat_scaling_included) + + # OCEAN + isnothing(sea_surface_temperature) || set!(progn.ocean.sea_surface_temperature, sea_surface_temperature, geometry, S; add) + isnothing(sea_ice_concentration) || set!(progn.ocean.sea_ice_concentration, sea_ice_concentration, geometry, S; add) - @assert length(var) == length(progn.layers) - @assert has(progn, varname) "PrognosticVariables has no variable $varname" + # LAND + isnothing(land_surface_temperature) || set!(progn.land.land_surface_temperature, land_surface_temperature, geometry, S; add) + isnothing(snow_depth) || set!(progn.land.snow_depth, snow_depth, geometry, S; add) + isnothing(soil_moisture_layer1) || set!(progn.land.soil_moisture_layer1, soil_moisture_layer1, geometry, S; add) + isnothing(soil_moisture_layer2) || set!(progn.land.soil_moisture_layer2, soil_moisture_layer2, geometry, S; add) + return nothing +end - for (progn_layer, var_layer) in zip(progn.layers, var) - _set_var_core!(getfield(progn_layer.timesteps[lf], varname), var_layer) +# set LTA <- LTA +function set!(var::LowerTriangularArray{T}, L::LowerTriangularArray, varargs...; add::Bool) where T + if add + if size(var) == size(L) + var .+= T.(L) + else + L_var = spectral_truncation(L, size(var, 1, as=Matrix), size(var, 2, as=Matrix)) + var .+= L_var + end + else + size(var) != size(L) || fill!(var, zero(T)) # copyto! copies over the largest subset, when size(var) > size(L), the copyto! isn't enough by itself + copyto!(var, L) end - - return progn + return var end -function _set_var_core!(var_old::LowerTriangularMatrix{T}, var_new::LowerTriangularMatrix{R}) where {T, R} - lmax, mmax = size(var_old, as=Matrix) .- (1, 1) - var_new_trunc = spectral_truncation!(var_new, mmax+1, mmax) - copyto!(var_old, var_new_trunc) -end - -""" - set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:AbstractGrid}; - lf::Integer=1) where NF - -Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` -with values given in `var` a vector with all information for all layers in grid space. -""" -function set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:AbstractGrid}; - lf::Integer=1) where NF - - @assert length(var) == length(progn.layers) - var_sph = [spectral(var_layer, one_more_degree=true) for var_layer in var] - return set_var!(progn, varname, var_sph; lf=lf) -end +# set LTA <- Grid +function set!(var::LowerTriangularArray, grids::AbstractGridArray, geometry::Union{Geometry, Nothing}=nothing, S::Union{Nothing, SpectralTransform}=nothing; add) + specs = isnothing(S) ? transform(grids) : transform(grids, S) + set!(var, specs; add) +end -""" - set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:AbstractGrid}, - M::ModelSetup; - lf::Integer=1) where NF - -Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` -with values given in `var` a vector with all information for all layers in grid space. -""" -function set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:AbstractGrid}, - M::ModelSetup; - lf::Integer=1) where NF +# set LTA <- func +function set!(var::LowerTriangularArray, f::Function, geometry::Geometry{NF}, S::Union{SpectralTransform, Nothing}=nothing; add::Bool) where NF + grid = ndims(var) == 1 ? zeros(geometry.Grid{NF}, geometry.nlat_half) : zeros(geometry.Grid{NF}, geometry.nlat_half, geometry.nlayers) + set!(grid, f, geometry, S; add=false) + set!(var, grid, geometry, S; add) +end - @assert length(var) == length(progn.layers) +# set LTA <- number +function set!(var::LowerTriangularArray{T}, s::Number, geometry::Geometry{NF}, S::Union{SpectralTransform, Nothing}=nothing; add::Bool) where {T, NF} - var_sph = [spectral(var_layer, M.spectral_transform) for var_layer in var] - - return set_var!(progn, varname, var_sph; lf=lf) -end - -""" - set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:AbstractMatrix}, - Grid::Type{<:AbstractGrid}=FullGaussianGrid; - lf::Integer=1) where NF - -Sets the prognostic variable with the name `varname` in all layers at leapfrog index `lf` -with values given in `var` a vector with all information for all layers in grid space. -""" -function set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - var::Vector{<:AbstractMatrix}, - Grid::Type{<:AbstractGrid}=FullGaussianGrid; - lf::Integer=1) where NF + # appropiate normalization, assume standard 2√π normalisation if no transform is given + norm_sphere = isnothing(S) ? 2sqrt(π) : S.norm_sphere - @assert length(var) == length(progn.layers) + # all elements are zero except for the 0,0 one + var_new = zero(var) - var_grid = [spectral(var_layer; Grid, one_more_degree=true) for var_layer in var] + for k in eachmatrix(var_new) + var_new[1, k] = norm_sphere * s + end - return set_var!(progn, varname, var_grid; lf=lf) + set!(var, var_new, geometry, S; add) end -""" - function set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - s::Number; - lf::Integer=1) where NF - -Sets all values of prognostic variable `varname` at leapfrog index `lf` to the scalar `s`. -""" -function set_var!(progn::PrognosticVariables{NF}, - varname::Symbol, - s::Number; - lf::Integer=1) where NF - - for progn_layer in progn.layers - fill!(getfield(progn_layer.timesteps[lf], varname), s) +# set Grid <- Grid +function set!(var::AbstractGridArray, grids::AbstractGridArray, geometry::Geometry, S::Union{Nothing, SpectralTransform}=nothing; add) + if add + if grids_match(var, grids) + var .+= grids + else + var .+= interpolate(typeof(var), geometry.nlat_half, grids) + end + else + interpolate!(var, grids) end - - return progn + return var end -""" - set_vorticity!(progn::PrognosticVariables, varargs...; kwargs...) - -See [`set_var!`](@ref) -""" -set_vorticity!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :vor, varargs...; kwargs...) - -""" - set_divergence!(progn::PrognosticVariables, varargs...; kwargs...) - -See [`set_var!`](@ref) -""" -set_divergence!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :div, varargs...; kwargs...) - -""" - set_temperature!(progn::PrognosticVariables, varargs...; kwargs...) - -See [`set_var!`](@ref) -""" -set_temperature!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :temp, varargs...; kwargs...) - -""" - set_humidity!(progn::PrognosticVariables, varargs...; kwargs...) - -See [`set_var!`](@ref) -""" -set_humidity!(progn::PrognosticVariables, varargs...; kwargs...) = set_var!(progn, :humid, varargs...; kwargs...) - -""" - set_pressure!(progn::PrognosticVariables{NF}, - pressure::LowerTriangularMatrix; - lf::Integer=1) where NF - -Sets the prognostic variable with the surface pressure in spectral space at leapfrog index `lf`. -""" -function set_pressure!(progn::PrognosticVariables, - pressure::LowerTriangularMatrix; - lf::Integer=1) - - _set_var_core!(progn.surface.timesteps[lf].pres, pressure) - - return progn +# set Grid <- LTA +function set!(var::AbstractGridArray, specs::LowerTriangularArray, geometry::Geometry, S::Union{Nothing, SpectralTransform}=nothing; add) + grids = isnothing(S) ? transform(specs) : transform(specs, S) + set!(var, grids, geometry, S; add) end -""" - set_pressure!(progn::PrognosticVariables{NF}, - pressure::AbstractGrid, - M::ModelSetup; - lf::Integer=1) where NF - -Sets the prognostic variable with the surface pressure in grid space at leapfrog index `lf`. -""" -set_pressure!(progn::PrognosticVariables, pressure::AbstractGrid, M::ModelSetup; lf::Integer=1) = - set_pressure!(progn, spectral(pressure, M.spectral_transform); lf) +# set Grid <- Func +function set!(var::AbstractGridArray, f::Function, geometry::Geometry, S::Union{Nothing, SpectralTransform}=nothing; add) + (; londs, latds, σ_levels_full) = geometry + kernel(a, b) = add ? a+b : b + for k in eachgrid(var) + for ij in eachgridpoint(var) + var[ij, k] = kernel(var[ij, k], f(londs[ij], latds[ij], σ_levels_full[k])) + end + end + return var +end -""" - set_pressure!(progn::PrognosticVariables{NF}, - pressure::AbstractGrid, - lf::Integer=1) where NF +# set Grid (surface/single level) <- Func +function set!( + var::AbstractGridArray{T,1}, + f::Function, + geometry::Geometry, + S::Union{Nothing, SpectralTransform}=nothing; + add +) where T + (; londs, latds) = geometry + kernel(a, b) = add ? a+b : b + for ij in eachgridpoint(var) + var[ij] = kernel(var[ij], f(londs[ij], latds[ij])) + end + return var +end -Sets the prognostic variable with the surface pressure in grid space at leapfrog index `lf`. -""" -set_pressure!(progn::PrognosticVariables, pressure::AbstractGrid; lf::Integer=1) = - set_pressure!(progn, spectral(pressure, one_more_degree=true); lf) +# set Grid <- Number +function set!( + var::AbstractGridArray{T}, + s::Number, + geometry::Union{Geometry, Nothing}=nothing, + S::Union{Nothing, SpectralTransform}=nothing; + add::Bool, +) where T + kernel(a, b) = add ? a+b : b + sT = T(s) + var .= kernel.(var, sT) +end -""" - set_pressure!(progn::PrognosticVariables{NF}, - pressure::AbstractMatrix, - Grid::Type{<:AbstractGrid}, - lf::Integer=1) where NF +# set vor_div <- func +function set_vordiv!( + vor::LowerTriangularArray, + div::LowerTriangularArray, + u_func, + v_func, + geometry::Geometry, + S::Union{Nothing, SpectralTransform}=nothing; + add::Bool, + coslat_scaling_included::Bool=false, +) + u_L = similar(vor) + set!(u_L, u_func, geometry, S) + v_L = similar(vor) + set!(v_L, v_func, geometry, S) + + set_vordiv!(vor, div, u_L, v_L, geometry, S; add, coslat_scaling_included) +end -Sets the prognostic variable with the surface pressure in grid space at leapfrog index `lf`. -""" -set_pressure!(progn::PrognosticVariables, pressure::AbstractMatrix; lf::Integer=1, - Grid::Type{<:AbstractGrid}=FullGaussianGrid) = set_pressure!(progn, spectral(pressure; Grid, one_more_degree=true); lf) - -""" - get_var(progn::PrognosticVariables, var_name::Symbol; lf::Integer=1) +# set vor_div <- grid +function set_vordiv!( + vor::LowerTriangularArray, + div::LowerTriangularArray, + u::AbstractGridArray, + v::AbstractGridArray, + geometry::Geometry, + S::Union{Nothing, SpectralTransform}=nothing; + add::Bool, + coslat_scaling_included::Bool=false, +) + u_ = coslat_scaling_included ? u : RingGrids.scale_coslat⁻¹(u) + v_ = coslat_scaling_included ? v : RingGrids.scale_coslat⁻¹(v) + + u_spec = isnothing(S) ? transform(u_) : transform(u_, S) + v_spec = isnothing(S) ? transform(v_) : transform(v_, S) + + set_vordiv!(vor, div, u_spec, v_spec, geometry, S; add, coslat_scaling_included=true) +end -Returns the prognostic variable `var_name` at leapfrog index `lf` as a `Vector{LowerTriangularMatrices}`. -""" -function get_var(progn::PrognosticVariables, var_name::Symbol; lf::Integer=1) - @assert has(progn, var_name) "PrognosticVariables has no variable $var_name" - return [getfield(layer.timesteps[lf], var_name) for layer in progn.layers] +# set vor_div <- LTA +function set_vordiv!( + vor::LowerTriangularArray, + div::LowerTriangularArray, + u::LowerTriangularArray, + v::LowerTriangularArray, + geometry::Geometry, + S::Union{Nothing, SpectralTransform}=nothing; + add::Bool, + coslat_scaling_included::Bool=false, +) + S = isnothing(S) ? SpectralTransform(geometry.spectral_grid) : S + + u_ = coslat_scaling_included ? u : transform(RingGrids.scale_coslat⁻¹(transform(u, S)), S) + v_ = coslat_scaling_included ? v : transform(RingGrids.scale_coslat⁻¹(transform(u, S)), S) + + if size(vor) != size(u_) != size(v_) + u_new = zero(vor) + copyto!(u_new, u_) + + v_new = zero(vor) + copyto!(v_new, v_) + + curl!(vor, u_new, v_new, S; add, radius=geometry.radius) + divergence!(div, u_new, v_new, S; add, radius=geometry.radius) + else + curl!(vor, u_, v_, S; add, radius=geometry.radius) + divergence!(div, u_, v_, S; add, radius=geometry.radius) + end end -get_vorticity(progn::PrognosticVariables; kwargs...) = get_var(progn, :vor; kwargs...) -get_divergence(progn::PrognosticVariables; kwargs...) = get_var(progn, :div; kwargs...) -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 set!(S::AbstractSimulation; kwargs...) + set!(S.prognostic_variables, S.model.geometry; S=S.model.spectral_transform, kwargs...) +end -function Base.show(io::IO, P::PrognosticVariables) - ζ = P.layers[end].timesteps[1].vor # create a view on surface relative vorticity - ζ_grid = gridded(ζ) # to grid space - print(io, plot(ζ_grid, title="Surface relative vorticity")) +function set!(progn::PrognosticVariables, model::AbstractModel; kwargs...) + progn.scale[] != 1 && @warn "Prognostic variables are scaled with $(progn.scale[]), but `set!` assumes unscaled variables." + set!(progn, model.geometry; S=model.spectral_transform, kwargs...) end \ No newline at end of file diff --git a/src/dynamics/scaling.jl b/src/dynamics/scaling.jl index 801f05508..3e2cd8104 100644 --- a/src/dynamics/scaling.jl +++ b/src/dynamics/scaling.jl @@ -7,17 +7,9 @@ function scale!( var::Symbol, scale::Real, ) - if var == :pres - for pres in progn.pres.timesteps - pres .*= scale - end - else - for layer in progn.layers - for step in layer.timesteps - variable = getfield(step, var) - variable .*= scale - end - end + variable_nsteps = getfield(progn, var) + for step in variable_nsteps + step .*= scale end end @@ -30,10 +22,8 @@ function scale!( var::Symbol, scale::Real, ) - for layer in diagn.layers - variable = getfield(layer.grid_variables, var) - variable .*= scale - end + variable = getfield(diagn.grid, var) + variable .*= scale end """ diff --git a/src/dynamics/spectral_grid.jl b/src/dynamics/spectral_grid.jl index af6c15c1c..270d562f0 100644 --- a/src/dynamics/spectral_grid.jl +++ b/src/dynamics/spectral_grid.jl @@ -1,12 +1,15 @@ abstract type AbstractSpectralGrid end -abstract type AbstractGeometry end +# computing const DEFAULT_NF = Float32 -const DEFAULT_MODEL = PrimitiveDry +const DEFAULT_DEVICE = CPU() +const DEFAULT_ARRAYTYPE = Array + +# numerics const DEFAULT_GRID = OctahedralGaussianGrid const DEFAULT_RADIUS = 6.371e6 const DEFAULT_TRUNC = 31 -const DEFAULT_NLEV = 8 +const DEFAULT_NLAYERS = 8 export SpectralGrid @@ -17,10 +20,16 @@ $(TYPEDFIELDS) `nlat_half` and `npoints` should not be chosen but are derived from `trunc`, `Grid` and `dealiasing`.""" -Base.@kwdef struct SpectralGrid <: AbstractSpectralGrid +@kwdef struct SpectralGrid <: AbstractSpectralGrid "[OPTION] number format used throughout the model" NF::Type{<:AbstractFloat} = DEFAULT_NF + "[OPTION] device archictecture to run on" + device::AbstractDevice = DEFAULT_DEVICE + + "[OPTION] array type to use for all variables" + ArrayType::Type{<:AbstractArray} = default_array_type(device) + # HORIZONTAL "[OPTION] horizontal resolution as the maximum degree of spherical harmonics" trunc::Int = DEFAULT_TRUNC @@ -35,7 +44,7 @@ Base.@kwdef struct SpectralGrid <: AbstractSpectralGrid radius::Float64 = DEFAULT_RADIUS "[OPTION] number of particles for particle advection [1]" - n_particles::Int = 0 + nparticles::Int = 0 # SIZE OF GRID from trunc, Grid, dealiasing: "number of latitude rings on one hemisphere (Equator incl)" @@ -49,31 +58,25 @@ Base.@kwdef struct SpectralGrid <: AbstractSpectralGrid # VERTICAL "[OPTION] number of vertical levels" - nlev::Int = DEFAULT_NLEV + nlayers::Int = DEFAULT_NLAYERS "[OPTION] coordinates used to discretize the vertical" - vertical_coordinates::VerticalCoordinates = SigmaCoordinates(; nlev) - - # make sure nlev and vertical_coordinates.nlev match - function SpectralGrid(NF, trunc, Grid, dealiasing, radius, n_particles, nlat_half, nlat, npoints, nlev, vertical_coordinates) - if nlev == vertical_coordinates.nlev - return new(NF, trunc, Grid, dealiasing, radius, n_particles, nlat_half, nlat, npoints, - nlev, vertical_coordinates) - else # use nlev from vert_coords: - return new(NF, trunc, Grid, dealiasing, radius, n_particles, nlat_half, nlat, npoints, - vertical_coordinates.nlev, vertical_coordinates) - end - end + vertical_coordinates::VerticalCoordinates = SigmaCoordinates(; nlayers) + + # ARRAY TYPES (horizontal dimension in grid/spectral is flattened to 1D) + SpectralVariable2D::Type{<:AbstractArray} = LowerTriangularArray{Complex{NF}, 1, ArrayType{Complex{NF}, 1}} + SpectralVariable3D::Type{<:AbstractArray} = LowerTriangularArray{Complex{NF}, 2, ArrayType{Complex{NF}, 2}} + SpectralVariable4D::Type{<:AbstractArray} = LowerTriangularArray{Complex{NF}, 3, ArrayType{Complex{NF}, 3}} + GridVariable2D::Type{<:AbstractArray} = RingGrids.nonparametric_type(Grid){NF, 1, ArrayType{NF, 1}} + GridVariable3D::Type{<:AbstractArray} = RingGrids.nonparametric_type(Grid){NF, 2, ArrayType{NF, 2}} + GridVariable4D::Type{<:AbstractArray} = RingGrids.nonparametric_type(Grid){NF, 3, ArrayType{NF, 3}} + ParticleVector::Type{<:AbstractArray} = ArrayType{Particle{NF}, 1} 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...) - function Base.show(io::IO, SG::SpectralGrid) - (; NF, trunc, Grid, radius, nlat, npoints, nlev, vertical_coordinates) = SG - (; n_particles) = SG + (; NF, trunc, Grid, radius, nlat, npoints, nlayers, vertical_coordinates) = SG + (; device, ArrayType) = SG + (; nparticles) = SG # resolution information res_ave = sqrt(4π*radius^2/npoints)/1000 # in [km] @@ -83,10 +86,11 @@ function Base.show(io::IO, SG::SpectralGrid) println(io, "├ Spectral: T$trunc LowerTriangularMatrix{Complex{$NF}}, radius = $radius m") println(io, "├ Grid: $nlat-ring $Grid{$NF}, $npoints grid points") println(io, "├ Resolution: $(s(res_ave))km (average)") - if n_particles > 0 - println(io, "├ Particles: $n_particles") + if nparticles > 0 + println(io, "├ Particles: $nparticles") end - print(io, "└ Vertical: $nlev-level $(typeof(vertical_coordinates))") + println(io, "├ Vertical: $nlayers-layer $(typeof(vertical_coordinates))") + print(io, "└ Device: $(typeof(device)) using $ArrayType") end """ diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl index 84e31c4a0..1fdb6df5e 100644 --- a/src/dynamics/tendencies.jl +++ b/src/dynamics/tendencies.jl @@ -1,12 +1,13 @@ -""" -$(TYPEDSIGNATURES) +"""$(TYPEDSIGNATURES) Calculate all tendencies for the BarotropicModel.""" -function dynamics_tendencies!( diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - time::DateTime, - model::Barotropic) - forcing!(diagn, progn, model.forcing, time, model) # = (Fᵤ, Fᵥ) forcing for u, v - drag!(diagn, progn, model.drag, time, model) # drag term for u, v +function dynamics_tendencies!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, # leapfrog index to evaluate tendencies at + model::Barotropic, +) + forcing!(diagn, progn, model.forcing, model, lf) # = (Fᵤ, Fᵥ) forcing for u, v + drag!(diagn, progn, model.drag, model, lf) # drag term for u, v vorticity_flux!(diagn, model) # = ∇×(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ) end @@ -14,233 +15,188 @@ end $(TYPEDSIGNATURES) Calculate all tendencies for the ShallowWaterModel.""" function dynamics_tendencies!( - diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - surface::SurfaceVariables, - pres::LowerTriangularMatrix, # spectral pressure/η for geopotential - time::DateTime, # time to evaluate the tendencies at + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, # leapfrog index to evaluate tendencies at model::ShallowWater, ) (; forcing, drag, planet, atmosphere, orography) = model (; spectral_transform, geometry) = model - # for compatibility with other ModelSetups pressure pres = interface displacement η here - forcing!(diagn, progn, forcing, time, model) # = (Fᵤ, Fᵥ, Fₙ) forcing for u, v, η - drag!(diagn, progn, drag, time, model) # drag term for momentum u, v - + # for compatibility with other AbstractModels pressure pres = interface displacement η here + forcing!(diagn, progn, forcing, model, lf) # = (Fᵤ, Fᵥ, Fₙ) forcing for u, v, η + drag!(diagn, progn, drag, model, lf) # drag term for u, v + # = ∇×(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ), tendency for vorticity # = ∇⋅(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ), tendency for divergence vorticity_flux!(diagn, model) - - geopotential!(diagn, pres, planet) # geopotential Φ = gη in shallow water + + geopotential!(diagn, progn.pres[lf], planet) # geopotential Φ = gη in shallow water bernoulli_potential!(diagn, spectral_transform) # = -∇²(E+Φ), tendency for divergence # = -∇⋅(uh, vh), tendency for "pressure" η - volume_flux_divergence!(diagn, surface, orography, atmosphere, geometry, spectral_transform) + volume_flux_divergence!(diagn, orography, atmosphere, geometry, spectral_transform) end -""" -$(TYPEDSIGNATURES) +"""$(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 +function dynamics_tendencies!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, # leapfrog index for tendencies + model::PrimitiveEquation, +) - O = model.orography - G = model.geometry - S = model.spectral_transform - GP = model.geopotential - A = model.atmosphere - I = model.implicit - (; surface ) = diagn + (; orography, geometry, spectral_transform, geopotential, atmosphere, implicit) = model # 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.implicit.α == 0 ? lf : 1 + lf_implicit = implicit.α == 0 ? lf : 1 - pressure_gradient!(diagn, progn, lf, S) # calculate ∇ln(pₛ) + # calculate ∇ln(pₛ), then (u_k, v_k)⋅∇ln(p_s) + pressure_gradient_flux!(diagn, progn, lf, spectral_transform) - @floop for (diagn_layer, progn_layer) in zip(diagn.layers, progn.layers) - pressure_flux!(diagn_layer, surface) # calculate (uₖ, vₖ)⋅∇ln(pₛ) + # calculate Tᵥ = T + Tₖμq in spectral as a approxmation to Tᵥ = T(1+μq) used for geopotential + linear_virtual_temperature!(diagn, progn, lf_implicit, model) - # 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, I) # temperature relative to profile - end + # temperature relative to profile + temperature_anomaly!(diagn, implicit) - geopotential!(diagn, GP, O) # from ∂Φ/∂ln(pₛ) = -RTᵥ for bernoulli_potential! - vertical_integration!(diagn, progn, lf_implicit, G)# get ū, v̄, D̄ on grid; D̄ in spectral - surface_pressure_tendency!(surface, S) # ∂ln(pₛ)/∂t = -(ū, v̄)⋅∇ln(pₛ) - D̄ - - @floop for layer in diagn.layers - vertical_velocity!(layer, surface, G) # calculate σ̇ for the vertical mass flux M = pₛσ̇ - # add the RTₖlnpₛ term to geopotential - linear_pressure_gradient!(layer, progn.surface, lf_implicit, A, I) - 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, 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 + # from ∂Φ/∂ln(pₛ) = -RTᵥ for bernoulli_potential! + geopotential!(diagn, geopotential, orography) -""" -$(TYPEDSIGNATURES) -Set the tendencies in `diagn` to zero.""" -function zero_tendencies!(diagn::DiagnosticVariables{NF, Grid, Model}) where {NF, Grid, Model<:Barotropic} - for layer in diagn.layers - fill!(layer.tendencies.u_tend_grid, 0) - fill!(layer.tendencies.v_tend_grid, 0) - fill!(layer.tendencies.vor_tend, 0) - end -end + # get ū, v̄, D̄ on grid; D̄ in spectral + vertical_integration!(diagn, progn, lf_implicit, geometry) -""" -$(TYPEDSIGNATURES) -Set the tendencies in `diagn` to zero.""" -function zero_tendencies!(diagn::DiagnosticVariables{NF, Grid, Model}) where {NF, Grid, Model<:ShallowWater} - for layer in diagn.layers - fill!(layer.tendencies.u_tend_grid, 0) - fill!(layer.tendencies.v_tend_grid, 0) - fill!(layer.tendencies.vor_tend, 0) - fill!(layer.tendencies.div_tend, 0) - end - fill!(diagn.surface.pres_tend_grid, 0) - fill!(diagn.surface.pres_tend, 0) -end + # ∂ln(pₛ)/∂t = -(ū, v̄)⋅∇ln(pₛ) - D̄ + surface_pressure_tendency!(diagn, spectral_transform) -""" -$(TYPEDSIGNATURES) -Set the tendencies in `diagn` to zero.""" -function zero_tendencies!(diagn::DiagnosticVariables{NF, Grid, Model}) where {NF, Grid, Model<:PrimitiveDry} - for layer in diagn.layers - fill!(layer.tendencies.u_tend_grid, 0) - fill!(layer.tendencies.v_tend_grid, 0) - fill!(layer.tendencies.vor_tend, 0) - fill!(layer.tendencies.div_tend, 0) - fill!(layer.tendencies.temp_tend_grid, 0) - end - fill!(diagn.surface.pres_tend_grid, 0) - fill!(diagn.surface.pres_tend, 0) -end + # calculate vertical velocity σ̇ in sigma coordinates for the vertical mass flux M = p_s*σ̇ + vertical_velocity!(diagn, geometry) + + # add the RTₖlnpₛ term to geopotential + linear_pressure_gradient!(diagn, progn, lf_implicit, atmosphere, implicit) -""" -$(TYPEDSIGNATURES) -Set the tendencies in `diagn` to zero.""" -function zero_tendencies!(diagn::DiagnosticVariables{NF, Grid, Model}) where {NF, Grid, Model<:PrimitiveWet} - for layer in diagn.layers - fill!(layer.tendencies.u_tend_grid, 0) - fill!(layer.tendencies.v_tend_grid, 0) - fill!(layer.tendencies.vor_tend, 0) - fill!(layer.tendencies.div_tend, 0) - fill!(layer.tendencies.temp_tend_grid, 0) - fill!(layer.tendencies.humid_tend_grid, 0) - end - fill!(diagn.surface.pres_tend_grid, 0) - fill!(diagn.surface.pres_tend, 0) -end + # use σ̇ for the vertical advection of u, v, T, q + vertical_advection!(diagn, model) -function pressure_gradient!(diagn::DiagnosticVariables, - progn::PrognosticVariables, - lf::Integer, # leapfrog index - S::SpectralTransform) - - (; 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_spec, ∇lnp_y_spec, pres, S) # CALCULATE ∇ln(pₛ) - gridded!(∇lnp_x, ∇lnp_x_spec, S, unscale_coslat=true) # transform to grid: zonal gradient - gridded!(∇lnp_y, ∇lnp_y_spec, S, unscale_coslat=true) # meridional gradient -end + # vorticity advection, pressure gradient term + vordiv_tendencies!(diagn, model) -function pressure_flux!(diagn::DiagnosticVariablesLayer, - surf::SurfaceVariables) + # hor. advection + adiabatic term + temperature_tendency!(diagn, model) - (; ∇lnp_x, ∇lnp_y ) = surf # zonal, meridional gradient of log surface pressure - (; u_grid, v_grid ) = diagn.grid_variables - (; uv∇lnp ) = diagn.dynamics_variables - - @. uv∇lnp = u_grid*∇lnp_x + v_grid*∇lnp_y # the (u, v)⋅∇lnpₛ term + # horizontal advection of humidity (nothing for wetcore) + humidity_tendency!(diagn, model) + + # add -∇²(E+ϕ+RTₖlnpₛ) term to div tendency + bernoulli_potential!(diagn, spectral_transform) + + return nothing end -"""Convert absolute and virtual temperature to anomalies wrt to the reference profile""" -function temperature_anomaly!( - diagn::DiagnosticVariablesLayer, - I::ImplicitPrimitiveEquation, - ) - - Tₖ = I.temp_profile[diagn.k] # reference temperature on this layer - (; temp_grid, temp_virt_grid ) = diagn.grid_variables - - @. temp_grid -= Tₖ # absolute temperature -> anomaly - @. temp_virt_grid -= Tₖ # virtual temperature -> anomaly +"""$(TYPEDSIGNATURES) +Compute the gradient ∇lnp_s of the logarithm of surface pressure, followed by +its flux, (u,v) * ∇lnp_s.""" +function pressure_gradient_flux!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, # leapfrog index + S::SpectralTransform, +) + # PRESSURE GRADIENT + pres = progn.pres[lf] # log of surface pressure at leapfrog step lf + ∇lnp_x_spec = diagn.dynamics.a_2D # reuse 2D work arrays for gradients + ∇lnp_y_spec = diagn.dynamics.b_2D # in spectral space + (; ∇lnp_x, ∇lnp_y) = diagn.dynamics # but store in grid space + + ∇!(∇lnp_x_spec, ∇lnp_y_spec, pres, S) # CALCULATE ∇ln(pₛ) + transform!(∇lnp_x, ∇lnp_x_spec, S, unscale_coslat=true) # transform to grid: zonal gradient + transform!(∇lnp_y, ∇lnp_y_spec, S, unscale_coslat=true) # meridional gradient + + (; u_grid, v_grid ) = diagn.grid + (; uv∇lnp ) = diagn.dynamics + + # PRESSURE GRADIENT FLUX + @inbounds for k in eachgrid(u_grid, v_grid, uv∇lnp) + for ij in eachgridpoint(u_grid, v_grid, uv∇lnp) + # the (u, v)⋅∇lnp_s term + uv∇lnp[ij, k] = u_grid[ij, k]*∇lnp_x[ij] + v_grid[ij, k]*∇lnp_y[ij] + end + end end -""" - vertical_integration!(Diag::DiagnosticVariables, G::Geometry) +"""$(TYPEDSIGNATURES) +Convert absolute and virtual temperature to anomalies wrt to the reference profile""" +function temperature_anomaly!( + diagn::DiagnosticVariables, + implicit::ImplicitPrimitiveEquation, +) + (; temp_profile) = implicit # reference temperature profile + (; temp_grid, temp_virt_grid ) = diagn.grid + + @inbounds for k in eachgrid(temp_grid, temp_virt_grid) + Tₖ = temp_profile[k] + for ij in eachgridpoint(temp_grid, temp_virt_grid) + temp_grid[ij, k] -= Tₖ # absolute temperature -> anomaly + temp_virt_grid[ij, k] -= Tₖ # virtual temperature -> anomaly + end + end +end +"""$(TYPEDSIGNATURES) Calculates the vertically averaged (weighted by the thickness of the σ level) velocities (*coslat) and divergence. E.g. - u_mean = ∑_k=1^nlev Δσ_k * u_k + u_mean = ∑_k=1^nlayers Δσ_k * u_k u, v are averaged in grid-point space, divergence in spectral space. """ -function vertical_integration!( diagn::DiagnosticVariables{NF}, - progn::PrognosticVariables{NF}, - lf::Int, # leapfrog index for D̄_spec - G::Geometry{NF}) where NF - - (; σ_levels_thick, nlev ) = G - (; ∇lnp_x, ∇lnp_y ) = diagn.surface # zonal, meridional grad of log surface pressure - - ū = diagn.surface.u_mean_grid # rename for convenience - v̄ = diagn.surface.v_mean_grid - D̄ = diagn.surface.div_mean_grid - D̄_spec = diagn.surface.div_mean +function vertical_integration!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, # leapfrog index for D̄_spec + geometry::Geometry, +) + (; σ_levels_thick, nlayers ) = geometry + (; ∇lnp_x, ∇lnp_y ) = diagn.dynamics # zonal, meridional grad of log surface pressure + (; u_grid, v_grid, div_grid) = diagn.grid + (; u_mean_grid, v_mean_grid, div_mean_grid, div_mean) = diagn.dynamics + (; div_sum_above, uv∇lnp_sum_above) = diagn.dynamics + div = progn.div[lf] - @boundscheck nlev == diagn.nlev || throw(BoundsError) + @boundscheck nlayers == diagn.nlayers || throw(BoundsError) - fill!(ū, 0) # reset accumulators from previous vertical average - fill!(v̄, 0) - fill!(D̄, 0) - fill!(D̄_spec, 0) + fill!(u_mean_grid, 0) # reset accumulators from previous vertical average + fill!(v_mean_grid, 0) + fill!(div_mean_grid, 0) + fill!(div_mean, 0) - @inbounds for k in 1:nlev + @inbounds for k in 1:nlayers # integrate from top to bottom # arrays for layer-thickness weighted column averages Δσₖ = σ_levels_thick[k] - u = diagn.layers[k].grid_variables.u_grid - v = diagn.layers[k].grid_variables.v_grid - D = diagn.layers[k].grid_variables.div_grid - D_spec = progn.layers[k].timesteps[lf].div - - # for the Σ_r=1^k-1 Δσᵣ(Dᵣ + u̲⋅∇lnpₛ) vertical integration - # Simmons and Burridge, 1981 eq 3.12 split into div and u̲⋅∇lnpₛ - D̄ᵣ = diagn.layers[k].dynamics_variables.div_sum_above - ūv̄∇lnpᵣ = diagn.layers[k].dynamics_variables.uv∇lnp_sum_above # GRID-POINT SPACE: u, v, D with thickness weighting Δσₖ # before this k's u, v, D are added to ū, v̄, D̄ store in the # sum_above fields for a 1:k-1 integration # which is =0 for k=1 as ū, v̄, D̄ accumulators are 0-initialised - @. D̄ᵣ = D̄ - @. ūv̄∇lnpᵣ = ū*∇lnp_x + v̄*∇lnp_y - - @. ū += u*Δσₖ # now add the k-th element to the sum - @. v̄ += v*Δσₖ - @. D̄ += D*Δσₖ + for ij in eachgridpoint(u_mean_grid, v_mean_grid, div_mean_grid) + # for the Σ_r=1^k-1 Δσᵣ(Dᵣ + u̲⋅∇lnpₛ) vertical integration + # Simmons and Burridge, 1981 eq 3.12 split into div and u̲⋅∇lnpₛ + div_sum_above[ij, k] = div_mean_grid[ij] + uv∇lnp_sum_above[ij, k] = u_mean_grid[ij]*∇lnp_x[ij] + v_mean_grid[ij]*∇lnp_y[ij] + + u_mean_grid[ij] += u_grid[ij, k]*Δσₖ # now add the k-th element to the sum + v_mean_grid[ij] += v_grid[ij, k]*Δσₖ + div_mean_grid[ij] += div_grid[ij, k]*Δσₖ + end # SPECTRAL SPACE: divergence - @. D̄_spec += D_spec*Δσₖ + for lm in eachharmonic(div, div_mean) + div_mean[lm] += div[lm, k]*Δσₖ + end end end @@ -261,64 +217,67 @@ of the logarithm of surface pressure ln(p_s) and D̄ the vertically averaged div 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, + diagn::DiagnosticVariables, S::SpectralTransform, ) - (; pres_tend, pres_tend_grid, ∇lnp_x, ∇lnp_y ) = surf + (; pres_tend, pres_tend_grid) = diagn.tendencies + (; ∇lnp_x, ∇lnp_y, u_mean_grid, v_mean_grid, div_mean) = diagn.dynamics - # vertical averages need to be computed first! - ū = surf.u_mean_grid # rename for convenience - v̄ = surf.v_mean_grid - D̄ = surf.div_mean # spectral - # 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, S) + @. pres_tend_grid = u_mean_grid*∇lnp_x + v_mean_grid*∇lnp_y + transform!(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 + # for semi-implicit div_mean is calc at time step i-1 in vertical_integration! + @. pres_tend = -pres_tend - div_mean # add the -div_mean term in spectral, swap sign pres_tend[1] = 0 # for mass conservation return nothing end function vertical_velocity!( - diagn::DiagnosticVariablesLayer, - surf::SurfaceVariables, - G::Geometry, + diagn::DiagnosticVariables, + geometry::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 + (; σ_levels_thick, σ_levels_half, nlayers) = geometry + (; σ_tend) = diagn.dynamics - # 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 + # sum of Δσ-weighted div, uv∇lnp from 1:k-1 + (; div_sum_above, uv∇lnp, uv∇lnp_sum_above) = diagn.dynamics + (; div_mean_grid) = diagn.dynamics # vertical avrgd div to be added to ūv̄∇lnp + (; σ_tend) = diagn.dynamics # vertical mass flux M = pₛσ̇ at k+1/2 + (; div_grid) = diagn.grid + ūv̄∇lnp = diagn.tendencies.pres_tend_grid # calc'd in surface_pressure_tendency! (excl -D̄) - ū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 + grids_match(σ_tend, div_sum_above, div_grid, uv∇lnp_sum_above, uv∇lnp) || + throw(DimensionMismatch(σ_tend, div_sum_above, div_grid, uv∇lnp_sum_above, uv∇lnp)) + + @inbounds for k in 1:nlayers-1 + Δσₖ = σ_levels_thick[k] + σₖ_half = σ_levels_half[k+1] - # 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 == G.nlev && (fill!(σ̇, 0); return nothing) + for ij in eachgridpoint(σ_tend) + # Hoskins and Simmons, 1975 just before eq. (6) + σ_tend[ij, k] = σₖ_half*(div_mean_grid[ij] + ūv̄∇lnp[ij]) - + (div_sum_above[ij, k] + Δσₖ*div_grid[ij, k]) - + (uv∇lnp_sum_above[ij, k] + Δσₖ*uv∇lnp[ij, k]) + end + end - # Hoskins and Simmons, 1975 just before eq. (6) - σ̇ .= σk_half*(D̄ .+ ūv̄∇lnp) .- - (div_sum_above .+ Δσₖ*div_grid) .- # for 1:k integral add level k σₖ-weighted div - (uv∇lnp_sum_above .+ Δσₖ*uv∇lnp) # and level k σₖ-weighted uv∇lnp here + # mass flux σ̇ is zero at k=1/2 (not explicitly stored) and k=nlayers+1/2 (stored in layer k) + # set to zero for bottom layer then + σ_tend[:, nlayers] .= 0 + return nothing end """ $(TYPEDSIGNATURES) Function barrier to unpack `model`.""" function vordiv_tendencies!( - diagn::DiagnosticVariablesLayer, - surf::SurfaceVariables, + diagn::DiagnosticVariables, model::PrimitiveEquation, ) (; coriolis, atmosphere, geometry, spectral_transform) = model - vordiv_tendencies!(diagn, surf, coriolis, atmosphere, geometry, spectral_transform) + vordiv_tendencies!(diagn, coriolis, atmosphere, geometry, spectral_transform) end """$(TYPEDSIGNATURES) @@ -340,8 +299,7 @@ spectral space `+ ...` because there's more terms added later for divergence.""" function vordiv_tendencies!( - diagn::DiagnosticVariablesLayer, - surf::SurfaceVariables, + diagn::DiagnosticVariables, coriolis::AbstractCoriolis, atmosphere::AbstractAtmosphere, geometry::AbstractGeometry, @@ -351,84 +309,65 @@ function vordiv_tendencies!( (; f) = coriolis # coriolis parameter (; coslat⁻¹) = geometry - (; u_tend_grid, v_tend_grid) = diagn.tendencies # already contains vertical advection - u = diagn.grid_variables.u_grid # velocity - v = diagn.grid_variables.v_grid # velocity - vor = diagn.grid_variables.vor_grid # relative vorticity - ∇lnp_x = surf.∇lnp_x # zonal gradient of logarithm of surface pressure - ∇lnp_y = surf.∇lnp_y # meridional gradient thereof - Tᵥ = diagn.grid_variables.temp_virt_grid # virtual temperature (anomaly!) + # tendencies already contain parameterizations + advection, therefore accumulate + (; u_tend_grid, v_tend_grid) = diagn.tendencies + (; u_grid, v_grid, vor_grid, temp_virt_grid) = diagn.grid # velocities, vorticity + (; ∇lnp_x, ∇lnp_y) = diagn.dynamics # zonal/meridional gradient of logarithm of surface pressure # precompute ring indices and boundscheck - rings = eachring(u_tend_grid, v_tend_grid, u, v, vor, ∇lnp_x, ∇lnp_y, Tᵥ) - - @inbounds for (j, ring) in enumerate(rings) - coslat⁻¹j = coslat⁻¹[j] - f_j = f[j] - for ij in ring - ω = vor[ij] + f_j # absolute vorticity - RTᵥ = R_dry*Tᵥ[ij] # dry gas constant * virtual temperature anomaly - u_tend_grid[ij] = (u_tend_grid[ij] + v[ij]*ω - RTᵥ*∇lnp_x[ij])*coslat⁻¹j - v_tend_grid[ij] = (v_tend_grid[ij] - u[ij]*ω - RTᵥ*∇lnp_y[ij])*coslat⁻¹j + rings = eachring(u_tend_grid, v_tend_grid, u_grid, v_grid, vor_grid, temp_virt_grid) + + @inbounds for k in eachgrid(u_tend_grid, v_tend_grid) + for (j, ring) in enumerate(rings) + coslat⁻¹j = coslat⁻¹[j] + f_j = f[j] + for ij in ring + ω = vor_grid[ij, k] + f_j # absolute vorticity + RTᵥ = R_dry*temp_virt_grid[ij, k] # dry gas constant * virtual temperature anomaly(!) + u_tend_grid[ij, k] = (u_tend_grid[ij, k] + v_grid[ij, k]*ω - RTᵥ*∇lnp_x[ij])*coslat⁻¹j + v_tend_grid[ij, k] = (v_tend_grid[ij, k] - u_grid[ij, k]*ω - RTᵥ*∇lnp_y[ij])*coslat⁻¹j + end 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 + u_tend = diagn.dynamics.a + v_tend = diagn.dynamics.b - spectral!(u_tend, u_tend_grid, S) - spectral!(v_tend, v_tend_grid, S) + transform!(u_tend, u_tend_grid, S) + transform!(v_tend, v_tend_grid, S) curl!(vor_tend, u_tend, v_tend, S) # ∂ζ/∂t = ∇×(u_tend, v_tend) divergence!(div_tend, u_tend, v_tend, S) # ∂D/∂t = ∇⋅(u_tend, v_tend) return nothing end -# function barrier -function tendencies_physics_only!( - diagn::DiagnosticVariablesLayer, - model::PrimitiveEquation -) - wet_core = model isa PrimitiveWet - tendencies_physics_only!(diagn, model.geometry, model.spectral_transform, wet_core) -end - -"""For dynamics=false, after calling parameterization_tendencies! call this function +"""$(TYPEDSIGNATURES) +For dynamics=false, after calling parameterization_tendencies! call this function to transform the physics tendencies from grid-point to spectral space including the necessary coslat⁻¹ scaling.""" -function tendencies_physics_only!( - diagn::DiagnosticVariablesLayer, - G::AbstractGeometry, - S::SpectralTransform, - wet_core::Bool = true +function physics_tendencies_only!( + diagn::DiagnosticVariables, + model::PrimitiveEquation, ) - (; coslat⁻¹) = G + (; coslat⁻¹) = model.geometry + S = model.spectral_transform # already contain parameterizations (; u_tend_grid, v_tend_grid, temp_tend_grid, humid_tend_grid) = diagn.tendencies - - # precompute ring indices and boundscheck - rings = eachring(u_tend_grid, v_tend_grid) - - @inbounds for (j, ring) in enumerate(rings) - coslat⁻¹j = coslat⁻¹[j] - for ij in ring - u_tend_grid[ij] *= coslat⁻¹j - v_tend_grid[ij] *= coslat⁻¹j - end - end + RingGrids._scale_lat!(u_tend_grid, coslat⁻¹) + RingGrids._scale_lat!(v_tend_grid, coslat⁻¹) # divergence and curl of that u, v_tend vector for vor, div tendencies (; vor_tend, div_tend, temp_tend, humid_tend) = diagn.tendencies - u_tend = diagn.dynamics_variables.a - v_tend = diagn.dynamics_variables.b + u_tend = diagn.dynamics.a + v_tend = diagn.dynamics.b - spectral!(u_tend, u_tend_grid, S) - spectral!(v_tend, v_tend_grid, S) - spectral!(temp_tend, temp_tend_grid, S) - wet_core && spectral!(humid_tend, humid_tend_grid, S) + transform!(u_tend, u_tend_grid, S) + transform!(v_tend, v_tend_grid, S) + transform!(temp_tend, temp_tend_grid, S) + model isa PrimitiveWet && transform!(humid_tend, humid_tend_grid, S) curl!(vor_tend, u_tend, v_tend, S) # ∂ζ/∂t = ∇×(u_tend, v_tend) divergence!(div_tend, u_tend, v_tend, S) # ∂D/∂t = ∇⋅(u_tend, v_tend) @@ -437,12 +376,12 @@ end # function barrier function temperature_tendency!( - diagn::DiagnosticVariablesLayer, + diagn::DiagnosticVariables, model::PrimitiveEquation, ) - (; adiabatic_conversion, atmosphere, geometry, spectral_transform, implicit) = model - temperature_tendency!(diagn, adiabatic_conversion, atmosphere, geometry, - spectral_transform, implicit) + (; adiabatic_conversion, atmosphere, implicit, geometry, spectral_transform) = model + temperature_tendency!(diagn, adiabatic_conversion, atmosphere, implicit, + geometry, spectral_transform) end """ @@ -455,38 +394,44 @@ Compute the temperature tendency `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, + diagn::DiagnosticVariables, adiabatic_conversion::AbstractAdiabaticConversion, atmosphere::AbstractAtmosphere, + implicit::ImplicitPrimitiveEquation, G::Geometry, S::SpectralTransform, - I::ImplicitPrimitiveEquation, ) (; 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 - - (; κ) = atmosphere # thermodynamic kappa = R_Dry/heat_capacity - Tᵥ = diagn.grid_variables.temp_virt_grid # anomaly wrt to Tₖ - Tₖ = I.temp_profile[diagn.k] # average layer temperature from reference profile - - # coefficients from Simmons and Burridge 1981 - σ_lnp_A = adiabatic_conversion.σ_lnp_A[diagn.k] # eq. 3.12, -1/Δσₖ*ln(σ_k+1/2/σ_k-1/2) - σ_lnp_B = adiabatic_conversion.σ_lnp_B[diagn.k] # eq. 3.12 -αₖ - + (; div_grid, temp_grid) = diagn.grid + (; uv∇lnp, uv∇lnp_sum_above, div_sum_above) = diagn.dynamics + (; κ) = atmosphere # thermodynamic kappa = R_Dry/heat_capacity + Tᵥ = diagn.grid.temp_virt_grid # anomaly wrt to Tₖ + (; temp_profile) = implicit + # semi-implicit: terms here are explicit+implicit evaluated at time step i # implicit_correction! then calculated the implicit terms from Vi-1 minus Vi # to move the implicit terms to i-1 which is cheaper then the alternative below - # Adiabatic conversion term following Simmons and Burridge 1981 but for σ coordinates - # += as tend already contains parameterizations + vertical advection - @. temp_tend_grid += temp_grid*div_grid + # +T'D term of hori advection - κ*(Tᵥ+Tₖ)*( # +κTᵥ*Dlnp/Dt, adiabatic term - σ_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 + @inbounds for k in eachgrid(temp_tend_grid, temp_grid, div_grid, Tᵥ, div_sum_above, uv∇lnp_sum_above) + Tₖ = temp_profile[k] # average layer temperature from reference profile + + # coefficients from Simmons and Burridge 1981 + σ_lnp_A = adiabatic_conversion.σ_lnp_A[k] # eq. 3.12, -1/Δσₖ*ln(σ_k+1/2/σ_k-1/2) + σ_lnp_B = adiabatic_conversion.σ_lnp_B[k] # eq. 3.12 -αₖ + + # Adiabatic conversion term following Simmons and Burridge 1981 but for σ coordinates + # += as tend already contains parameterizations + vertical advection + for ij in eachgridpoint(temp_tend_grid, temp_grid, div_grid, Tᵥ) + temp_tend_grid[ij, k] += + temp_grid[ij, k] * div_grid[ij, k] + # +T'D term of hori advection + κ * (Tᵥ[ij, k] + Tₖ)*( # +κTᵥ*Dlnp/Dt, adiabatic term + σ_lnp_A * (div_sum_above[ij, k] + uv∇lnp_sum_above[ij, k]) + # eq. 3.12 1st term + σ_lnp_B * (div_grid[ij, k] + uv∇lnp[ij, k]) + # eq. 3.12 2nd term + uv∇lnp[ij, k]) # eq. 3.13 + end + end - spectral!(temp_tend, temp_tend_grid, S) + transform!(temp_tend, temp_tend_grid, S) # now add the -∇⋅((u, v)*T') term flux_divergence!(temp_tend, temp_grid, diagn, G, S, add=true, flipsign=true) @@ -494,42 +439,44 @@ function temperature_tendency!( return nothing end -function humidity_tendency!(diagn::DiagnosticVariablesLayer, +function humidity_tendency!(diagn::DiagnosticVariables, model::PrimitiveWet) G = model.geometry S = model.spectral_transform (; humid_tend, humid_tend_grid ) = diagn.tendencies - (; humid_grid ) = diagn.grid_variables + (; humid_grid ) = diagn.grid # add horizontal advection to parameterization + vertical advection tendencies horizontal_advection!(humid_tend, humid_tend_grid, humid_grid, diagn, G, S, add=true) end # no humidity tendency for dry core -humidity_tendency!(::DiagnosticVariablesLayer, ::PrimitiveDry) = nothing +humidity_tendency!(::DiagnosticVariables, ::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}, + A_tend::LowerTriangularArray, # Ouput: tendency to write into + A_tend_grid::AbstractGridArray, # Input: tendency incl prev terms + A_grid::AbstractGridArray, # Input: grid field to be advected + diagn::DiagnosticVariables, G::Geometry, S::SpectralTransform; - add::Bool=true # add/overwrite A_tend_grid? -) where NF + add::Bool=true, # add/overwrite A_tend_grid? +) - (; div_grid) = diagn.grid_variables + (; div_grid) = diagn.grid @inline kernel(a, b, c) = add ? a+b*c : b*c - # +A*div term of the advection operator - @inbounds for ij in eachgridpoint(A_tend_grid, A_grid, div_grid) - # add as tend already contains parameterizations + vertical advection - A_tend_grid[ij] = kernel(A_tend_grid[ij], A_grid[ij], div_grid[ij]) + for k in eachgrid(A_tend_grid, A_grid, div_grid) + # +A*div term of the advection operator + @inbounds for ij in eachgridpoint(A_tend_grid, A_grid, div_grid) + # add as tend already contains parameterizations + vertical advection + A_tend_grid[ij, k] = kernel(A_tend_grid[ij, k], A_grid[ij, k], div_grid[ij, k]) + end end - spectral!(A_tend, A_tend_grid, S) # for +A*div in spectral space + transform!(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, G, S, add=true, flipsign=true) @@ -545,36 +492,39 @@ Computes ∇⋅((u, v)*A) with the option to add/overwrite `A_tend` and to - `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}, - 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 +function flux_divergence!( + A_tend::LowerTriangularArray, # Ouput: tendency to write into + A_grid::AbstractGridArray, # Input: grid field to be advected + diagn::DiagnosticVariables, # for u_grid, v_grid + G::Geometry, + S::SpectralTransform; + add::Bool=true, # add result to A_tend or overwrite for false + flipsign::Bool=true, # compute -∇⋅((u, v)*A) (true) or ∇⋅((u, v)*A)? +) + (; u_grid, v_grid) = diagn.grid (; coslat⁻¹) = G # reuse general work arrays a, b, a_grid, b_grid - uA = diagn.dynamics_variables.a # = u*A in spectral - vA = diagn.dynamics_variables.b # = v*A in spectral - uA_grid = diagn.dynamics_variables.a_grid # = u*A on grid - vA_grid = diagn.dynamics_variables.b_grid # = v*A on grid - - rings = eachring(uA_grid, vA_grid, u_grid, v_grid, A_grid) # precompute ring indices - - @inbounds for (j, ring) in enumerate(rings) - coslat⁻¹j = coslat⁻¹[j] - for ij in ring - Acoslat⁻¹j = A_grid[ij]*coslat⁻¹j - uA_grid[ij] = u_grid[ij]*Acoslat⁻¹j - vA_grid[ij] = v_grid[ij]*Acoslat⁻¹j + uA = diagn.dynamics.a # = u*A in spectral + vA = diagn.dynamics.b # = v*A in spectral + uA_grid = diagn.dynamics.a_grid # = u*A on grid + vA_grid = diagn.dynamics.b_grid # = v*A on grid + + # precomputed ring indices and check grids_match + rings = eachring(A_grid, u_grid, v_grid) + @inbounds for k in eachgrid(u_grid, v_grid) + for (j, ring) in enumerate(rings) + coslat⁻¹j = coslat⁻¹[j] + for ij in ring + Acoslat⁻¹j = A_grid[ij, k]*coslat⁻¹j + uA_grid[ij, k] = u_grid[ij, k]*Acoslat⁻¹j + vA_grid[ij, k] = v_grid[ij, k]*Acoslat⁻¹j + end end end - spectral!(uA, uA_grid, S) - spectral!(vA, vA_grid, S) + transform!(uA, uA_grid, S) + transform!(vA, vA_grid, S) divergence!(A_tend, uA, vA, S; add, flipsign) return nothing @@ -595,7 +545,7 @@ with 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, +function vorticity_flux_curldiv!( diagn::DiagnosticVariables, coriolis::AbstractCoriolis, geometry::Geometry, S::SpectralTransform; @@ -605,31 +555,32 @@ function vorticity_flux_curldiv!( diagn::DiagnosticVariablesLayer, (; f) = coriolis (; coslat⁻¹) = 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 + (; u_tend_grid, v_tend_grid) = diagn.tendencies # already contains forcing + u = diagn.grid.u_grid # velocity + v = diagn.grid.v_grid # velocity + vor = diagn.grid.vor_grid # relative vorticity - # precompute ring indices and boundscheck + # precompute ring indices and check grids match rings = eachring(u_tend_grid, v_tend_grid, u, v, vor) - - @inbounds for (j, ring) in enumerate(rings) - coslat⁻¹j = coslat⁻¹[j] - f_j = f[j] - for ij in ring - ω = vor[ij] + f_j # 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 + @inbounds for k in eachgrid(u) + for (j, ring) in enumerate(rings) + coslat⁻¹j = coslat⁻¹[j] + f_j = f[j] + for ij in ring + ω = vor[ij, k] + f_j # absolute vorticity + u_tend_grid[ij, k] = (u_tend_grid[ij, k] + v[ij, k]*ω)*coslat⁻¹j + v_tend_grid[ij, k] = (v_tend_grid[ij, k] - u[ij, k]*ω)*coslat⁻¹j + end 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 + u_tend = diagn.dynamics.a + v_tend = diagn.dynamics.b - spectral!(u_tend, u_tend_grid, S) - spectral!(v_tend, v_tend_grid, S) + transform!(u_tend, u_tend_grid, S) + transform!(v_tend, v_tend_grid, S) curl!(vor_tend, u_tend, v_tend, S; add) # ∂ζ/∂t = ∇×(u_tend, v_tend) div && divergence!(div_tend, u_tend, v_tend, S; add) # ∂D/∂t = ∇⋅(u_tend, v_tend) @@ -650,7 +601,7 @@ with 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) +function vorticity_flux!(diagn::DiagnosticVariables, model::ShallowWater) C = model.coriolis G = model.geometry S = model.spectral_transform @@ -670,7 +621,7 @@ with 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) +function vorticity_flux!(diagn::DiagnosticVariables, model::Barotropic) C = model.coriolis G = model.geometry S = model.spectral_transform @@ -687,21 +638,22 @@ Computes the Laplace operator ∇² of the Bernoulli potential `B` in spectral s This version is used for both ShallowWater and PrimitiveEquation, only the geopotential calculation in geopotential! differs.""" -function bernoulli_potential!( diagn::DiagnosticVariablesLayer{NF}, - S::SpectralTransform, - ) where NF - - (; u_grid, v_grid ) = diagn.grid_variables - (; geopot ) = diagn.dynamics_variables - bernoulli = diagn.dynamics_variables.a # reuse work arrays for Bernoulli potential - bernoulli_grid = diagn.dynamics_variables.a_grid +function bernoulli_potential!( + diagn::DiagnosticVariables, + S::SpectralTransform, +) + (; u_grid, v_grid ) = diagn.grid + (; geopot ) = diagn.dynamics + bernoulli = diagn.dynamics.a # reuse work arrays a, a_grid + bernoulli_grid = diagn.dynamics.a_grid (; div_tend ) = diagn.tendencies - half = convert(NF, 0.5) + half = convert(eltype(bernoulli_grid), 0.5) @. bernoulli_grid = half*(u_grid^2 + v_grid^2) # = ½(u² + v²) on grid - spectral!(bernoulli, bernoulli_grid, S) # to spectral space + transform!(bernoulli, bernoulli_grid, S) # to spectral space bernoulli .+= geopot # add geopotential Φ ∇²!(div_tend, bernoulli, S, add=true, flipsign=true) # add -∇²(½(u² + v²) + ϕ) + return nothing end """ @@ -715,35 +667,42 @@ So that the second term inside the Laplace operator can be added to the geopoten 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, + diagn::DiagnosticVariables, + progn::PrognosticVariables, lf::Int, # leapfrog index to evaluate tendencies on atmosphere::AbstractAtmosphere, - I::ImplicitPrimitiveEquation, + implicit::ImplicitPrimitiveEquation, ) - (; R_dry) = atmosphere # dry gas constant - Tₖ = I.temp_profile[diagn.k] # reference profile at layer k - (; pres) = surface.timesteps[lf] # logarithm of surface pressure - (; geopot) = diagn.dynamics_variables + (; R_dry) = atmosphere # dry gas constant + (; temp_profile) = implicit # reference profile at layer k + pres = progn.pres[lf] # logarithm of surface pressure at leapfrog index lf + (; geopot) = diagn.dynamics # -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 # vordiv_tendencies! include as R_dry*Tₖ*lnpₛ into the geopotential on which the operator # -∇² is applied in bernoulli_potential! - @. geopot += R_dry*Tₖ*pres + @inbounds for k in eachmatrix(geopot) + R_dryTₖ = R_dry*temp_profile[k] + for lm in eachharmonic(pres) + geopot[lm, k] += R_dryTₖ*pres[lm] + end + end end """ $(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, - orog::AbstractOrography, - atmosphere::AbstractAtmosphere, - G::AbstractGeometry, - S::SpectralTransform) - - (; pres_grid, pres_tend ) = surface +function volume_flux_divergence!( + diagn::DiagnosticVariables, + orog::AbstractOrography, + atmosphere::AbstractAtmosphere, + G::AbstractGeometry, + S::SpectralTransform +) + + (; pres_grid) = diagn.grid + (; pres_tend ) = diagn.tendencies (; orography ) = orog H = atmosphere.layer_thickness @@ -758,52 +717,25 @@ function volume_flux_divergence!( diagn::DiagnosticVariablesLayer, flux_divergence!(pres_tend, pres_grid, diagn, G, S, add=true, flipsign=true) end -""" -$(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; - kwargs... -) - - # all variables on layers - for (progn_layer, diagn_layer) in zip(progn.layers, diagn.layers) - progn_layer_lf = progn_layer.timesteps[lf] - gridded!(diagn_layer, progn_layer_lf, model; kwargs...) - end - - # surface only for ShallowWaterModel or PrimitiveEquation - S = model.spectral_transform - (; pres_grid) = diagn.surface - (; pres) = progn.surface.timesteps[lf] - model isa Barotropic || gridded!(pres_grid, pres, S) - - return nothing -end - """ $(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.""" -function SpeedyTransforms.gridded!( - diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, +function SpeedyTransforms.transform!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, model::Barotropic; kwargs... -) - - (; vor_grid, u_grid, v_grid ) = diagn.grid_variables - (; vor ) = progn # relative vorticity - U = diagn.dynamics_variables.a # reuse work arrays for velocities in spectral - V = diagn.dynamics_variables.b # U = u*coslat, V=v*coslat +) + (; vor_grid, u_grid, v_grid ) = diagn.grid + vor = progn.vor[lf] # relative vorticity at leapfrog step lf + U = diagn.dynamics.a # reuse work arrays for velocities in spectral + V = diagn.dynamics.b # U = u*coslat, V=v*coslat S = model.spectral_transform - gridded!(vor_grid, vor, S) # get vorticity on grid from spectral vor + transform!(vor_grid, vor, S) # get vorticity on grid from spectral vor # get spectral U, V from spectral vorticity via stream function Ψ # U = u*coslat = -coslat*∂Ψ/∂lat @@ -811,8 +743,8 @@ function SpeedyTransforms.gridded!( UV_from_vor!(U, V, vor, S) # transform from U, V in spectral to u, v on grid (U, V = u, v*coslat) - gridded!(u_grid, U, S, unscale_coslat=true) - gridded!(v_grid, V, S, unscale_coslat=true) + transform!(u_grid, U, S, unscale_coslat=true) + transform!(v_grid, V, S, unscale_coslat=true) return nothing end @@ -823,30 +755,34 @@ 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.""" -function SpeedyTransforms.gridded!( - diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, +function SpeedyTransforms.transform!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, model::ShallowWater; kwargs... ) - - (; vor_grid, div_grid, u_grid, v_grid ) = diagn.grid_variables - (; vor, div) = progn - U = diagn.dynamics_variables.a # reuse work arrays for velocities spectral - V = diagn.dynamics_variables.b # U = u*coslat, V=v*coslat - S = model.spectral_transform + (; vor_grid, div_grid, pres_grid, u_grid, v_grid ) = diagn.grid + vor = progn.vor[lf] # relative vorticity at leapfrog step lf + div = progn.div[lf] # divergence at leapfrog step lf + pres = progn.pres[lf] # interface displacement η at leapfrog step lf + U = diagn.dynamics.a # reuse work arrays for velocities spectral + V = diagn.dynamics.b # U = u*coslat, V=v*coslat + S = model.spectral_transform + + transform!(vor_grid, vor, S) # get vorticity on grid from spectral vor + transform!(div_grid, div, S) # get divergence on grid from spectral div + transform!(pres_grid, pres, S) # get η on grid from spectral η + # get spectral U, V from vorticity and divergence via stream function Ψ and vel potential ϕ # U = u*coslat = -coslat*∂Ψ/∂lat + ∂ϕ/dlon # V = v*coslat = coslat*∂ϕ/∂lat + ∂Ψ/dlon UV_from_vordiv!(U, V, vor, div, S) - gridded!(vor_grid, vor, S) # get vorticity on grid from spectral vor - gridded!(div_grid, div, S) # get divergence on grid from spectral div - # transform from U, V in spectral to u, v on grid (U, V = u, v*coslat) - gridded!(u_grid, U, S, unscale_coslat=true) - gridded!(v_grid, V, S, unscale_coslat=true) + transform!(u_grid, U, S, unscale_coslat=true) + transform!(v_grid, V, S, unscale_coslat=true) return nothing end @@ -857,19 +793,25 @@ 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::PrognosticVariablesLayer, +function SpeedyTransforms.transform!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, model::PrimitiveEquation; - initialize=false -) - - (; vor_grid, div_grid, u_grid, v_grid ) = diagn.grid_variables - (; temp_grid, humid_grid ) = diagn.grid_variables - (; temp_grid_prev, humid_grid_prev, u_grid_prev, v_grid_prev) = diagn.grid_variables - (; vor, div, temp, humid) = progn - U = diagn.dynamics_variables.a # reuse work arrays for velocities spectral - V = diagn.dynamics_variables.b # U = u*coslat, V=v*coslat + initialize::Bool = false, +) + (; vor_grid, div_grid, pres_grid, u_grid, v_grid, temp_grid, humid_grid) = diagn.grid + (; temp_grid_prev, humid_grid_prev, u_grid_prev, v_grid_prev) = diagn.grid + + vor = progn.vor[lf] # relative vorticity at leapfrog step lf + div = progn.div[lf] # divergence at leapfrog step lf + temp = progn.temp[lf] # temperature at leapfrog step lf + humid = progn.humid[lf] # humidity at leapfrog step lf + pres = progn.pres[lf] # logarithm of surface pressure at leapfrog step lf + + U = diagn.dynamics.a # reuse work arrays + V = diagn.dynamics.b # U = u*coslat, V=v*coslat + S = model.spectral_transform # retain previous time step for vertical advection and some parameterizations if initialize == false # only store prev after initial step @@ -879,40 +821,45 @@ function SpeedyTransforms.gridded!( @. v_grid_prev = v_grid end - S = model.spectral_transform - wet_core = model isa PrimitiveWet + transform!(vor_grid, vor, S) # get vorticity on grid from spectral vor + transform!(div_grid, div, S) # get divergence on grid from spectral div + transform!(temp_grid, temp, S) # -- temperature -- + transform!(pres_grid, pres, S) # -- pressure -- + + if model isa PrimitiveWet + transform!(humid_grid, humid, S) + hole_filling!(humid_grid, model.hole_filling, model) # remove negative humidity + end # get spectral U, V from vorticity and divergence via stream function Ψ and vel potential ϕ # U = u*coslat = -coslat*∂Ψ/∂lat + ∂ϕ/dlon # V = v*coslat = coslat*∂ϕ/∂lat + ∂Ψ/dlon UV_from_vordiv!(U, V, vor, div, S) - - gridded!(vor_grid, vor, S) # get vorticity on grid from spectral vor - gridded!(div_grid, div, S) # get divergence on grid from spectral div - gridded!(temp_grid, temp, S) # (absolute) temperature - if wet_core # specific humidity (wet core only) - gridded!(humid_grid, humid, S) - hole_filling!(humid_grid, model.hole_filling, model) # remove negative humidity - end - + # transform from U, V in spectral to u, v on grid (U, V = u, v*coslat) + transform!(u_grid, U, S, unscale_coslat=true) + transform!(v_grid, V, S, unscale_coslat=true) + # include humidity effect into temp for everything stability-related temperature_average!(diagn, temp, S) - virtual_temperature!(diagn, temp, model) # temp = virt temp for dry core - - # transform from U, V in spectral to u, v on grid (U, V = u, v*coslat) - gridded!(u_grid, U, S, unscale_coslat=true) - gridded!(v_grid, V, S, unscale_coslat=true) + virtual_temperature!(diagn, model) # temp = virt temp for dry core if initialize # at initial step store prev <- current - # subtract the reference temperature profile as temp_grid is too after every time step - @. temp_grid_prev = temp_grid - model.implicit.temp_profile[diagn.k] + # subtract the reference temperature profile Tₖ as temp_grid is too after every time step + # the Tₖ is added for the physics parameterizations again + # technically Tₖ is from model.implicit (which is held constant throughout) integration + # but given it's the initial step here using the instantaneous diagn.temp_average is the same + @inbounds for k in eachgrid(temp_grid_prev, temp_grid) + Tₖ = diagn.temp_average[k] + for ij in eachgridpoint(temp_grid_prev, temp_grid) + temp_grid_prev[ij, k] = temp_grid[ij, k] - Tₖ + end + end + @. humid_grid_prev = humid_grid @. u_grid_prev = u_grid @. v_grid_prev = v_grid end - - return nothing end """ @@ -920,10 +867,12 @@ $(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, + diagn::DiagnosticVariables, + temp::LowerTriangularArray, S::SpectralTransform, ) - # average from l=m=0 harmonic divided by norm of the sphere - diagn.temp_average[] = real(temp[1])/S.norm_sphere + @inbounds for k in eachmatrix(temp) + # average from l=m=0 harmonic divided by norm of the sphere + diagn.temp_average[k] = real(temp[1, k])/S.norm_sphere + end end \ No newline at end of file diff --git a/src/dynamics/time_integration.jl b/src/dynamics/time_integration.jl index d1393a0b2..8ff4372c3 100644 --- a/src/dynamics/time_integration.jl +++ b/src/dynamics/time_integration.jl @@ -1,23 +1,27 @@ +const DEFAULT_NSTEPS = 2 export Leapfrog """ Leapfrog time stepping defined by the following fields $(TYPEDFIELDS) """ -Base.@kwdef mutable struct Leapfrog{NF<:AbstractFloat} <: AbstractTimeStepper +@kwdef mutable struct Leapfrog{NF<:AbstractFloat} <: AbstractTimeStepper # DIMENSIONS "spectral resolution (max degree of spherical harmonics)" trunc::Int + "Number of timesteps stored simultaneously in prognostic variables" + nsteps::Int = 2 + # OPTIONS - "time step in minutes for T31, scale linearly to `trunc`" + "Time step in minutes for T31, scale linearly to `trunc`" Δt_at_T31::Second = Minute(30) - "radius of sphere [m], used for scaling" + "Radius of sphere [m], used for scaling" radius::NF = DEFAULT_RADIUS - "adjust Δt_at_T31 with the output_dt to reach output_dt exactly in integer time steps" + "Adjust Δt_at_T31 with the output_dt to reach output_dt exactly in integer time steps" adjust_with_output::Bool = true # NUMERICS @@ -100,7 +104,7 @@ Initialize leapfrogging `L` by recalculating the timestep given the output time `output_dt` from `model.output`. Recalculating will slightly adjust the time step to be a divisor such that an integer number of time steps matches exactly with the output time step.""" -function initialize!(L::Leapfrog, model::ModelSetup) +function initialize!(L::Leapfrog, model::AbstractModel) (; output_dt) = model.output if L.adjust_with_output @@ -122,19 +126,21 @@ end $(TYPEDSIGNATURES) Performs one leapfrog time step with (`lf=2`) or without (`lf=1`) Robert+Williams filter (see Williams (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 Williams filter - L::Leapfrog{NF}, # struct with constants - ) where {NF<:AbstractFloat} # number format NF +function leapfrog!( + A_old::LowerTriangularArray, # prognostic variable at t + A_new::LowerTriangularArray, # prognostic variable at t+dt + tendency::LowerTriangularArray, # 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 Williams filter + L::Leapfrog{NF}, # struct with constants +) where NF # 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 Williams filter - (; robert_filter, williams_filter) = L # coefficients for the Robert and Williams filter - dt_NF = convert(NF, dt) # time step dt in number format NF + @boundscheck size(A_old) == size(A_new) == size(tendency) || throw(BoundsError()) + + A_lf = lf == 1 ? A_old : A_new # view on either t or t+dt to dis/enable Williams filter + (; robert_filter, williams_filter) = L # coefficients for the Robert and Williams filter + dt_NF = convert(NF, dt) # time step dt in number format NF # LEAP FROG time step with or without Robert+Williams filter # Robert time filter to compress computational mode, Williams filter for 3rd order accuracy @@ -144,7 +150,7 @@ function leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}}, # prognostic w1 = lf == 1 ? zero(NF) : robert_filter*williams_filter/2 # = ν*α/2 in Williams (2009, Eq. 8) w2 = lf == 1 ? zero(NF) : robert_filter*(1-williams_filter)/2 # = ν(1-α)/2 in Williams (2009, Eq. 9) - @inbounds for lm in eachharmonic(A_old, A_new, A_lf, tendency) + @inbounds for lm in eachindex(A_old, A_new, tendency) a_old = A_old[lm] # double filtered value from previous time step (t-Δt) a_new = a_old + dt_NF*tendency[lm] # Leapfrog/Euler step depending on dt=Δt, 2Δt (unfiltered at t+Δt) a_update = a_old - 2A_lf[lm] + a_new # Eq. 8&9 in Williams (2009), calculate only once @@ -153,40 +159,24 @@ function leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}}, # prognostic end 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(::PrimitiveDry) = (:vor, :div, :temp) -leapfrog_layer_vars(::PrimitiveWet) = (:vor, :div, :temp, :humid) - -function leapfrog!( progn::PrognosticLayerTimesteps, - diagn::DiagnosticVariablesLayer, - dt::Real, # time step (mostly =2Δt, but for init steps =Δt, Δt/2) - lf::Int, # leapfrog index to dis/enable Williams filter - model::ModelSetup) - - for var in leapfrog_layer_vars(model) - var_old = getproperty(progn.timesteps[1], var) - var_new = getproperty(progn.timesteps[2], var) - var_tend = getproperty(diagn.tendencies, Symbol(var, :_tend)) - spectral_truncation!(var_tend) # set lmax+1 mode to zero +# variables that are leapfrogged in the respective models, e.g. :vor_tend, :div_tend, etc... +tendency_names(model::AbstractModel) = tuple((Symbol(var, :_tend) for var in prognostic_variables(model))...) + +function leapfrog!( + progn::PrognosticVariables, + tend::Tendencies, + dt::Real, # time step (mostly =2Δt, but for init steps =Δt, Δt/2) + lf::Int, # leapfrog index to dis/enable Williams filter + model::AbstractModel, +) + for (varname, tendname) in zip(prognostic_variables(model), tendency_names(model)) + var_old, var_new = getfield(progn, varname) + var_tend = getfield(tend, tendname) + spectral_truncation!(var_tend) leapfrog!(var_old, var_new, var_tend, dt, lf, model.time_stepping) end end -function leapfrog!( progn::PrognosticSurfaceTimesteps, - diagn::SurfaceVariables, - dt::Real, # time step (mostly =2Δt, but for init steps =Δt, Δt/2) - lf::Int, # leapfrog index to dis/enable Williams filter - model::ModelSetup) - - (; pres_tend) = diagn - pres_old = progn.timesteps[1].pres - pres_new = progn.timesteps[2].pres - spectral_truncation!(pres_tend) # set lmax+1 mode to zero - leapfrog!(pres_old, pres_new, pres_tend, dt, lf, model.time_stepping) -end - """ $(TYPEDSIGNATURES) Performs the first two initial time steps (Euler forward, unfiltered leapfrog) to populate the @@ -194,7 +184,7 @@ prognostic variables with two time steps (t=0, Δt) that can then be used in the function first_timesteps!( progn::PrognosticVariables, # all prognostic variables diagn::DiagnosticVariables, # all pre-allocated diagnostic variables - model::ModelSetup, # everything that is constant at runtime + model::AbstractModel, # everything that is constant at runtime ) (; clock) = progn clock.n_timesteps == 0 && return nothing # exit immediately for no time steps @@ -225,7 +215,7 @@ function first_timesteps!( timestep!(clock, Δt_millisec) # do output and callbacks after the first proper (from i=0 to i=1) time step - write_output!(model.output, clock.time, diagn) + output!(model.output, progn, diagn, model) callback!(model.callbacks, progn, diagn, model) # from now on precomputed implicit terms with 2Δt @@ -234,31 +224,26 @@ function first_timesteps!( return nothing end -""" -$(TYPEDSIGNATURES) -Calculate a single time step for the `model <: Barotropic`.""" +"""$(TYPEDSIGNATURES) +Calculate a single time step for the barotropic model.""" function timestep!( progn::PrognosticVariables, # all prognostic variables diagn::DiagnosticVariables, # all pre-allocated diagnostic variables dt::Real, # time step (mostly =2Δt, but for first_timesteps! =Δt, Δt/2) model::Barotropic, # everything that's constant at runtime - lf1::Integer=2, # leapfrog index 1 (dis/enables Robert+Williams filter) - lf2::Integer=2, # leapfrog index 2 (time step used for tendencies) + lf1::Integer = 2, # leapfrog index 1 (dis/enables Robert+Williams filter) + lf2::Integer = 2, # leapfrog index 2 (time step used for tendencies) ) - model.feedback.nars_detected && return nothing # exit immediately if NaRs already present - (; time) = progn.clock # current time + model.feedback.nars_detected && return nothing # exit immediately if NaNs/Infs already present # set the tendencies back to zero for accumulation - 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) - progn_lf = progn_layer.timesteps[lf2] # pick the leapfrog time step lf2 for tendencies - dynamics_tendencies!(diagn_layer, progn_lf, time, model) - horizontal_diffusion!(diagn_layer, progn_layer, model) - leapfrog!(progn_layer, diagn_layer, dt, lf1, model) - gridded!(diagn_layer, progn_lf, model) - end + fill!(diagn.tendencies, 0, Barotropic) + + # TENDENCIES, DIFFUSION, LEAPFROGGING AND TRANSFORM SPECTRAL STATE TO GRID + dynamics_tendencies!(diagn, progn, lf2, model) + horizontal_diffusion!(diagn, progn, model) + leapfrog!(progn, diagn.tendencies, dt, lf1, model) + transform!(diagn, progn, lf2, model) # PARTICLE ADVECTION (always skip 1st step of first_timesteps!) not_first_timestep = lf2 == 2 @@ -273,36 +258,22 @@ function timestep!( diagn::DiagnosticVariables, # all pre-allocated diagnostic variables dt::Real, # time step (mostly =2Δt, but for first_timesteps! =Δt, Δt/2) model::ShallowWater, # everything that's constant at runtime - lf1::Integer=2, # leapfrog index 1 (dis/enables Robert+Williams filter) - lf2::Integer=2, # leapfrog index 2 (time step used for tendencies) + lf1::Integer = 2, # leapfrog index 1 (dis/enables Robert+Williams filter) + lf2::Integer = 2, # leapfrog index 2 (time step used for tendencies) ) - model.feedback.nars_detected && return nothing # exit immediately if NaRs already present - (; time) = progn.clock # current time # set the tendencies back to zero for accumulation - zero_tendencies!(diagn) - - progn_layer = progn.layers[1] # only calculate tendencies for the first layer - diagn_layer = diagn.layers[1] # multi-layer shallow water not supported - - progn_lf = progn_layer.timesteps[lf2] # pick the leapfrog time step lf2 for tendencies - (; pres) = progn.surface.timesteps[lf2] - (; implicit, spectral_transform) = model + fill!(diagn.tendencies, 0, ShallowWater) # GET TENDENCIES, CORRECT THEM FOR SEMI-IMPLICIT INTEGRATION - dynamics_tendencies!(diagn_layer, progn_lf, diagn.surface, pres, time, model) - implicit_correction!(diagn_layer, progn_layer, diagn.surface, progn.surface, implicit) + dynamics_tendencies!(diagn, progn, lf2, model) + implicit_correction!(diagn, progn, model.implicit) # APPLY DIFFUSION, STEP FORWARD IN TIME, AND TRANSFORM NEW TIME STEP TO GRID - horizontal_diffusion!(progn_layer, diagn_layer, model) - leapfrog!(progn_layer, diagn_layer, dt, lf1, model) - gridded!(diagn_layer, progn_lf, model) - - # SURFACE LAYER (pressure), no diffusion though - (; pres_grid) = diagn.surface - leapfrog!(progn.surface, diagn.surface, dt, lf1, model) - gridded!(pres_grid, pres, spectral_transform) + horizontal_diffusion!(diagn, progn, model) + leapfrog!(progn, diagn.tendencies, dt, lf1, model) + transform!(diagn, progn, lf2, model) # PARTICLE ADVECTION (always skip 1st step of first_timesteps!) not_first_timestep = lf2 == 2 @@ -317,54 +288,37 @@ function timestep!( diagn::DiagnosticVariables, # all pre-allocated diagnostic variables dt::Real, # time step (mostly =2Δt, but for first_timesteps! =Δt, Δt/2) model::PrimitiveEquation, # everything that's constant at runtime - lf1::Integer=2, # leapfrog index 1 (dis/enables Robert+Williams filter) - lf2::Integer=2, # leapfrog index 2 (time step used for tendencies) + lf1::Integer = 2, # leapfrog index 1 (dis/enables Robert+Williams filter) + lf2::Integer = 2, # leapfrog index 2 (time step used for tendencies) ) model.feedback.nars_detected && return nothing # exit immediately if NaRs already present (; time) = progn.clock # current time # set the tendencies back to zero for accumulation - zero_tendencies!(diagn) + fill!(diagn.tendencies, 0, PrimitiveWet) if model.physics # switch on/off all physics parameterizations # time step ocean (temperature and TODO sea ice) and land (temperature and soil moisture) - ocean_timestep!(progn.ocean, time, model) - land_timestep!(progn.land, time, model) - soil_moisture_availability!(diagn.surface, progn.land, model) + ocean_timestep!(progn, diagn, model) + land_timestep!(progn, diagn, model) + soil_moisture_availability!(diagn, progn, model) # calculate all parameterizations parameterization_tendencies!(diagn, progn, time, model) end - if model.dynamics # switch on/off all dynamics - dynamics_tendencies!(diagn, progn, model, lf2) # dynamical core - implicit_correction!(diagn, model.implicit, progn) # semi-implicit time stepping corrections + if model.dynamics # switch on/off all dynamics + dynamics_tendencies!(diagn, progn, lf2, model) # dynamical core + implicit_correction!(diagn, model.implicit, progn) # semi-implicit time stepping corrections else # just transform physics tendencies to spectral space - for k in 1:diagn.nlev - diagn_layer = diagn.layers[k] - tendencies_physics_only!(diagn_layer, model) - end + physics_tendencies_only!(diagn, model) end - # LOOP OVER ALL LAYERS for diffusion, leapfrog time integration - # and progn state from spectral to grid for next time step - @floop for k in 1:diagn.nlev+1 - if k <= diagn.nlev # model levels - diagn_layer = diagn.layers[k] - progn_layer = progn.layers[k] - progn_layer_lf = progn_layer.timesteps[lf2] - - horizontal_diffusion!(progn_layer, diagn_layer, model) # for vor, div, temp, humid - leapfrog!(progn_layer, diagn_layer, dt, lf1, model) # time step forward for vor, div, temp, humid - gridded!(diagn_layer, progn_layer_lf, model) # propagate spectral state to grid - else # surface level, time step for pressure - leapfrog!(progn.surface, diagn.surface, dt, lf1, model) - (; pres_grid) = diagn.surface - pres_lf = progn.surface.timesteps[lf2].pres - gridded!(pres_grid, pres_lf, model.spectral_transform) - end - end + # APPLY DIFFUSION, STEP FORWARD IN TIME, AND TRANSFORM NEW TIME STEP TO GRID + horizontal_diffusion!(diagn, progn, model) + leapfrog!(progn, diagn.tendencies, dt, lf1, model) + transform!(diagn, progn, lf2, model) # PARTICLE ADVECTION (always skip 1st step of first_timesteps!) not_first_timestep = lf2 == 2 @@ -378,12 +332,11 @@ 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 model components + model::AbstractModel, # all model components ) (; clock) = progn (; Δt, Δt_millisec) = model.time_stepping - (; time_stepping) = model # SCALING: we use vorticity*radius, divergence*radius in the dynamical core scale!(progn, diagn, model.spectral_grid.radius) @@ -392,9 +345,9 @@ function time_stepping!( # propagate spectral state to grid variables for initial condition output (; output, feedback) = model lf = 1 # use first leapfrog index - gridded!(diagn, progn, lf, model, initialize=true) + transform!(diagn, progn, lf, model, initialize=true) initialize!(progn.particles, progn, diagn, model.particle_advection) - initialize!(output, feedback, time_stepping, clock, diagn, model) + initialize!(output, feedback, progn, diagn, model) initialize!(model.callbacks, progn, diagn, model) # FIRST TIMESTEPS: EULER FORWARD THEN 1x LEAPFROG @@ -410,7 +363,7 @@ function time_stepping!( timestep!(clock, Δt_millisec) # time of lf=2 and diagn after timestep! progress!(feedback, progn) # updates the progress meter bar - write_output!(output, clock.time, diagn) + output!(output, progn, diagn, model) callback!(model.callbacks, progn, diagn, model) end @@ -419,8 +372,10 @@ function time_stepping!( unscale!(progn) # undo radius-scaling for vor, div from the dynamical core unscale!(diagn) # undo radius-scaling for vor, div from the dynamical core close(output) # close netCDF file - write_restart_file(progn, output) # as JLD2 + write_restart_file(output, progn) # as JLD2 finish!(model.callbacks, progn, diagn, model) - return progn # to trigger UnicodePlot via show(::IO, ::PrognosticVariables) + # return a UnicodePlot of surface vorticity + surface_vorticity = diagn.grid.vor_grid[:, end] + return plot(surface_vorticity, title="Surface relative vorticity [1/s]") end \ No newline at end of file diff --git a/src/dynamics/vertical_advection.jl b/src/dynamics/vertical_advection.jl index 4e84df1a4..2929d3fac 100644 --- a/src/dynamics/vertical_advection.jl +++ b/src/dynamics/vertical_advection.jl @@ -20,91 +20,88 @@ WENOVerticalAdvection(spectral_grid) = WENOVerticalAdvection{ @inline retrieve_time_step(::DiffusiveVerticalAdvection, variables, var) = retrieve_previous_time_step(variables, var) @inline retrieve_time_step(::DispersiveVerticalAdvection, variables, var) = retrieve_current_time_step(variables, var) -@inline function retrieve_current_stencil(k, layers, var, nlev, ::VerticalAdvection{NF, B}) where {NF, B} - k_stencil = max.(min.(nlev, k-B:k+B), 1) - ξ_stencil = Tuple(retrieve_current_time_step(layers[k].grid_variables, var) for k in k_stencil) - return ξ_stencil +@inline function retrieve_stencil(k, nlayers, ::VerticalAdvection{NF, B}) where {NF, B} + # creates allocation-free tuples for k-B:k+B but clamped into (1, nlayers) + # e.g. (1, 1, 2), (1, 2, 3), (2, 3, 4) ... (for k=1, 2, 3; B=1) + return ntuple(i -> clamp(i+k-B-1, 1, nlayers), 2B+1) end -@inline function retrieve_previous_stencil(k, layers, var, nlev, ::VerticalAdvection{NF, B}) where {NF, B} - k_stencil = max.(min.(nlev, k-B:k+B), 1) - ξ_stencil = Tuple(retrieve_previous_time_step(layers[k].grid_variables, var) for k in k_stencil) - return ξ_stencil -end - -@inline retrieve_stencil(k, layers, var, nlev, scheme::DiffusiveVerticalAdvection) = retrieve_previous_stencil(k, layers, var, nlev, scheme) -@inline retrieve_stencil(k, layers, var, nlev, scheme::DispersiveVerticalAdvection) = retrieve_current_stencil(k, layers, var, nlev, scheme) - -function vertical_advection!( layer::DiagnosticVariablesLayer, - diagn::DiagnosticVariables, - model::PrimitiveEquation) - - (; k ) = layer # which layer are we on? - - wet_core = model isa PrimitiveWet - (; σ_levels_thick, nlev ) = model.geometry - - scheme = model.vertical_advection - - # 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⁻ = max(1, k-1) # just saturate, because M_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_above = diagn.layers[k⁻].dynamics_variables.σ_tend - σ_tend_below = layer.dynamics_variables.σ_tend - - # layer thickness Δσ on level k - Δσₖ = σ_levels_thick[k] +function vertical_advection!( + diagn::DiagnosticVariables, + model::PrimitiveEquation, +) + Δσ = model.geometry.σ_levels_thick + advection_scheme = model.vertical_advection + (; σ_tend) = diagn.dynamics for var in (:u, :v, :temp) - ξ_tend = getproperty(layer.tendencies, Symbol(var, :_tend_grid)) - ξ_sten = retrieve_stencil(k, diagn.layers, var, nlev, scheme) - ξ = retrieve_time_step(scheme, layer.grid_variables, var) - - _vertical_advection!(ξ_tend, σ_tend_above, σ_tend_below, ξ_sten, ξ, Δσₖ, scheme) + ξ_tend = getproperty(diagn.tendencies, Symbol(var, :_tend_grid)) + ξ = retrieve_time_step(advection_scheme, diagn.grid, var) + _vertical_advection!(ξ_tend, σ_tend, ξ, Δσ, advection_scheme) end - if wet_core - ξ_tend = getproperty(layer.tendencies, :humid_tend_grid) - ξ_sten = retrieve_current_stencil(k, diagn.layers, :humid, nlev, scheme) - ξ = retrieve_current_time_step(layer.grid_variables, :humid) - - _vertical_advection!(ξ_tend, σ_tend_above, σ_tend_below, ξ_sten, ξ, Δσₖ, scheme) + if model isa PrimitiveWet # advect humidity only with primitive wet core + ξ_tend = getproperty(diagn.tendencies, :humid_tend_grid) + ξ = retrieve_time_step(advection_scheme, diagn.grid, :humid) + _vertical_advection!(ξ_tend, σ_tend, ξ, Δσ, advection_scheme) end 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 - σ_tend_below::Grid, # vertical velocity at k+1/2 - ξ_sten, # ξ stencil for vertical advection (from k-B to k+B) - ξ::Grid, # ξ at level k - Δσₖ::NF, # layer thickness on σ levels - adv::VerticalAdvection{NF, B} # vertical advection scheme - ) where {NF<:AbstractFloat, Grid<:AbstractGrid{NF}, B} - Δσₖ⁻¹ = 1/Δσₖ # precompute - - # += as the tendencies already contain the parameterizations - for ij in eachgridpoint(ξ_tend) - σ̇⁻ = σ_tend_above[ij] # velocity into layer k from above - σ̇⁺ = σ_tend_below[ij] # velocity out of layer k to below - - ξᶠ⁺ = reconstructed_at_face(ij, adv, σ̇⁺, ξ_sten[2:end]) - ξᶠ⁻ = reconstructed_at_face(ij, adv, σ̇⁻, ξ_sten[1:end-1]) - - ξ_tend[ij] -= Δσₖ⁻¹ * (σ̇⁺ * ξᶠ⁺ - σ̇⁻ * ξᶠ⁻ - ξ[ij] * (σ̇⁺ - σ̇⁻)) +function _vertical_advection!( + ξ_tend::AbstractGridArray, # tendency of quantity ξ + σ_tend::AbstractGridArray, # vertical velocity at k+1/2 + ξ::AbstractGridArray, # ξ + Δσ, # layer thickness on σ levels + adv::VerticalAdvection # vertical advection scheme of order B +) + grids_match(ξ_tend, σ_tend, ξ) || throw(DimensionMismatch(ξ_tend, σ_tend, ξ)) + + nlayers = size(ξ, 2) + @inbounds for k in 1:nlayers + Δσₖ⁻¹ = inv(Δσ[k]) # inverse layer thickness, compute inv only once + + # for k=1 "above" term (at k-1/2) is 0, for k==nlayers "below" term (at k+1/2) is zero + # avoid out-of-bounds indexing with k⁻, k⁺ + k⁻ = max(1, k-1) # TODO check that this actually zeros velocity at k=1/2 + k⁺ = k + + k_stencil = retrieve_stencil(k, nlayers, adv) + + for ij in eachgridpoint(ξ_tend) + σ̇⁻ = σ_tend[ij, k⁻] # velocity into layer k from above + σ̇⁺ = σ_tend[ij, k⁺] # velocity out of layer k to below + + ξᶠ⁺ = reconstructed_at_face(ξ, ij, k_stencil[2:end], σ̇⁺, adv) + ξᶠ⁻ = reconstructed_at_face(ξ, ij, k_stencil[1:end-1], σ̇⁻, adv) + + # -= as the tendencies already contain the parameterizations + ξ_tend[ij, k] -= Δσₖ⁻¹ * (σ̇⁺ * ξᶠ⁺ - σ̇⁻ * ξᶠ⁻ - ξ[ij, k] * (σ̇⁺ - σ̇⁻)) + end end end -@inline reconstructed_at_face(ij, ::UpwindVerticalAdvection{NF, 1}, u, ξ) where NF = ifelse(u > 0, ξ[1][ij], ξ[2][ij]) -@inline reconstructed_at_face(ij, ::UpwindVerticalAdvection{NF, 2}, u, ξ) where NF = ifelse(u > 0, (2ξ[1][ij] + 5ξ[2][ij] - ξ[3][ij]) / 6, - (2ξ[4][ij] + 5ξ[3][ij] - ξ[2][ij]) / 6) -@inline reconstructed_at_face(ij, ::UpwindVerticalAdvection{NF, 3}, u, ξ) where NF = ifelse(u > 0, (2ξ[1][ij] - 13ξ[2][ij] + 47ξ[3][ij] + 27ξ[4][ij] - 3ξ[5][ij]) / 60, - (2ξ[6][ij] - 13ξ[5][ij] + 47ξ[4][ij] + 27ξ[3][ij] - 3ξ[2][ij]) / 60) +# 1st order upwind +@inline reconstructed_at_face(ξ, ij, k, u, ::UpwindVerticalAdvection{NF, 1}) where NF = + ifelse(u > 0, ξ[ij, k[1]], + ξ[ij, k[2]]) + +# 3rd order upwind +@inline reconstructed_at_face(ξ, ij, k, u, ::UpwindVerticalAdvection{NF, 2}) where NF = + ifelse(u > 0, (2ξ[ij, k[1]] + 5ξ[ij, k[2]] - ξ[ij, k[3]]) / 6, + (2ξ[ij, k[4]] + 5ξ[ij, k[3]] - ξ[ij, k[2]]) / 6) + +# 5th order upwind +@inline reconstructed_at_face(ξ, ij, k, u, ::UpwindVerticalAdvection{NF, 3}) where NF = + ifelse(u > 0, (2ξ[ij, k[1]] - 13ξ[ij, k[2]] + 47ξ[ij, k[3]] + 27ξ[ij, k[4]] - 3ξ[ij, k[5]]) / 60, + (2ξ[ij, k[6]] - 13ξ[ij, k[5]] + 47ξ[ij, k[4]] + 27ξ[ij, k[3]] - 3ξ[ij, k[2]]) / 60) + +# 2nd order centered +@inline reconstructed_at_face(ξ, ij, k, u, ::CenteredVerticalAdvection{NF, 1}) where NF = + (ξ[ij, k[1]] + ξ[ij, k[2]]) / 2 -@inline reconstructed_at_face(ij, ::CenteredVerticalAdvection{NF, 1}, u, ξ) where NF = ( ξ[1][ij] + ξ[2][ij]) / 2 -@inline reconstructed_at_face(ij, ::CenteredVerticalAdvection{NF, 2}, u, ξ) where NF = (-ξ[1][ij] + 7ξ[2][ij] + 7ξ[3][ij] - ξ[4][ij]) / 12 +# 4th order centered +@inline reconstructed_at_face(ξ, ij, k, u, ::CenteredVerticalAdvection{NF, 2}) where NF = + (-ξ[ij, k[1]] + 7ξ[ij, k[2]] + 7ξ[ij, k[3]] - ξ[ij, k[4]]) / 12 const ε = 1e-6 const d₀ = 3/10 @@ -135,15 +132,15 @@ const d₂ = 1/10 return p₀(S₀) * w₀ + p₁(S₁) * w₁ + p₂(S₂) * w₂ end -@inline function reconstructed_at_face(ij, ::WENOVerticalAdvection{NF}, u, ξ) where NF +@inline function reconstructed_at_face(ξ, ij, k, u, ::WENOVerticalAdvection{NF}) where NF if u > 0 - S₀ = (ξ[3][ij], ξ[4][ij], ξ[5][ij]) - S₁ = (ξ[2][ij], ξ[3][ij], ξ[4][ij]) - S₂ = (ξ[1][ij], ξ[2][ij], ξ[3][ij]) + S₀ = (ξ[ij, k[3]], ξ[ij, k[4]], ξ[ij, k[5]]) + S₁ = (ξ[ij, k[2]], ξ[ij, k[3]], ξ[ij, k[4]]) + S₂ = (ξ[ij, k[1]], ξ[ij, k[2]], ξ[ij, k[3]]) else - S₀ = (ξ[4][ij], ξ[3][ij], ξ[2][ij]) - S₁ = (ξ[5][ij], ξ[4][ij], ξ[3][ij]) - S₂ = (ξ[6][ij], ξ[5][ij], ξ[4][ij]) + S₀ = (ξ[ij, k[4]], ξ[ij, k[3]], ξ[ij, k[2]]) + S₁ = (ξ[ij, k[5]], ξ[ij, k[4]], ξ[ij, k[3]]) + S₂ = (ξ[ij, k[6]], ξ[ij, k[5]], ξ[ij, k[4]]) end return weno_reconstruction(S₀, S₁, S₂, NF) end diff --git a/src/dynamics/vertical_coordinates.jl b/src/dynamics/vertical_coordinates.jl index 4a5a0e1b1..d96edd041 100644 --- a/src/dynamics/vertical_coordinates.jl +++ b/src/dynamics/vertical_coordinates.jl @@ -1,46 +1,46 @@ abstract type VerticalCoordinates end -Base.@kwdef struct NoVerticalCoordinates <: VerticalCoordinates - nlev::Int = 1 +@kwdef struct NoVerticalCoordinates <: VerticalCoordinates + nlayers::Int = 1 end export SigmaCoordinates -Base.@kwdef struct SigmaCoordinates <: VerticalCoordinates - nlev::Int = 8 - σ_half::Vector{Float64} = default_sigma_coordinates(nlev) +@kwdef struct SigmaCoordinates <: VerticalCoordinates + nlayers::Int = 8 + σ_half::Vector{Float64} = default_sigma_coordinates(nlayers) - SigmaCoordinates(nlev::Integer, σ_half::AbstractVector) = sigma_okay(nlev, σ_half) ? - new(nlev, σ_half) : error("σ_half = $σ_half cannot be used for $nlev-level SigmaCoordinates") + SigmaCoordinates(nlayers::Integer, σ_half::AbstractVector) = sigma_okay(nlayers, σ_half) ? + new(nlayers, σ_half) : error("σ_half = $σ_half cannot be used for $nlayers-level SigmaCoordinates") end -# obtain nlev from length of predefined σ_half levels -SigmaCoordinates(σ_half::AbstractVector) = SigmaCoordinates(nlev=length(σ_half)-1; σ_half) +# obtain nlayers from length of predefined σ_half levels +SigmaCoordinates(σ_half::AbstractVector) = SigmaCoordinates(nlayers=length(σ_half)-1; σ_half) SigmaCoordinates(σ_half::AbstractRange) = SigmaCoordinates(collect(σ_half)) function Base.show(io::IO, σ::SigmaCoordinates) - println(io, "$(σ.nlev)-level SigmaCoordinates") - nchars = length(string(σ.nlev)) + println(io, "$(σ.nlayers)-layer SigmaCoordinates") + nchars = length(string(σ.nlayers)) format = Printf.Format("%$(nchars)d") - for k=1:σ.nlev + for k=1:σ.nlayers println(io, "├─ ", @sprintf("%1.4f", σ.σ_half[k]), " k = ", Printf.format(format, k-1), ".5") σk = (σ.σ_half[k] + σ.σ_half[k+1])/2 println(io, "│× ", @sprintf("%1.4f", σk), " k = ", Printf.format(format, k)) end - print(io, "└─ ", @sprintf("%1.4f", σ.σ_half[end]), " k = ", Printf.format(format, σ.nlev), ".5") + print(io, "└─ ", @sprintf("%1.4f", σ.σ_half[end]), " k = ", Printf.format(format, σ.nlayers), ".5") end """ $(TYPEDSIGNATURES) -Vertical sigma coordinates defined by their nlev+1 half levels `σ_levels_half`. Sigma coordinates are +Vertical sigma coordinates defined by their nlayers+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. Default levels are equally spaced from 0 to 1 (including).""" -function default_sigma_coordinates(nlev::Integer) +function default_sigma_coordinates(nlayers::Integer) # equi-distant in σ - σ_half = collect(range(0, 1, nlev+1)) + σ_half = collect(range(0, 1, nlayers+1)) # Frierson 2006 spacing, higher resolution in surface boundary layer and in stratosphere - # z = collect(range(1, 0, nlev+1)) + # z = collect(range(1, 0, nlayers+1)) # σ_half = @. exp(-5*(0.05*z + 0.95*z^3)) # σ_half[1] = 0 return σ_half @@ -48,11 +48,11 @@ end """ $(TYPEDSIGNATURES) -Check that nlev and σ_half match.""" -function sigma_okay(nlev::Integer, σ_half::AbstractVector) +Check that nlayers and σ_half match.""" +function sigma_okay(nlayers::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 nlayers == (length(σ_half) - 1) "nlayers has to be length of σ_half - 1" @assert isincreasing(σ_half) "Vertical sigma coordinates are not increasing." return true end @@ -62,6 +62,6 @@ 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 +default_nlayers(::Type{<:Barotropic}) = 1 +default_nlayers(::Type{<:ShallowWater}) = 1 +default_nlayers(::Type{<:PrimitiveEquation}) = 8 diff --git a/src/dynamics/virtual_temperature.jl b/src/dynamics/virtual_temperature.jl index a99c71410..6daab5b81 100644 --- a/src/dynamics/virtual_temperature.jl +++ b/src/dynamics/virtual_temperature.jl @@ -1,12 +1,4 @@ -# function barrier -function virtual_temperature!( diagn::DiagnosticVariablesLayer, - temp::LowerTriangularMatrix, # only needed for dispatch compat with DryCore - model::PrimitiveWet) - virtual_temperature!(diagn, temp, model.atmosphere) -end - -""" -$(TYPEDSIGNATURES) +"""$(TYPEDSIGNATURES) Calculates the virtual temperature Tᵥ as Tᵥ = T(1+μq) @@ -17,37 +9,30 @@ With absolute temperature T, specific humidity q and in grid-point space.""" function virtual_temperature!( - diagn::DiagnosticVariablesLayer, - temp::LowerTriangularMatrix, # only needed for dispatch compat with DryCore - atmosphere::AbstractAtmosphere, + diagn::DiagnosticVariables, + model::PrimitiveWet, ) - (; temp_grid, humid_grid, temp_virt_grid) = diagn.grid_variables - (; temp_virt) = diagn.dynamics_variables - μ = atmosphere.μ_virt_temp + (; temp_grid, humid_grid, temp_virt_grid) = diagn.grid + μ = model.atmosphere.μ_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]) + @inbounds for ijk in eachindex(temp_virt_grid, temp_grid, humid_grid) + temp_virt_grid[ijk] = temp_grid[ijk]*(1 + μ*humid_grid[ijk]) 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 """ $(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, - model::PrimitiveDry) - - (; temp_grid, temp_virt_grid) = diagn.grid_variables - (; temp_virt) = diagn.dynamics_variables - - copyto!(temp_virt_grid, temp_grid) +function virtual_temperature!( + diagn::DiagnosticVariables, + model::PrimitiveDry, +) + (; temp_grid, temp_virt_grid) = diagn.grid + # Tᵥ = T(1 + μ*q) with humid=q=0 + temp_virt_grid .= temp_grid + return nothing end function virtual_temperature!( @@ -78,24 +63,14 @@ as humidity is zero in this `model`.""" function linear_virtual_temperature!( diagn::DiagnosticVariablesLayer, progn::PrognosticLayerTimesteps, - model::PrimitiveDry, lf::Integer, + model::PrimitiveDry, ) (; 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.atmosphere, lf) -end - """ $(TYPEDSIGNATURES) Calculates a linearised virtual temperature Tᵥ as @@ -108,15 +83,28 @@ specific humidity q and μ = (1-ξ)/ξ, ξ = R_dry/R_vapour. in spectral space.""" -function linear_virtual_temperature!( diagn::DiagnosticVariablesLayer, - progn::PrognosticLayerTimesteps, - atmosphere::AbstractAtmosphere, - lf::Int) - - (; temp_virt) = diagn.dynamics_variables - μ = atmosphere.μ_virt_temp - Tₖ = diagn.temp_average[] - (; temp, humid) = progn.timesteps[lf] +function linear_virtual_temperature!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, + model::PrimitiveEquation, +) + (; temp_virt) = diagn.dynamics + μ = model.atmosphere.μ_virt_temp + (; temp_average) = diagn + temp = progn.temp[lf] + humid = progn.humid[lf] - @. temp_virt = temp + (Tₖ*μ)*humid + # 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 + # transform!(temp_virt, temp_virt_grid, S) + + for k in eachmatrix(temp_virt, temp, humid) + Tₖ = temp_average[k] + for lm in eachharmonic(temp_virt, temp, humid) + temp_virt[lm, k] = temp[lm, k] + (Tₖ*μ)*humid[lm, k] + end + end end \ No newline at end of file diff --git a/src/gpu.jl b/src/gpu.jl index fbf890eac..e36362150 100644 --- a/src/gpu.jl +++ b/src/gpu.jl @@ -1,98 +1,91 @@ """ abstract type AbstractDevice -Supertype of all devices SpeedyWeather.jl can ran on +Supertype of all devices SpeedyWeather.jl can run on """ abstract type AbstractDevice end +export CPU, GPU + """ - CPUDevice <: AbstractDevice + CPU <: AbstractDevice Indicates that SpeedyWeather.jl runs on a single CPU """ -struct CPUDevice <: AbstractDevice end +struct CPU <: AbstractDevice end """ - GPUDevice <: AbstractDevice + GPU <: AbstractDevice Indicates that SpeedyWeather.jl runs on a single GPU """ -struct GPUDevice <: AbstractDevice end +struct GPU <: AbstractDevice end +"""$(TYPEDSIGNATURES) +Return default used device for internal purposes, either `CPU` or `GPU` if a GPU is available. """ - Device() +Device() = CUDA.functional() ? GPU() : CPU() -Return default used device for internal purposes, either `CPUDevice` or `GPUDevice` if a GPU is available. -""" -Device() = CUDA.functional() ? GPUDevice() : CPUDevice() - -""" - Device_KernelAbstractions() +"""$(TYPEDSIGNATURES) +Default array type on `device`.""" +default_array_type(device::AbstractDevice) = default_array_type(typeof(device)) +default_array_type(::Type{GPU}) = CuArray +default_array_type(::Type{CPU}) = Array +"""$(TYPEDSIGNATURES) Return default used device for KernelAbstractions, either `CPU` or `CUDADevice` if a GPU is available """ Device_KernelAbstractions() = CUDA.functional() ? KernelAbstractions.CUDADevice : KernelAbstractions.CPU -""" - Device_KernelAbstractions(::AbstractDevice) - +"""$(TYPEDSIGNATURES) Return used device for use with KernelAbstractions """ -Device_KernelAbstractions(::CPUDevice) = KernelAbstractions.CPU -Device_KernelAbstractions(::GPUDevice) = KernelAbstractions.CUDADevice - -""" - DeviceSetup{S<:AbstractDevice} +Device_KernelAbstractions(::CPU) = KernelAbstractions.CPU +Device_KernelAbstractions(::GPU) = KernelAbstractions.CUDADevice +"""$(TYPEDSIGNATURES) Holds information about the device the model is running on and workgroup size. - -* `device::AbstractDevice`: Device the model is running on -* `device_KA::KernelAbstractions.Device`: Device for use with KernelAbstractions -* `n`: workgroup size -""" +$(TYPEDFIELDS)""" struct DeviceSetup{S<:AbstractDevice, T} - device::S # for internal purposes - device_KA::T # for KernelAbstractions - n::Int # workgroup size + "::AbstractDevice, device the model is running on." + device::S + + "::KernelAbstractions.Device, device for use with KernelAbstractions" + device_KA::T + + "workgroup size" + n::Int end DeviceSetup() = DeviceSetup(Device(), Device_KernelAbstractions(Device()), workgroup_size(Device())) DeviceSetup(device::AbstractDevice) = DeviceSetup(device, Device_KernelAbstractions(device), workgroup_size(device)) DeviceSetup(device::AbstractDevice, n::Integer) = DeviceSetup(device, Device_KernelAbstractions(device), n) -""" - workgroup_size(dev::AbstractDevice) - -Returns a workgroup size depending on `dev`. +"""$(TYPEDSIGNATURES) +Returns a workgroup size depending on `device`. WIP: Will be expanded in the future to also include grid information. """ function workgroup_size(device::AbstractDevice) - return device isa GPUDevice ? 32 : 4 + return device isa GPU ? 32 : 4 end -""" - DeviceArray(device::AbstractDevice, x) - -Adapts `x` to a `CuArray` when `device<:GPUDevice` is used, otherwise a regular `Array`. Uses `adapt`, thus also can return SubArrays etc. -""" -DeviceArray(::GPUDevice, x) = Adapt.adapt(CuArray, x) -DeviceArray(::CPUDevice, x) = Adapt.adapt(Array, x) +"""$(TYPEDSIGNATURES) +Adapts `x` to a `CuArray` when `device::GPU` is used, otherwise a regular `Array`. +Uses `adapt`, thus also can return SubArrays etc.""" +DeviceArray(::GPU, x) = Adapt.adapt(CuArray, x) +DeviceArray(::CPU, x) = Adapt.adapt(Array, x) DeviceArray(dev::DeviceSetup, x) = DeviceArray(dev.device, x) -""" - DeviceArrayNotAdapt(device::AbstractDevice, x) - -Returns a `CuArray` when `device<:GPUDevice` is used, otherwise a regular `Array`. Doesn't uses `adapt`, therefore always returns CuArray/Array -""" -DeviceArrayNotAdapt(::GPUDevice, x) = CuArray(x) -DeviceArrayNotAdapt(::CPUDevice, x) = Array(x) +"""$(TYPEDSIGNATURES) +Returns a `CuArray` when `device<:GPU` is used, otherwise a regular `Array`. +Doesn't uses `adapt`, therefore always returns CuArray/Array.""" +DeviceArrayNotAdapt(::GPU, x) = CuArray(x) +DeviceArrayNotAdapt(::CPU, x) = Array(x) DeviceArrayNotAdapt(dev::DeviceSetup, x) = DeviceArrayNotAdapt(dev.device, x) -""" - launch_kernel!(device_setup::DeviceSetup, kernel!, ndrange, kernel_args...) - -Launches the `kernel!` on the `device_setup` with `ndrange` computations over the kernel and arguments `kernel_args` -""" +"""$(TYPEDSIGNATURES) +Launches the `kernel!` on the `device_setup` with `ndrange` computations over the +kernel and arguments `kernel_args`.""" function launch_kernel!(device_setup::DeviceSetup, kernel!, ndrange, kernel_args...) device = device_setup.device_KA() n = device_setup.n diff --git a/src/models/abstract_models.jl b/src/models/abstract_models.jl index 2b52b9524..36a786f91 100644 --- a/src/models/abstract_models.jl +++ b/src/models/abstract_models.jl @@ -1,10 +1,10 @@ -export Barotropic, ShallowWater, PrimitiveEquation, PrimitiveDry, PrimitiveWet, ModelSetup +export Barotropic, ShallowWater, PrimitiveEquation, PrimitiveDry, PrimitiveWet, AbstractModel abstract type AbstractSimulation{Model} end -abstract type ModelSetup end -abstract type Barotropic <: ModelSetup end -abstract type ShallowWater <: ModelSetup end -abstract type PrimitiveEquation <: ModelSetup end +abstract type AbstractModel end +abstract type Barotropic <: AbstractModel end +abstract type ShallowWater <: AbstractModel end +abstract type PrimitiveEquation <: AbstractModel end abstract type PrimitiveDry <: PrimitiveEquation end abstract type PrimitiveWet <: PrimitiveEquation end @@ -20,15 +20,16 @@ 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(M::Type{<:ModelSetup}, var_name::Symbol) = var_name in (:vor, :div, :temp, :humid, :pres) -has(M::ModelSetup, var_name) = has(typeof(M), var_name) +has(Model::Type{<:AbstractModel}, var_name::Symbol) = var_name in prognostic_variables(Model) +has(model::AbstractModel, var_name) = has(typeof(model), var_name) +prognostic_variables(model::AbstractModel) = prognostic_variables(typeof(model)) # model class is the abstract supertype model_class(::Type{<:Barotropic}) = Barotropic model_class(::Type{<:ShallowWater}) = ShallowWater model_class(::Type{<:PrimitiveDry}) = PrimitiveDry model_class(::Type{<:PrimitiveWet}) = PrimitiveWet -model_class(model::ModelSetup) = model_class(typeof(model)) +model_class(model::AbstractModel) = model_class(typeof(model)) # model type is the parameter-free type of a model # TODO what happens if we have several concrete types under each abstract type? @@ -36,9 +37,9 @@ model_type(::Type{<:Barotropic}) = BarotropicModel model_type(::Type{<:ShallowWater}) = ShallowWaterModel model_type(::Type{<:PrimitiveDry}) = PrimitiveDryModel model_type(::Type{<:PrimitiveWet}) = PrimitiveWetModel -model_type(model::ModelSetup) = model_type(typeof(model)) +model_type(model::AbstractModel) = model_type(typeof(model)) -function Base.show(io::IO, M::ModelSetup) +function Base.show(io::IO, M::AbstractModel) println(io, "$(model_type(M)) <: $(model_class(M))") properties = propertynames(M) n = length(properties) diff --git a/src/models/barotropic.jl b/src/models/barotropic.jl index e6f63ee45..68a43ed91 100644 --- a/src/models/barotropic.jl +++ b/src/models/barotropic.jl @@ -10,7 +10,7 @@ with `spectral_grid::SpectralGrid` used to initalize all non-default components passed on as keyword arguments, e.g. `planet=Earth(spectral_grid)`. Fields, representing model components, are $(TYPEDFIELDS)""" -Base.@kwdef mutable struct BarotropicModel{ +@kwdef mutable struct BarotropicModel{ # TODO add constraints again when we stop supporting julia v1.9 DS, # <:DeviceSetup, GE, # <:AbstractGeometry, @@ -25,12 +25,12 @@ Base.@kwdef mutable struct BarotropicModel{ ST, # <:SpectralTransform{NF}, IM, # <:AbstractImplicit, HD, # <:AbstractHorizontalDiffusion, - OW, # <:AbstractOutputWriter, + OU, # <:AbstractOutput, FB, # <:AbstractFeedback, } <: Barotropic spectral_grid::SpectralGrid - device_setup::DS = DeviceSetup(CPUDevice()) + device_setup::DS = DeviceSetup(spectral_grid.device) # DYNAMICS geometry::GE = Geometry(spectral_grid) @@ -49,12 +49,12 @@ Base.@kwdef mutable struct BarotropicModel{ horizontal_diffusion::HD = HyperDiffusion(spectral_grid) # OUTPUT - output::OW = OutputWriter(spectral_grid, Barotropic) + output::OU = NetCDFOutput(spectral_grid, Barotropic) callbacks::Dict{Symbol, AbstractCallback} = Dict{Symbol, AbstractCallback}() feedback::FB = Feedback() end -has(::Type{<:Barotropic}, var_name::Symbol) = var_name in (:vor,) +prognostic_variables(::Type{<:Barotropic}) = (:vor,) default_concrete_model(::Type{Barotropic}) = BarotropicModel """ @@ -65,8 +65,8 @@ at in `time_stepping!`.""" function initialize!(model::Barotropic; time::DateTime = DEFAULT_DATE) (; spectral_grid) = model - spectral_grid.nlev > 1 && @warn "Only nlev=1 supported for BarotropicModel, \ - SpectralGrid with nlev=$(spectral_grid.nlev) provided." + spectral_grid.nlayers > 1 && @error "Only nlayers=1 supported for BarotropicModel, \ + SpectralGrid with nlayers=$(spectral_grid.nlayers) provided." # initialize components initialize!(model.time_stepping, model) @@ -85,6 +85,6 @@ function initialize!(model::Barotropic; time::DateTime = DEFAULT_DATE) initialize!(model.particle_advection, model) initialize!(prognostic_variables.particles, model) - diagnostic_variables = DiagnosticVariables(spectral_grid, model) + diagnostic_variables = DiagnosticVariables(spectral_grid) return Simulation(prognostic_variables, diagnostic_variables, model) end \ No newline at end of file diff --git a/src/models/primitive_dry.jl b/src/models/primitive_dry.jl index b9ebece98..520ed9813 100644 --- a/src/models/primitive_dry.jl +++ b/src/models/primitive_dry.jl @@ -10,7 +10,7 @@ with `spectral_grid::SpectralGrid` used to initalize all non-default components passed on as keyword arguments, e.g. `planet=Earth(spectral_grid)`. Fields, representing model components, are $(TYPEDFIELDS)""" -Base.@kwdef mutable struct PrimitiveDryModel{ +@kwdef mutable struct PrimitiveDryModel{ # TODO add constraints again when we stop supporting julia v1.9 DS, # <:DeviceSetup, GE, # <:AbstractGeometry, @@ -41,12 +41,12 @@ Base.@kwdef mutable struct PrimitiveDryModel{ IM, # <:AbstractImplicit, HD, # <:AbstractHorizontalDiffusion, VA, # <:AbstractVerticalAdvection, - OW, # <:AbstractOutputWriter, + OU, # <:AbstractOutput, FB, # <:AbstractFeedback, } <: PrimitiveDry spectral_grid::SpectralGrid - device_setup::DS = DeviceSetup(CPUDevice()) + device_setup::DS = DeviceSetup(spectral_grid.device) # DYNAMICS dynamics::Bool = true @@ -87,12 +87,12 @@ Base.@kwdef mutable struct PrimitiveDryModel{ vertical_advection::VA = CenteredVerticalAdvection(spectral_grid) # OUTPUT - output::OW = OutputWriter(spectral_grid, PrimitiveDry) + output::OU = NetCDFOutput(spectral_grid, PrimitiveDry) callbacks::Dict{Symbol, AbstractCallback} = Dict{Symbol, AbstractCallback}() feedback::FB = Feedback() end -has(::Type{<:PrimitiveDry}, var_name::Symbol) = var_name in (:vor, :div, :temp, :pres) +prognostic_variables(::Type{<:PrimitiveDry}) = (:vor, :div, :temp, :pres) default_concrete_model(::Type{PrimitiveDry}) = PrimitiveDryModel """ @@ -136,15 +136,16 @@ function initialize!(model::PrimitiveDry; time::DateTime = DEFAULT_DATE) (; clock) = prognostic_variables clock.time = time # set the current time clock.start = time # and store the start time - + + diagnostic_variables = DiagnosticVariables(spectral_grid) + # particle advection initialize!(model.particle_advection, model) initialize!(prognostic_variables.particles, model) # initialize ocean and land and synchronize clocks - initialize!(prognostic_variables.ocean, clock.time, model) - initialize!(prognostic_variables.land, clock.time, model) + initialize!(prognostic_variables.ocean, time, model) + initialize!(prognostic_variables.land, prognostic_variables, diagnostic_variables, model) - diagnostic_variables = DiagnosticVariables(spectral_grid, model) return Simulation(prognostic_variables, diagnostic_variables, model) end \ No newline at end of file diff --git a/src/models/primitive_wet.jl b/src/models/primitive_wet.jl index 010028419..da033ca81 100644 --- a/src/models/primitive_wet.jl +++ b/src/models/primitive_wet.jl @@ -10,7 +10,7 @@ with `spectral_grid::SpectralGrid` used to initalize all non-default components passed on as keyword arguments, e.g. `planet=Earth(spectral_grid)`. Fields, representing model components, are $(TYPEDFIELDS)""" -Base.@kwdef mutable struct PrimitiveWetModel{ +@kwdef mutable struct PrimitiveWetModel{ # TODO add constraints again when we stop supporting julia v1.9 DS, # <:DeviceSetup, GE, # <:AbstractGeometry, @@ -47,12 +47,12 @@ Base.@kwdef mutable struct PrimitiveWetModel{ HD, # <:AbstractHorizontalDiffusion, VA, # <:AbstractVerticalAdvection, HF, # <:AbstractHoleFilling, - OW, # <:AbstractOutputWriter, + OU, # <:AbstractOutput, FB, # <:AbstractFeedback, } <: PrimitiveWet spectral_grid::SpectralGrid - device_setup::DS = DeviceSetup(CPUDevice()) + device_setup::DS = DeviceSetup(spectral_grid.device) # DYNAMICS dynamics::Bool = true @@ -99,12 +99,12 @@ Base.@kwdef mutable struct PrimitiveWetModel{ hole_filling::HF = ClipNegatives(spectral_grid) # OUTPUT - output::OW = OutputWriter(spectral_grid, PrimitiveWet) + output::OU = NetCDFOutput(spectral_grid, PrimitiveWet) callbacks::Dict{Symbol, AbstractCallback} = Dict{Symbol, AbstractCallback}() feedback::FB = Feedback() end - -has(::Type{<:PrimitiveWet}, var_name::Symbol) = var_name in (:vor, :div, :temp, :pres, :humid) + +prognostic_variables(::Type{<:PrimitiveWet}) = (:vor, :div, :temp, :humid, :pres) default_concrete_model(::Type{PrimitiveWet}) = PrimitiveWetModel """ @@ -154,14 +154,15 @@ function initialize!(model::PrimitiveWet; time::DateTime = DEFAULT_DATE) clock.time = time # set the current time clock.start = time # and store the start time + diagnostic_variables = DiagnosticVariables(spectral_grid) + # particle advection initialize!(model.particle_advection, model) initialize!(prognostic_variables.particles, model) # initialize ocean and land and synchronize clocks - initialize!(prognostic_variables.ocean, clock.time, model) - initialize!(prognostic_variables.land, clock.time, model) + initialize!(prognostic_variables.ocean, time, model) + initialize!(prognostic_variables.land, prognostic_variables, diagnostic_variables, model) - diagnostic_variables = DiagnosticVariables(spectral_grid, model) return Simulation(prognostic_variables, diagnostic_variables, model) end \ No newline at end of file diff --git a/src/models/shallow_water.jl b/src/models/shallow_water.jl index dbebc1c40..66fe8452f 100644 --- a/src/models/shallow_water.jl +++ b/src/models/shallow_water.jl @@ -10,7 +10,7 @@ with `spectral_grid::SpectralGrid` used to initalize all non-default components passed on as keyword arguments, e.g. `planet=Earth(spectral_grid)`. Fields, representing model components, are $(TYPEDFIELDS)""" -Base.@kwdef mutable struct ShallowWaterModel{ +@kwdef mutable struct ShallowWaterModel{ # TODO add constraints again when we stop supporting julia v1.9 DS, # <:DeviceSetup, GE, # <:AbstractGeometry, @@ -26,12 +26,12 @@ Base.@kwdef mutable struct ShallowWaterModel{ ST, # <:SpectralTransform{NF}, IM, # <:AbstractImplicit, HD, # <:AbstractHorizontalDiffusion, - OW, # <:AbstractOutputWriter, + OU, # <:AbstractOutput, FB, # <:AbstractFeedback, } <: ShallowWater spectral_grid::SpectralGrid - device_setup::DS = DeviceSetup(CPUDevice()) + device_setup::DS = DeviceSetup(spectral_grid.device) # DYNAMICS geometry::GE = Geometry(spectral_grid) @@ -51,12 +51,12 @@ Base.@kwdef mutable struct ShallowWaterModel{ horizontal_diffusion::HD = HyperDiffusion(spectral_grid) # OUTPUT - output::OW = OutputWriter(spectral_grid, ShallowWater) + output::OU = NetCDFOutput(spectral_grid, ShallowWater) callbacks::Dict{Symbol, AbstractCallback} = Dict{Symbol, AbstractCallback}() feedback::FB = Feedback() end -has(::Type{<:ShallowWater}, var_name::Symbol) = var_name in (:vor, :div, :pres) +prognostic_variables(::Type{<:ShallowWater}) = (:vor, :div, :pres) default_concrete_model(::Type{ShallowWater}) = ShallowWaterModel """ @@ -67,8 +67,8 @@ in `time_stepping!` and `model.implicit` which is done in `first_timesteps!`.""" function initialize!(model::ShallowWater; time::DateTime = DEFAULT_DATE) (; spectral_grid) = model - spectral_grid.nlev > 1 && @warn "Only nlev=1 supported for ShallowWaterModel, \ - SpectralGrid with nlev=$(spectral_grid.nlev) provided." + spectral_grid.nlayers > 1 && @error "Only nlayers=1 supported for ShallowWaterModel, \ + SpectralGrid with nlayers=$(spectral_grid.nlayers) provided." # initialize components initialize!(model.time_stepping, model) @@ -89,6 +89,6 @@ function initialize!(model::ShallowWater; time::DateTime = DEFAULT_DATE) initialize!(model.particle_advection, model) initialize!(prognostic_variables.particles, model) - diagnostic_variables = DiagnosticVariables(spectral_grid, model) + diagnostic_variables = DiagnosticVariables(spectral_grid) return Simulation(prognostic_variables, diagnostic_variables, model) end \ No newline at end of file diff --git a/src/models/simulation.jl b/src/models/simulation.jl index b763a783f..449bbdcc0 100644 --- a/src/models/simulation.jl +++ b/src/models/simulation.jl @@ -5,7 +5,7 @@ $(TYPEDSIGNATURES) Simulation is a container struct to be used with `run!(::Simulation)`. It contains $(TYPEDFIELDS)""" -struct Simulation{Model<:ModelSetup} <: AbstractSimulation{Model} +struct Simulation{Model<:AbstractModel} <: AbstractSimulation{Model} "define the current state of the model" prognostic_variables::PrognosticVariables @@ -18,9 +18,9 @@ end function Base.show(io::IO, S::AbstractSimulation) println(io, "Simulation{$(model_type(S.model))}") - println(io, "├ $(typeof(S.prognostic_variables))") - println(io, "├ $(typeof(S.diagnostic_variables))") - print(io, "└ model::$(model_type(S.model))") + println(io, "├ prognostic_variables::PrognosticVariables{...}") + println(io, "├ diagnostic_variables::DiagnosticVariables{...}") + print(io, "└ model::$(model_type(S.model)){...}") end export run! @@ -29,7 +29,7 @@ export run! $(TYPEDSIGNATURES) Run a SpeedyWeather.jl `simulation`. The `simulation.model` is assumed to be initialized.""" function run!( simulation::AbstractSimulation; - period = Day(10), + period::Dates.Period = Day(10), output::Bool = false, n_days::Union{Nothing, Real} = nothing) @@ -46,7 +46,7 @@ function run!( simulation::AbstractSimulation; initialize!(clock, model.time_stepping) # store the start date, reset counter # OUTPUT - model.output.output = output # enable/disable output + model.output.active = output # enable/disable output # run it, yeah! time_stepping!(prognostic_variables, diagnostic_variables, model) diff --git a/src/models/tree.jl b/src/models/tree.jl index 6080494c8..8b925fe45 100644 --- a/src/models/tree.jl +++ b/src/models/tree.jl @@ -25,7 +25,7 @@ as long as they are defined within the modules argument (default SpeedyWeather). Other keyword arguments are `max_level::Integer=10`, `with_types::Bool=false` or `with_size::Bool=false.""" function tree( - M::ModelSetup; + M::AbstractModel; modules=SpeedyWeather, with_size::Bool = false, kwargs... diff --git a/src/output/abstract_types.jl b/src/output/abstract_types.jl deleted file mode 100644 index 9487db7d0..000000000 --- a/src/output/abstract_types.jl +++ /dev/null @@ -1,2 +0,0 @@ -abstract type AbstractOutputWriter end -abstract type AbstractFeedback end \ No newline at end of file diff --git a/src/output/callbacks.jl b/src/output/callbacks.jl index 381d4b4fb..d668a5c87 100644 --- a/src/output/callbacks.jl +++ b/src/output/callbacks.jl @@ -61,15 +61,15 @@ end """ $(TYPEDSIGNATURES) -Add a or several callbacks to a model::ModelSetup. To be used like +Add a or several callbacks to a model::AbstractModel. To be used like add!(model, :my_callback => callback) add!(model, :my_callback1 => callback, :my_callback2 => other_callback) """ -add!(model::ModelSetup, key_callbacks::Pair{Symbol, <:AbstractCallback}...) = +add!(model::AbstractModel, key_callbacks::Pair{Symbol, <:AbstractCallback}...) = add!(model.callbacks, key_callbacks...) add!(D::CALLBACK_DICT, key::Symbol, callback::AbstractCallback) = add!(D, Pair(key, callback)) -add!(model::ModelSetup, key::Symbol, callback::AbstractCallback) = +add!(model::AbstractModel, key::Symbol, callback::AbstractCallback) = add!(model.callbacks, Pair(key, callback)) @@ -102,7 +102,7 @@ key which is randomly created like callback_????. To be used like add!(model.callbacks, callback) add!(model.callbacks, callback1, callback2).""" -add!(model::ModelSetup, callbacks::AbstractCallback...) = +add!(model::AbstractModel, callbacks::AbstractCallback...) = add!(model.callbacks, callbacks..., verbose = model.feedback.verbose) # delete!(dict, key) already defined in Base @@ -128,11 +128,11 @@ function initialize!( callback::GlobalSurfaceTemperatureCallback{NF}, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) where NF - callback.temp = Vector{NF}(undef, progn.clock.n_timesteps+1) # replace with vector of correct length - callback.temp[1] = diagn.layers[diagn.nlev].temp_average[] # set initial conditions - callback.timestep_counter = 1 # (re)set counter to 1 + callback.temp = Vector{NF}(undef, progn.clock.n_timesteps+1) # replace with vector of correct length + callback.temp[1] = diagn.temp_average[diagn.nlayers] # set initial conditions + callback.timestep_counter = 1 # (re)set counter to 1 end """ @@ -143,11 +143,11 @@ function callback!( callback::GlobalSurfaceTemperatureCallback, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) callback.timestep_counter += 1 i = callback.timestep_counter - callback.temp[i] = diagn.layers[diagn.nlev].temp_average[] + callback.temp[i] = diagn.temp_average[diagn.nlayers] end # nothing to finish diff --git a/src/output/feedback.jl b/src/output/feedback.jl index c7def8213..77d4b6d61 100644 --- a/src/output/feedback.jl +++ b/src/output/feedback.jl @@ -1,3 +1,5 @@ +abstract type AbstractFeedback end + export Feedback """ @@ -70,7 +72,7 @@ end """ $(TYPEDSIGNATURES) Initializes the a `Feedback` struct.""" -function initialize!(feedback::Feedback, clock::Clock, model::ModelSetup) +function initialize!(feedback::Feedback, clock::Clock, model::AbstractModel) # set to false to recheck for NaRs feedback.nars_detected = false @@ -158,10 +160,10 @@ function nar_detection!(feedback::Feedback, progn::PrognosticVariables) feedback.nars_detected && return nothing # escape immediately if nans already detected i = feedback.progress_meter.counter # time step - (; vor ) = progn.layers[end].timesteps[2] # only check for surface vorticity + vor0 = progn.vor[2][1, end] # only check 0-0 mode of surface vorticity # just check first harmonic, spectral transform propagates NaRs globally anyway - nars_detected_here = ~isfinite(vor[1]) + nars_detected_here = ~isfinite(vor0) nars_detected_here && @warn "NaN or Inf detected at time step $i" feedback.nars_detected = nars_detected_here end diff --git a/src/output/netcdf_output.jl b/src/output/netcdf_output.jl new file mode 100644 index 000000000..c6e9a0922 --- /dev/null +++ b/src/output/netcdf_output.jl @@ -0,0 +1,866 @@ +abstract type AbstractOutput <: AbstractModelComponent end +abstract type AbstractOutputVariable end + +# default number format for output +const DEFAULT_OUTPUT_NF = Float32 +const DEFAULT_OUTPUT_DT = Hour(6) +const OUTPUT_VARIABLES_DICT = Dict{Symbol, AbstractOutputVariable} +OutputVariablesDict() = OUTPUT_VARIABLES_DICT() + +const DEFAULT_MISSING_VALUE = NaN +const DEFAULT_COMPRESSION_LEVEL = 1 +const DEFAULT_SHUFFLE = false +const DEFAULT_KEEPBITS = 15 + +export NetCDFOutput + +"""Output writer for a netCDF file with (re-)gridded variables. +Interpolates non-rectangular grids. Fields are +$(TYPEDFIELDS)""" +@kwdef mutable struct NetCDFOutput{ + Grid2D, + Grid3D, + Interpolator, +} <: AbstractOutput + + # FILE OPTIONS + active::Bool = false + + "[OPTION] path to output folder, run_???? will be created within" + path::String = pwd() + + "[OPTION] run identification number/string" + id::String = "0001" + run_path::String = "" # will be determined in initalize! + + "[OPTION] name of the output netcdf file" + filename::String = "output.nc" + + "[OPTION] also write restart file if output==true?" + write_restart::Bool = true + pkg_version::VersionNumber = pkgversion(SpeedyWeather) + + # WHAT/WHEN OPTIONS + startdate::DateTime = DateTime(2000, 1, 1) + + "[OPTION] output frequency, time step" + output_dt::Second = Second(DEFAULT_OUTPUT_DT) + + "[OPTION] dictionary of variables to output, e.g. u, v, vor, div, pres, temp, humid" + variables::OUTPUT_VARIABLES_DICT = OutputVariablesDict() + + # 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 + + # the netcdf file to be written into, will be created + netcdf_file::Union{NCDataset, Nothing} = nothing + + const interpolator::Interpolator + + # SCRATCH GRIDS TO INTERPOLATE ONTO + const grid2D::Grid2D + const grid3D::Grid3D +end + +""" +$(TYPEDSIGNATURES) +Constructor for NetCDFOutput based on `S::SpectralGrid` and optionally +the `Model` type (e.g. `ShallowWater`, `PrimitiveWet`) as second positional argument. +The output grid is optionally determined by keyword arguments `output_Grid` (its type, full grid required), +`nlat_half` (resolution) and `output_NF` (number format). By default, uses the full grid +equivalent of the grid and resolution used in `SpectralGrid` `S`.""" +function NetCDFOutput( + S::SpectralGrid, + Model::Type{<:AbstractModel} = Barotropic; + output_Grid::Type{<:AbstractFullGrid} = RingGrids.full_grid_type(S.Grid), + nlat_half::Integer = S.nlat_half, + output_NF::DataType = DEFAULT_OUTPUT_NF, + output_dt::Period = Second(DEFAULT_OUTPUT_DT), # only needed for dispatch + kwargs...) + + # INPUT GRID + input_Grid = S.Grid + input_nlat_half = S.nlat_half + + # OUTPUT GRID + nlon = RingGrids.get_nlon(output_Grid, nlat_half) + nlat = RingGrids.get_nlat(output_Grid, nlat_half) + npoints = nlon*nlat + (; nlayers) = S + + # CREATE INTERPOLATOR + interpolator = DEFAULT_INTERPOLATOR(DEFAULT_OUTPUT_NF, input_Grid, input_nlat_half, npoints) + + # CREATE GRIDS TO + output_Grid2D = RingGrids.nonparametric_type(output_Grid){output_NF, 1} + output_Grid3D = RingGrids.nonparametric_type(output_Grid){output_NF, 2} + grid2D = output_Grid2D(undef, nlat_half) + grid3D = output_Grid3D(undef, nlat_half, nlayers) + + output = NetCDFOutput(; + output_dt=Second(output_dt), # convert to seconds for dispatch + interpolator, + grid2D, + grid3D, + kwargs...) + + add_default!(output.variables, Model) + return output +end + +function Base.show(io::IO, output::NetCDFOutput{Grid}) where Grid + println(io, "NetCDFOutput{$Grid}") + println(io, "├ status: $(output.active ? "active" : "inactive/uninitialized")") + println(io, "├ write restart file: $(output.write_restart) (if active)") + println(io, "├ interpolator: $(typeof(output.interpolator))") + println(io, "├ path: $(joinpath(output.run_path, output.filename))") + println(io, "├ frequency: $(output.output_dt)") + print(io, "└┐ variables:") + nvars = length(output.variables) + for (i, (key, var)) in enumerate(output.variables) + print(io, "\n $(i==nvars ? "└" : "├") $key: $(var.long_name) [$(var.unit)]") + end +end + +""" +$(TYPEDSIGNATURES) +Add `outputvariables` to a dictionary defining the variables subject to NetCDF output.""" +function add!(D::OUTPUT_VARIABLES_DICT, outputvariables::AbstractOutputVariable...) + for outputvariable in outputvariables # loop over all variables in arguments + key = Symbol(outputvariable.name) # use name as key::Symbol + D[key] = outputvariable + end + return D +end + +"""$(TYPEDSIGNATURES) +Add `outputvariables` to the dictionary in `output::NetCDFOutput`, i.e. at `output.variables`.""" +function add!(output::NetCDFOutput, outputvariables::AbstractOutputVariable...) + add!(output.variables, outputvariables...) + return output +end + +"""$(TYPEDSIGNATURES) +Add `outputvariables` to the dictionary in `output::NetCDFOutput` of `model`, i.e. at `model.output.variables`.""" +function add!(model::AbstractModel, outputvariables::AbstractOutputVariable...) + add!(model.output, outputvariables...) + return nothing +end + +"""$(TYPEDSIGNATURES) +Delete output variables from `output` by their (short name) (Symbol or String), corresponding +to the keys in the dictionary.""" +function Base.delete!(output::NetCDFOutput, keys::Union{String, Symbol}...) + for key in keys + delete!(output.variables, Symbol(key)) + end + return output +end + +"""$(TYPEDSIGNATURES) +Add default variables to output for a `Barotropic` model: Vorticity, zonal and meridional velocity.""" +function add_default!( + output_variables::OUTPUT_VARIABLES_DICT, + Model::Type{<:Barotropic}, +) + add!(output_variables, VorticityOutput(), ZonalVelocityOutput(), MeridionalVelocityOutput()) +end + +"""$(TYPEDSIGNATURES) +Add default variables to output for a `ShallowWater` model, same as for a `Barotropic` model but also +the interface displacement.""" +function add_default!( + variables::Dict{Symbol, AbstractOutputVariable}, + Model::Type{<:ShallowWater}, +) + add_default!(variables, Barotropic) + add!(variables, InterfaceDisplacementOutput()) +end + +"""$(TYPEDSIGNATURES) +Add default variables to output for a `PrimitiveDry` model, same as for a `Barotropic` model but also +the surface pressure and temperature.""" +function add_default!( + variables::Dict{Symbol, AbstractOutputVariable}, + Model::Type{<:PrimitiveDry}, +) + add_default!(variables, Barotropic) + add!(variables, SurfacePressureOutput(), TemperatureOutput()) +end + +"""$(TYPEDSIGNATURES) +Add default variables to output for a `PrimitiveWet` model, same as for a `PrimitiveDry` model but also +the specific humidity.""" +function add_default!( + variables::Dict{Symbol, AbstractOutputVariable}, + Model::Type{<:PrimitiveWet}, +) + add_default!(variables, PrimitiveDry) + add!(variables, HumidityOutput(), ConvectivePrecipitationOutput(), LargeScalePrecipitationOutput(), CloudTopOutput()) +end + +"""$(TYPEDSIGNATURES) +Initialize NetCDF `output` by creating a netCDF file and storing the initial conditions +of `diagn` (and `progn`). To be called just before the first timesteps.""" +function initialize!( + output::NetCDFOutput{Grid2D, Grid3D, Interpolator}, + feedback::AbstractFeedback, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) where {Grid2D, Grid3D, Interpolator} + + output.active || return nothing # exit immediately for no output + + # GET RUN ID, CREATE FOLDER + # get new id only if not already specified + 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! + + # OUTPUT FREQUENCY + output.output_every_n_steps = max(1, round(Int, + Millisecond(output.output_dt).value/model.time_stepping.Δt_millisec.value)) + output.output_dt = Second(round(Int, output.output_every_n_steps*model.time_stepping.Δt_sec)) + + # RESET COUNTERS + output.timestep_counter = 0 # time step counter + output.output_counter = 0 # output step counter + + # CREATE NETCDF FILE, vector of NcVars for output + (; run_path, filename) = output + dataset = NCDataset(joinpath(run_path, filename), "c") + output.netcdf_file = dataset + + # DEFINE NETCDF DIMENSIONS TIME and write current (=initial) time + (; startdate) = output + time_string = "hours since $(Dates.format(startdate, "yyyy-mm-dd HH:MM:0.0"))" + defDim(dataset, "time", Inf) # unlimited time dimension + defVar(dataset, "time", Float64, ("time",), + attrib=Dict("units"=>time_string, "long_name"=>"time", + "standard_name"=>"time", "calendar"=>"proleptic_gregorian")) + output!(output, progn.clock.time) # write initial time + + # DEFINE NETCDF DIMENSIONS SPACE + Grid = typeof(output.grid2D) + nlat_half = output.grid2D.nlat_half + lond = get_lond(Grid, nlat_half) + latd = get_latd(Grid, nlat_half) + + # INTERPOLATION: PRECOMPUTE LOCATION INDICES + latds, londs = RingGrids.get_latdlonds(Grid, nlat_half) + RingGrids.update_locator!(output.interpolator, latds, londs) + + σ = model.geometry.σ_levels_full + defVar(dataset, "lon", lond, ("lon",), attrib=Dict("units"=>"degrees_east", "long_name"=>"longitude")) + defVar(dataset, "lat", latd, ("lat",), attrib=Dict("units"=>"degrees_north", "long_name"=>"latitude")) + defVar(dataset, "layer", σ, ("layer",), attrib=Dict("units"=>"1", "long_name"=>"sigma layer")) + all_dims = ["lon", "lat", "layer", "time"] + + # VARIABLES, define every output variable in the netCDF file and write initial conditions + output_NF = eltype(output.grid2D) + for (key, var) in output.variables + missing_value = hasfield(typeof(var), :missing_value) ? var.missing_value : DEFAULT_MISSING_VALUE + attributes = Dict("long_name"=>var.long_name, "units"=>var.unit, "_FillValue"=>output_NF(missing_value)) + dims = collect(dim for (dim, this_dim) in zip(all_dims, var.dims_xyzt) if this_dim) + + # pick defaults if compression + deflatelevel = hasfield(typeof(var), :compression_level) ? var.compression_level : DEFAULT_COMPRESSION_LEVEL + shuffle = hasfield(typeof(var), :shuffle) ? var.shuffle : DEFAULT_SHUFFLE + + defVar(dataset, var.name, output_NF, dims, attrib=attributes; deflatelevel, shuffle) + output!(output, var, progn, diagn, model) + end + + # also export parameters into run????/parameters.txt + parameters_txt = open(joinpath(output.run_path, "parameters.txt"), "w") + for property in propertynames(model) + println(parameters_txt, "model.$property") + println(parameters_txt, getfield(model, property,), "\n") + end + close(parameters_txt) +end + +Base.close(output::NetCDFOutput) = NCDatasets.close(output.netcdf_file) +Base.close(::Nothing) = nothing # in case of no netCDF output nothing to close + +""" +$(TYPEDSIGNATURES) +Writes the variables from `progn` or `diagn` of time step `i` at time `time` into `output.netcdf_file`. +Simply escapes for no netcdf output or if output shouldn't be written on this time step. +Interpolates onto output grid and resolution as specified in `output`, converts to output +number format, truncates the mantissa for higher compression and applies lossless compression.""" +function output!( + output::NetCDFOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + output.timestep_counter += 1 # increase counter + (; active, output_every_n_steps, timestep_counter ) = output + active || return nothing # escape immediately for no netcdf output + timestep_counter % output_every_n_steps == 0 || return nothing # escape if output not written on this step + + output!(output, progn.clock.time) # increase counter write time + output!(output, output.variables, progn, diagn, model) # write variables +end + +""" +$(TYPEDSIGNATURES) +Write the current time `time::DateTime` to the netCDF file in `output`.""" +function output!( + output::NetCDFOutput, + time::DateTime, +) + output.output_counter += 1 # output counter increases when writing time + i = output.output_counter + + (; netcdf_file, startdate ) = output + time_passed = Millisecond(time-startdate) + time_hrs = time_passed.value/3600_000 # [ms] -> [hrs] + netcdf_file["time"][i] = time_hrs + NCDatasets.sync(netcdf_file) + + return nothing +end + +"""$(TYPEDSIGNATURES) +Loop over every variable in `output.variables` to call the respective `output!` method +to write into the `output.netcdf_file`.""" +function output!( + output::NetCDFOutput, + output_variables::OUTPUT_VARIABLES_DICT, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + for (key, var) in output_variables + output!(output, var, progn, diagn, model) + end +end + +function Base.show(io::IO, outputvariable::AbstractOutputVariable) + print(io, "$(typeof(outputvariable)) <: SpeedyWeather.AbstractOutputVariable") + for field in propertynames(outputvariable) + value = getfield(outputvariable, field) + print(io, "\n├ $field::$(typeof(value)) = $value") + end +end + +## VORTICITY ------------- + +"""Defines netCDF output of vorticity. Fields are +$(TYPEDFIELDS) + +Custom variable output defined similarly with required fields marked, +optional fields otherwise use variable-independent defaults. Initialize with `VorticityOutput()` +and non-default fields can always be passed on as keyword arguments, +e.g. `VorticityOutput(long_name="relative vorticity", compression_level=0)`.""" +@kwdef mutable struct VorticityOutput <: AbstractOutputVariable + + "[Required] short name of variable (unique) used in netCDF file and key for dictionary" + name::String = "vor" + + "[Required] unit of variable" + unit::String = "s^-1" + + "[Required] long name of variable used in netCDF file" + long_name::String = "relative vorticity" + + "[Required] NetCDF dimensions the variable uses, lon, lat, layer, time" + dims_xyzt::NTuple{4, Bool} = (true, true, true, true) + + "[Optional] missing value for the variable, if not specified uses NaN" + missing_value::Float64 = NaN + + "[Optional] compression level of the lossless compressor, 1=lowest/fastest, 9=highest/slowest, 3=default" + compression_level::Int = 3 + + "[Optional] bitshuffle the data for compression, false = default" + shuffle::Bool = true + + "[Optional] number of mantissa bits to keep for compression (default: 15)" + keepbits::Int = 5 +end + +"""$(TYPEDSIGNATURES) +Output the vorticity field `vor` from `diagn.grid` into the netCDF file `output.netcdf_file`. +Interpolates the vorticity field onto the output grid and resolution as specified in `output`. +Method required for all output variables `<: AbstractOutputVariable` with dispatch over the +second argument.""" +function output!( + output::NetCDFOutput, + variable::VorticityOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + vor = output.grid3D # use output grid, vorticity is 3D + (; vor_grid) = diagn.grid # unpack vorticity from gridded diagnostic variables + RingGrids.interpolate!(vor, vor_grid, output.interpolator) + + unscale!(vor, diagn.scale[]) # was vor*radius, back to vor + round!(vor, variable.keepbits) # bitrounding for compression + i = output.output_counter # write into timestep i + output.netcdf_file[variable.name][:, :, :, i] = vor + return nothing +end + +## U velocity ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct ZonalVelocityOutput <: AbstractOutputVariable + name::String = "u" + unit::String = "m/s" + long_name::String = "zonal wind" + dims_xyzt::NTuple{4, Bool} = (true, true, true, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 7 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `ZonalVelocityOutput` to write the zonal velocity field `u` from `diagn.grid`, +see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::ZonalVelocityOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + u = output.grid3D + (; u_grid) = diagn.grid + RingGrids.interpolate!(u, u_grid, output.interpolator) + + round!(u, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, :, i] = u + return nothing +end + +## V velocity ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct MeridionalVelocityOutput <: AbstractOutputVariable + name::String = "v" + unit::String = "m/s" + long_name::String = "meridional wind" + dims_xyzt::NTuple{4, Bool} = (true, true, true, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 7 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::MeridionalVelocityOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + v = output.grid3D + (; v_grid) = diagn.grid + RingGrids.interpolate!(v, v_grid, output.interpolator) + + round!(v, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, :, i] = v + return nothing +end + +## DIVERGENCE ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct DivergenceOutput <: AbstractOutputVariable + name::String = "div" + unit::String = "s^-1" + long_name::String = "divergence" + dims_xyzt::NTuple{4, Bool} = (true, true, true, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 5 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::DivergenceOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + div = output.grid3D + (; div_grid) = diagn.grid + RingGrids.interpolate!(div, div_grid, output.interpolator) + + unscale!(div, diagn.scale[]) # was vor*radius, back to vor + round!(div, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, :, i] = div + return nothing +end + +## INTERFACE DISPLACEMENT ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct InterfaceDisplacementOutput <: AbstractOutputVariable + name::String = "eta" + unit::String = "m" + long_name::String = "interface displacement" + dims_xyzt::NTuple{4, Bool} = (true, true, false, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 7 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::InterfaceDisplacementOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + eta = output.grid2D + (; pres_grid) = diagn.grid + RingGrids.interpolate!(eta, pres_grid, output.interpolator) + + round!(eta, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, i] = eta + return nothing +end + +## SURFACE PRESSURE ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct SurfacePressureOutput <: AbstractOutputVariable + name::String = "pres" + unit::String = "hPa" + long_name::String = "surface pressure" + dims_xyzt::NTuple{4, Bool} = (true, true, false, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 12 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::SurfacePressureOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + pres = output.grid2D + (; pres_grid) = diagn.grid + RingGrids.interpolate!(pres, pres_grid, output.interpolator) + + @inbounds for ij in eachindex(pres) + pres[ij] = exp(pres[ij]) / 100 # from log(Pa) to hPa + end + + round!(pres, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, i] = pres + return nothing +end + +## TEMPERATURE ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct TemperatureOutput <: AbstractOutputVariable + name::String = "temp" + unit::String = "degC" + long_name::String = "temperature" + dims_xyzt::NTuple{4, Bool} = (true, true, true, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 10 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::TemperatureOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + temp = output.grid3D + (; temp_grid) = diagn.grid + RingGrids.interpolate!(temp, temp_grid, output.interpolator) + temp .-= 273.15 # convert from K to ˚C + + round!(temp, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, :, i] = temp + return nothing +end + +## HUMIDITY ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct HumidityOutput <: AbstractOutputVariable + name::String = "humid" + unit::String = "kg/kg" + long_name::String = "specific humidity" + dims_xyzt::NTuple{4, Bool} = (true, true, true, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 7 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::HumidityOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + humid = output.grid3D + (; humid_grid) = diagn.grid + RingGrids.interpolate!(humid, humid_grid, output.interpolator) + + round!(humid, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, :, i] = humid + return nothing +end + +## OROGRAPHY ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct OrographyOutput <: AbstractOutputVariable + name::String = "orography" + unit::String = "m" + long_name::String = "orography" + dims_xyzt::NTuple{4, Bool} = (true, true, false, false) + missing_value::Float64 = NaN + compression_level::Int = DEFAULT_COMPRESSION_LEVEL + shuffle::Bool = DEFAULT_SHUFFLE + keepbits::Int = 15 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::OrographyOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + # escape immediately when initialization counter > 1 to not write orography again + output.output_counter > 1 || return nothing + + orog = output.grid2D + (; orography) = model.orography + RingGrids.interpolate!(orog, orography, output.interpolator) + + round!(orog, variable.keepbits) + output.netcdf_file[variable.name][:, :] = orog + return nothing +end + +## PRECIPITATION ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct ConvectivePrecipitationOutput <: AbstractOutputVariable + name::String = "precip_conv" + unit::String = "mm/hr" + long_name::String = "convective precipitation" + dims_xyzt::NTuple{4, Bool} = (true, true, false, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 7 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::ConvectivePrecipitationOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + precip = output.grid2D + (; precip_convection) = diagn.physics + RingGrids.interpolate!(precip, precip_convection, output.interpolator) + + # after output set precip accumulator back to zero + precip_convection .= 0 + + # convert from [m] to [mm/hr] rain rate over output time step (e.g. 6hours) + s = (1000*Hour(1)/output.output_dt) + precip .*= s + + round!(precip, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, i] = precip + return nothing +end + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct LargeScalePrecipitationOutput <: AbstractOutputVariable + name::String = "precip_cond" + unit::String = "mm/hr" + long_name::String = "large-scale precipitation" + dims_xyzt::NTuple{4, Bool} = (true, true, false, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 7 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::LargeScalePrecipitationOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + precip = output.grid2D + (; precip_large_scale) = diagn.physics + RingGrids.interpolate!(precip, precip_large_scale, output.interpolator) + + # after output set precip accumulator back to zero + precip_large_scale .= 0 + + # convert from [m] to [mm/hr] rain rate over output time step (e.g. 6hours) + s = (1000*Hour(1)/output.output_dt) + precip .*= s + + round!(precip, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, i] = precip + return nothing +end + +## CLOUDS ------------- + +"""Defines netCDF output for a specific variables, see `VorticityOutput` for details. +Fields are $(TYPEDFIELDS)""" +@kwdef mutable struct CloudTopOutput <: AbstractOutputVariable + name::String = "cloud_top" + unit::String = "m" + long_name::String = "cloud top height" + dims_xyzt::NTuple{4, Bool} = (true, true, false, true) + missing_value::Float64 = NaN + compression_level::Int = 3 + shuffle::Bool = true + keepbits::Int = 7 +end + +"""$(TYPEDSIGNATURES) +`output!` method for `variable`, see `output!(::NetCDFOutput, ::VorticityOutput, ...)` for details.""" +function output!( + output::NetCDFOutput, + variable::CloudTopOutput, + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::AbstractModel, +) + cloud = output.grid2D + (; cloud_top) = diagn.physics + RingGrids.interpolate!(cloud, cloud_top, output.interpolator) + + round!(cloud, variable.keepbits) + i = output.output_counter # output time step to write + output.netcdf_file[variable.name][:, :, i] = cloud + return nothing +end + +""" +$(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, id::String) + # if run_???? folder doesn't exist yet don't change the id + run_id = string("run_", run_id_to_string(id)) + !(run_id in readdir(path)) && return id + + # otherwise 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 + +run_id_to_string(run_id::Integer) = @sprintf("%04d", run_id) +run_id_to_string(run_id::String) = run_id + +""" +$(TYPEDSIGNATURES) +Returns the full path of the output file after it was created. +""" +get_full_output_file_path(output::AbstractOutput) = joinpath(output.run_path, output.filename) + +""" +$(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}, model::AbstractModel) + @assert model.output.active "Output is turned off" + return Array(NCDataset(get_full_output_file_path(model.output))[string(var_name)]) +end \ No newline at end of file diff --git a/src/output/output.jl b/src/output/output.jl deleted file mode 100644 index f0fcfd384..000000000 --- a/src/output/output.jl +++ /dev/null @@ -1,568 +0,0 @@ -abstract type AbstractKeepbits end - -"""Number of mantissa bits to keep for each prognostic variable when compressed for -netCDF and .jld2 data output. -$(TYPEDFIELDS)""" -Base.@kwdef struct Keepbits - u::Int = 7 - v::Int = 7 - vor::Int = 5 - div::Int = 5 - temp::Int = 10 - pres::Int = 12 - humid::Int = 7 - precip_cond::Int = 7 - precip_conv::Int = 7 - cloud::Int = 7 -end - -function Base.show(io::IO, K::Keepbits) - println(io, "$(typeof(K))") - keys = propertynames(K) - print_fields(io, K, keys) -end - -# default number format for output -const DEFAULT_OUTPUT_NF = Float32 -const DEFAULT_OUTPUT_DT = Hour(6) - -export OutputWriter - -""" -$(TYPEDSIGNATURES) -NetCDF output writer. Contains all output options and auxiliary fields for output interpolation. -To be initialised with `OutputWriter(::SpectralGrid, ::Type{<:ModelSetup}, kwargs...)` to pass on the -resolution information and the model type which chooses which variables to output. Options include -$(TYPEDFIELDS)""" -Base.@kwdef mutable struct OutputWriter{NF<:Union{Float32, Float64}, Model<:ModelSetup} <: AbstractOutputWriter - - spectral_grid::SpectralGrid - - # FILE OPTIONS - output::Bool = false - - "[OPTION] path to output folder, run_???? will be created within" - path::String = pwd() - - "[OPTION] run identification number/string" - id::String = "0001" - run_path::String = "" # will be determined in initalize! - - "[OPTION] name of the output netcdf file" - filename::String = "output.nc" - - "[OPTION] also write restart file if output==true?" - write_restart::Bool = true - pkg_version::VersionNumber = pkgversion(SpeedyWeather) - - # WHAT/WHEN OPTIONS - startdate::DateTime = DateTime(2000, 1, 1) - - "[OPTION] output frequency, time step" - output_dt::Second = DEFAULT_OUTPUT_DT - - "[OPTION] which variables to output, u, v, vor, div, pres, temp, humid" - output_vars::Vector{Symbol} = default_output_vars(Model) - - "[OPTION] missing value to be used in netcdf output" - missing_value::NF = NaN - - # COMPRESSION OPTIONS - "[OPTION] lossless compression level; 1=low but fast, 9=high but slow" - compression_level::Int = 3 - - "[OPTION] shuffle/bittranspose filter for compression" - shuffle::Bool = true - - "[OPTION] mantissa bits to keep for every variable" - keepbits::Keepbits = Keepbits() - - # 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 - - # the netcdf file to be written into, will be created - netcdf_file::Union{NCDataset, Nothing} = nothing - - # INPUT GRID (the one used in the dynamical core) - input_Grid::Type{<:AbstractGrid} = spectral_grid.Grid - - # Output as matrix (particularly for reduced grids) - "[OPTION] sort grid points into a matrix (interpolation-free), for OctahedralClenshawGrid, OctaHEALPixGrid only" - as_matrix::Bool = 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 - "[OPTION] the grid used for output, full grids only" - output_Grid::Type{<:AbstractFullGrid} = RingGrids.full_grid_type(input_Grid) - - "[OPTION] the resolution of the output grid, default: same nlat_half as in the dynamical core" - nlat_half::Int = spectral_grid.nlat_half - 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(NF, input_Grid, spectral_grid.nlat_half, npoints) - - # fields to output (only one layer, reuse over layers) - 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) - const precip_cond::Matrix{NF} = fill(missing_value, nlon, nlat) - const precip_conv::Matrix{NF} = fill(missing_value, nlon, nlat) - const cloud::Matrix{NF} = fill(missing_value, nlon, nlat) -end - -# generator function pulling grid resolution and time stepping from ::SpectralGrid and ::AbstractTimeStepper -function OutputWriter( - spectral_grid::SpectralGrid, - ::Type{Model}; - NF::Type{<:Union{Float32, Float64}} = DEFAULT_OUTPUT_NF, - kwargs... -) where {Model<:ModelSetup} - return OutputWriter{NF, Model}(; spectral_grid, kwargs...) -end - -# default variables to output by model -default_output_vars(::Type{<:Barotropic}) = [:vor, :u, :v] -default_output_vars(::Type{<:ShallowWater}) = [:vor, :u, :v] -default_output_vars(::Type{<:PrimitiveDry}) = [:vor, :u, :v, :temp, :pres] -default_output_vars(::Type{<:PrimitiveWet}) = [:vor, :u, :v, :temp, :humid, :pres, :precip, :cloud] - -# print all fields with type <: Number -function Base.show(io::IO, O::AbstractOutputWriter) - println(io, "$(typeof(O))") - keys = propertynames(O) - - # remove interpolator from being printed, TODO: implement show for AbstractInterpolator - keys_filtered = filter(key -> ~(getfield(O, key) isa AbstractInterpolator), keys) - print_fields(io, O, keys_filtered) -end - -""" -$(TYPEDSIGNATURES) -Creates a netcdf file on disk and the corresponding `netcdf_file` object preallocated with output variables -and dimensions. `write_output!` then writes consecuitive time steps into this file. -""" -function initialize!( - output::OutputWriter{output_NF, Model}, - feedback::AbstractFeedback, - time_stepping::AbstractTimeStepper, - clock::Clock, - diagn::DiagnosticVariables, - model::ModelSetup, -) where {output_NF, Model} - - output.output || return nothing # exit immediately for no output - - # GET RUN ID, CREATE FOLDER - # get new id only if not already specified - 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! - - # OUTPUT FREQUENCY - output.output_every_n_steps = max(1, round(Int, - Millisecond(output.output_dt).value/time_stepping.Δt_millisec.value)) - output.output_dt = Second(round(Int, output.output_every_n_steps*time_stepping.Δt_sec)) - - # RESET COUNTERS - output.timestep_counter = 0 # time step counter - output.output_counter = 0 # output step counter - - # CREATE NETCDF FILE, vector of NcVars for output - (; run_path, filename, output_vars) = output - dataset = NCDataset(joinpath(run_path, filename), "c") - output.netcdf_file = dataset - - # DEFINE NETCDF DIMENSIONS TIME - (; startdate) = output - time_string = "hours since $(Dates.format(startdate, "yyyy-mm-dd HH:MM:0.0"))" - defDim(dataset, "time", Inf) # unlimited time dimension - defVar(dataset, "time", Float64, ("time",), - attrib=Dict("units"=>time_string, "long_name"=>"time", - "standard_name"=>"time", "calendar"=>"proleptic_gregorian")) - - # DEFINE NETCDF DIMENSIONS SPACE - (; input_Grid, output_Grid, nlat_half) = 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) = 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 - - # 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_NF.(model.geometry.σ_levels_full) - defVar(dataset, lon_name, lond, (lon_name,), attrib=Dict("units"=>lon_units, "long_name"=>lon_longname)) - defVar(dataset, lat_name, latd, (lat_name,), attrib=Dict("units"=>lat_units, "long_name"=>lat_longname)) - defVar(dataset, "lev", σ, ("lev",), attrib=Dict("units"=>"1", "long_name"=>"sigma levels")) - - # VARIABLES, define every variable here that could be output - (; compression_level) = output - missing_value = convert(output_NF, output.missing_value) - - # given pres the right name, depending on ShallowWaterModel or PrimitiveEquationModel - pres_name = Model <: ShallowWater ? "interface displacement" : "surface pressure" - pres_unit = Model <: ShallowWater ? "m" : "hPa" - - # zonal wind - u_attribs = Dict("long_name"=>"zonal wind", "units"=>"m/s", "_FillValue"=>missing_value) - :u in output_vars && defVar(dataset, "u", output_NF, (lon_name, lat_name, "lev", "time"), attrib=u_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # meridional wind - v_attribs = Dict("long_name"=>"meridional wind", "units"=>"m/s", "_FillValue"=>missing_value) - :v in output_vars && defVar(dataset, "v", output_NF, (lon_name, lat_name, "lev", "time"), attrib=v_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # vorticity - vor_attribs = Dict("long_name"=>"relative vorticity", "units"=>"1/s", "_FillValue"=>missing_value) - :vor in output_vars && defVar(dataset, "vor", output_NF, (lon_name, lat_name, "lev", "time"), attrib=vor_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # divergence - div_attribs = Dict("long_name"=>"divergence", "units"=>"1/s", "_FillValue"=>missing_value) - :div in output_vars && defVar(dataset, "div", output_NF, (lon_name, lat_name, "lev", "time"), attrib=div_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # pressure / interface displacement - pres_attribs = Dict("long_name"=>pres_name, "units"=>pres_unit, "_FillValue"=>missing_value) - :pres in output_vars && defVar(dataset, "pres", output_NF, (lon_name, lat_name, "time"), attrib=pres_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # temperature - temp_attribs = Dict("long_name"=>"temperature", "units"=>"degC", "_FillValue"=>missing_value) - :temp in output_vars && defVar(dataset, "temp", output_NF, (lon_name, lat_name, "lev", "time"), attrib=temp_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # humidity - humid_attribs = Dict("long_name"=>"specific humidity", "units"=>"kg/kg", "_FillValue"=>missing_value) - :humid in output_vars && defVar(dataset, "humid", output_NF, (lon_name, lat_name, "lev", "time"), attrib=humid_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # orography - if :orography in output_vars # write orography directly to file - orog_attribs = Dict("long_name"=>"orography", "units"=>"m", "_FillValue"=>missing_value) - orog_grid = model.orography.orography - orog_matrix = output.u - output.as_matrix && (orog_matrix = Matrix(orog_grid)) - output.as_matrix || RingGrids.interpolate!(output_Grid(output.u, input_as=Matrix), orog_grid, output.interpolator) - defVar(dataset, "orography", orog_matrix, (lon_name, lat_name), attrib=orog_attribs) - end - - # large-scale condensation - precip_cond_attribs = Dict("long_name"=>"large-scale precipitation", "units"=>"mm/hr", "_FillValue"=>missing_value) - :precip in output_vars && defVar(dataset, "precip_cond", output_NF, (lon_name, lat_name, "time"), attrib=precip_cond_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # convective precipitation - precip_conv_attribs = Dict("long_name"=>"convective precipitation", "units"=>"mm/hr", "_FillValue"=>missing_value) - :precip in output_vars && defVar(dataset, "precip_conv", output_NF, (lon_name, lat_name, "time"), attrib=precip_conv_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # convective precipitation - cloud_top_attribs = Dict("long_name"=>"cloud top height", "units"=>"m", "_FillValue"=>missing_value) - :cloud in output_vars && defVar(dataset, "cloud_top", output_NF, (lon_name, lat_name, "time"), attrib=cloud_top_attribs, - deflatelevel=compression_level, shuffle=output.shuffle) - - # WRITE INITIAL CONDITIONS TO FILE - write_netcdf_variables!(output, diagn) - write_netcdf_time!(output, clock.time) - - # 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 - - -""" -$(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, id::String) - # if run_???? folder doesn't exist yet don't change the id - run_id = string("run_", run_id_to_string(id)) - !(run_id in readdir(path)) && return id - - # otherwise 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 - -run_id_to_string(run_id::Integer) = @sprintf("%04d", run_id) -run_id_to_string(run_id::String) = run_id - - -""" -$(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 - output || return nothing # escape immediately for no netcdf output - timestep_counter % output_every_n_steps == 0 || return nothing # escape if output not written on this step - - # WRITE VARIABLES - write_netcdf_variables!(outputter, diagn) - write_netcdf_time!(outputter, time) -end - -""" -$(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 ) = output - i = output.output_counter - - time_passed = Millisecond(time-startdate) - time_hrs = time_passed.value/3600_000 # [ms] -> [hrs] - netcdf_file["time"][i] = time_hrs - NCDatasets.sync(netcdf_file) - - return nothing -end - -""" -$(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} - - output.output_counter += 1 # increase output step counter - (; output_vars) = output # Vector{Symbol} of variables to output - i = output.output_counter - - (; u, v, vor, div, pres, temp, humid, precip_cond, precip_conv, cloud) = output - (; output_Grid, interpolator) = output - (; quadrant_rotation, matrix_quadrant) = output - (; netcdf_file, keepbits) = 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.as_matrix # resort gridded variables interpolation-free into a matrix - - # create (matrix, grid) tuples for simultaneous grid -> matrix conversion - # TODO this currently does the Matrix! conversion to all variables, not just output_vars - # as arrays are always initialised - MGs = ((M, G) for (M, G) in zip((u, v, vor, div, temp, humid), - (u_grid, v_grid, vor_grid, div_grid, temp_grid, humid_grid)) - if length(M) > 0) - - RingGrids.Matrix!(MGs...; quadrant_rotation, matrix_quadrant) - - else # or interpolate onto a full grid - :u in output_vars && RingGrids.interpolate!(output_Grid(u, input_as=Matrix), u_grid, interpolator) - :v in output_vars && RingGrids.interpolate!(output_Grid(v, input_as=Matrix), v_grid, interpolator) - :vor in output_vars && RingGrids.interpolate!(output_Grid(vor, input_as=Matrix), vor_grid, interpolator) - :div in output_vars && RingGrids.interpolate!(output_Grid(div, input_as=Matrix), div_grid, interpolator) - :temp in output_vars && RingGrids.interpolate!(output_Grid(temp, input_as=Matrix), temp_grid, interpolator) - :humid in output_vars && RingGrids.interpolate!(output_Grid(humid, input_as=Matrix), humid_grid, interpolator) - end - - # UNSCALE THE SCALED VARIABLES - 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 - :u in output_vars && round!(u, keepbits.u) - :v in output_vars && round!(v, keepbits.v) - :vor in output_vars && round!(vor, keepbits.vor) - :div in output_vars && round!(div, keepbits.div) - :temp in output_vars && round!(temp, keepbits.temp) - :humid in output_vars && round!(humid, keepbits.humid) - - # WRITE VARIABLES TO FILE, APPEND IN TIME DIMENSION - :u in output_vars && (netcdf_file["u"][:, :, k, i] = u) - :v in output_vars && (netcdf_file["v"][:, :, k, i] = v) - :vor in output_vars && (netcdf_file["vor"][:, :, k, i] = vor) - :div in output_vars && (netcdf_file["div"][:, :, k, i] = div) - :temp in output_vars && (netcdf_file["temp"][:, :, k, i] = temp) - :humid in output_vars && (netcdf_file["humid"][:, :, k, i] = humid) - end - - # surface pressure, i.e. interface displacement η - (; pres_grid, precip_large_scale, precip_convection, cloud_top ) = diagn.surface - - if output.as_matrix - if :pres in output_vars || :precip_cond in output_vars || :precip_conv in output_vars || :cloud in output_vars - MGs = ((M, G) for (M, G) in zip((pres, precip_cond, precip_conv, cloud), - (pres_grid, precip_large_scale, precip_convection, cloud_top))) - - RingGrids.Matrix!(MGs...; quadrant_rotation, matrix_quadrant) - end - else - :pres in output_vars && RingGrids.interpolate!(output_Grid(pres, input_as=Matrix), pres_grid, interpolator) - :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_cond, input_as=Matrix), precip_large_scale, interpolator) - :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_conv, input_as=Matrix), precip_convection, interpolator) - :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_conv, input_as=Matrix), precip_convection, interpolator) - :cloud in output_vars && RingGrids.interpolate!(output_Grid(cloud, input_as=Matrix), cloud_top, interpolator) - end - - # after output set precip accumulators back to zero - precip_large_scale .= 0 - precip_convection .= 0 - - # convert from [m] to [mm/hr] rain rate over output time step (e.g. 6hours) - s = (1000*Hour(1)/output.output_dt) - precip_cond *= s - precip_conv *= s - - if Model <: PrimitiveEquation - @. pres = exp(pres)/100 # convert from log(pₛ) to surface pressure pₛ [hPa] - end - - :pres in output_vars && round!(pres, keepbits.pres) - :precip in output_vars && round!(precip_cond, keepbits.precip_cond) - :precip in output_vars && round!(precip_conv, keepbits.precip_conv) - :cloud in output_vars && round!(cloud, keepbits.cloud) - - :pres in output_vars && (netcdf_file["pres"][:, :, i] = pres) - :precip in output_vars && (netcdf_file["precip_cond"][:, :, i] = precip_cond) - :precip in output_vars && (netcdf_file["precip_conv"][:, :, i] = precip_conv) - :cloud in output_vars && (netcdf_file["cloud_top"][:, :, i] = cloud) - - return nothing -end - -Base.close(output::OutputWriter) = NCDatasets.close(output.netcdf_file) -Base.close(::Nothing) = nothing - -""" -$(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 bitrounded for compression and the 2nd leapfrog time step is discarded. -Variables in restart file are unscaled.""" -function write_restart_file(progn::PrognosticVariables{T}, - output::OutputWriter) where T - - (; 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 - for layer in progn.layers - - # copy over leapfrog 2 to 1 - copyto!(layer.timesteps[1].vor, layer.timesteps[2].vor) - copyto!(layer.timesteps[1].div, layer.timesteps[2].div) - copyto!(layer.timesteps[1].temp, layer.timesteps[2].temp) - copyto!(layer.timesteps[1].humid, layer.timesteps[2].humid) - - # bitround 1st leapfrog step to output precision - if T <: Base.IEEEFloat # currently not defined for other formats... - round!(layer.timesteps[1].vor, keepbits.vor) - round!(layer.timesteps[1].div, keepbits.div) - round!(layer.timesteps[1].temp, keepbits.temp) - round!(layer.timesteps[1].humid, keepbits.humid) - end - - # remove 2nd leapfrog step by filling with zeros - fill!(layer.timesteps[2].vor, 0) - fill!(layer.timesteps[2].div, 0) - fill!(layer.timesteps[2].temp, 0) - fill!(layer.timesteps[2].humid, 0) - end - - # same for surface pressure - copyto!(progn.surface.timesteps[1].pres, progn.surface.timesteps[2].pres) - T <: Base.IEEEFloat && round!(progn.surface.timesteps[1].pres, keepbits.pres) - fill!(progn.surface.timesteps[2].pres, 0) - - jldopen(joinpath(run_path, "restart.jld2"), "w"; compress=true) do f - f["prognostic_variables"] = progn - f["version"] = output.pkg_version - f["description"] = "Restart file created for SpeedyWeather.jl" - end -end - -""" -$(TYPEDSIGNATURES) -Returns the full path of the output file after it was created. -""" -get_full_output_file_path(output::OutputWriter) = joinpath(output.run_path, output.filename) - -""" -$(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}, model::ModelSetup) - @assert model.output.output "Output is turned off" - return Array(NCDataset(get_full_output_file_path(model.output))[string(var_name)]) -end diff --git a/src/output/particle_tracker.jl b/src/output/particle_tracker.jl index 9789bfe8f..5969b10ae 100644 --- a/src/output/particle_tracker.jl +++ b/src/output/particle_tracker.jl @@ -26,31 +26,31 @@ Base.@kwdef mutable struct ParticleTracker{NF} <: AbstractCallback keepbits::Int = 15 "Number of particles to track" - n_particles::Int = 0 + nparticles::Int = 0 "The netcdf file to be written into, will be created at initialization" netcdf_file::Union{NCDataset, Nothing} = nothing # tracking arrays - lon::Vector{NF} = zeros(NF, n_particles) - lat::Vector{NF} = zeros(NF, n_particles) - σ::Vector{NF} = zeros(NF, n_particles) + lon::Vector{NF} = zeros(NF, nparticles) + lat::Vector{NF} = zeros(NF, nparticles) + σ::Vector{NF} = zeros(NF, nparticles) end ParticleTracker(SG::SpectralGrid; kwargs...) = - ParticleTracker{SG.NF}(;n_particles = SG.n_particles, kwargs...) + ParticleTracker{SG.NF}(;nparticles = SG.nparticles, kwargs...) function initialize!( callback::ParticleTracker{NF}, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) where NF initialize!(callback.schedule, progn.clock) # if model.output doesn't output create a folder anyway to store the particles.nc file - if model.output.output == false + if model.output.active == false (;output, feedback) = model output.id = get_run_id(output.path, output.id) output.run_path = create_output_folder(output.path, output.id) @@ -76,13 +76,13 @@ function initialize!( "standard_name"=>"time", "calendar"=>"proleptic_gregorian")) # PARTICLE DIMENSION - n_particles = length(progn.particles) - callback.n_particles = n_particles - defDim(dataset, "particle", n_particles) + nparticles = length(progn.particles) + callback.nparticles = nparticles + defDim(dataset, "particle", nparticles) defVar(dataset, "particle", Int64, ("particle",), attrib=Dict("units"=>"1", "long_name"=>"particle identification number")) - # coordinates of particles (the variables inside netCDF) + # coordinates of particles (the variables inside netCDF) defVar(dataset, "lon", NF, ("particle", "time"), attrib = Dict("long_name"=>"longitude", "units"=>"degrees_north"), deflatelevel=callback.compression_level, shuffle=callback.shuffle) @@ -95,14 +95,14 @@ function initialize!( Dict("long_name"=>"vertical sigma coordinate", "units"=>"1"), deflatelevel=callback.compression_level, shuffle=callback.shuffle) - # pull particle locations into output work arrays + # pull particle locations into output work arrays for (p,particle) in enumerate(progn.particles) callback.lon[p] = particle.lon callback.lat[p] = particle.lat callback.σ[p] = particle.σ end - # rounding + # rounding (; keepbits) = callback round!(callback.lon, keepbits) round!(callback.lat, keepbits) @@ -120,19 +120,19 @@ function callback!( callback::ParticleTracker, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) isscheduled(callback.schedule, progn.clock) || return nothing # else escape immediately i = callback.schedule.counter+1 # +1 for initial conditions (not scheduled) - # pull particle locations into output work arrays + # pull particle locations into output work arrays for (p,particle) in enumerate(progn.particles) callback.lon[p] = particle.lon callback.lat[p] = particle.lat callback.σ[p] = particle.σ end - # rounding + # rounding (; keepbits) = callback round!(callback.lon, keepbits) round!(callback.lat, keepbits) @@ -140,7 +140,7 @@ function callback!( # write current particle locations to file (;time, start) = progn.clock - time_passed_hrs = Millisecond(time - start).value/3600_000 # [ms] -> [hrs] + time_passed_hrs = Millisecond(time - start).value/3600_000 # [ms] -> [hrs] callback.netcdf_file["time"][i] = time_passed_hrs callback.netcdf_file["lon"][:, i] = callback.lon callback.netcdf_file["lat"][:, i] = callback.lat diff --git a/src/output/restart_file.jl b/src/output/restart_file.jl new file mode 100644 index 000000000..34155ee61 --- /dev/null +++ b/src/output/restart_file.jl @@ -0,0 +1,44 @@ +""" +$(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 bitrounded for compression and the 2nd leapfrog time step is discarded. +Variables in restart file are unscaled.""" +function write_restart_file( + output::AbstractOutput, + progn::PrognosticVariables{T}, +) where T + + # exit immediately if no output or no restart file desired + output.active && output.write_restart || return nothing + + # move 2nd leapfrog to 1st to compress restart file + copyto!(progn.vor[1], progn.vor[2]) + copyto!(progn.div[1], progn.div[2]) + copyto!(progn.temp[1], progn.temp[2]) + copyto!(progn.humid[1], progn.humid[2]) + copyto!(progn.pres[1], progn.pres[2]) + + # bitround 1st leapfrog step to output precision + if T <: Base.IEEEFloat # currently not defined for other formats... + round!(progn.vor[1], 7) # hardcode some defaults for now + round!(progn.div[1], 7) + round!(progn.temp[1], 12) + round!(progn.humid[1], 10) + round!(progn.pres[1], 14) + end + + # remove 2nd leapfrog step by filling with zeros + fill!(progn.vor[2], 0) + fill!(progn.div[2], 0) + fill!(progn.temp[2], 0) + fill!(progn.humid[2], 0) + fill!(progn.pres[2], 0) + + jldopen(joinpath(output.run_path, "restart.jld2"), "w"; compress=true) do f + f["prognostic_variables"] = progn + f["version"] = output.pkg_version + f["description"] = "Restart file created for SpeedyWeather.jl" + end +end \ No newline at end of file diff --git a/src/output/schedule.jl b/src/output/schedule.jl index 4a28ab6dd..1ea0d14e7 100644 --- a/src/output/schedule.jl +++ b/src/output/schedule.jl @@ -53,7 +53,7 @@ function initialize!(scheduler::Schedule, clock::Clock) # PERIODIC SCHEDULE, always AFTER scheduler.every time period has passed if scheduler.every.value < typemax(Int) - every_n_timesteps = max(1,round(Int, scheduler.every/clock.Δt)) + every_n_timesteps = max(1, round(Int, scheduler.every/clock.Δt)) schedule[every_n_timesteps:every_n_timesteps:end] .= true prev_every = readable_secs(scheduler.every.value) diff --git a/src/physics/boundary_layer.jl b/src/physics/boundary_layer.jl index 41d6202b5..c9f705ecb 100644 --- a/src/physics/boundary_layer.jl +++ b/src/physics/boundary_layer.jl @@ -17,23 +17,23 @@ export LinearDrag """Linear boundary layer drag following Held and Suarez, 1996 BAMS $(TYPEDFIELDS)""" -Base.@kwdef struct LinearDrag{NF<:AbstractFloat} <: AbstractBoundaryLayer +@kwdef struct LinearDrag{NF<:AbstractFloat} <: AbstractBoundaryLayer # DIMENSIONS - nlev::Int + nlayers::Int # PARAMETERS σb::NF = 0.7 # sigma coordinate below which linear drag is applied time_scale::Second = Hour(24) # time scale for linear drag coefficient at σ=1 (=1/kf in HS96) # PRECOMPUTED CONSTANTS - drag_coefs::Vector{NF} = zeros(NF, nlev) + drag_coefs::Vector{NF} = zeros(NF, nlayers) end """ $(TYPEDSIGNATURES) -Generator function using `nlev` from `SG::SpectralGrid`""" -LinearDrag(SG::SpectralGrid; kwargs...) = LinearDrag{SG.NF}(nlev=SG.nlev; kwargs...) +Generator function using `nlayers` from `SG::SpectralGrid`""" +LinearDrag(SG::SpectralGrid; kwargs...) = LinearDrag{SG.NF}(nlayers=SG.nlayers; kwargs...) """ $(TYPEDSIGNATURES) @@ -94,7 +94,7 @@ export BulkRichardsonDrag """Boundary layer drag coefficient from the bulk Richardson number, following Frierson, 2006, Journal of the Atmospheric Sciences. $(TYPEDFIELDS)""" -Base.@kwdef struct BulkRichardsonDrag{NF} <: AbstractBoundaryLayer +@kwdef struct BulkRichardsonDrag{NF} <: AbstractBoundaryLayer "von Kármán constant [1]" κ::NF = 0.4 @@ -164,11 +164,13 @@ function bulk_richardson_surface( atmosphere::AbstractAtmosphere, ) cₚ = atmosphere.heat_capacity - (; u, v, geopot, temp_virt, nlev) = column - - V² = u[nlev]^2 + v[nlev]^2 - Θ₀ = cₚ*temp_virt[nlev] - Θ₁ = Θ₀ + geopot[nlev] - bulk_richardson = geopot[nlev]*(Θ₁ - Θ₀) / (Θ₀*V²) + (; u, v, geopot, temp_virt) = column + surface = column.nlayers # surface index = nlayers + + V² = u[surface]^2 + v[surface]^2 + Θ₀ = cₚ*temp_virt[surface] + Θ₁ = Θ₀ + geopot[surface] + bulk_richardson = geopot[surface]*(Θ₁ - Θ₀) / (Θ₀*V²) + return bulk_richardson end diff --git a/src/physics/column_variables.jl b/src/physics/column_variables.jl index a79c3460e..f541da440 100644 --- a/src/physics/column_variables.jl +++ b/src/physics/column_variables.jl @@ -37,7 +37,7 @@ function get_column!( (; σ_levels_full, ln_σ_levels_full) = geometry (; temp_profile) = implicit # reference temperature - @boundscheck C.nlev == D.nlev || throw(BoundsError) + @boundscheck C.nlayers == D.nlayers || throw(BoundsError) C.latd = geometry.latds[ij] # pull latitude, longitude [˚N, ˚E] for gridpoint ij from Geometry C.lond = geometry.londs[ij] @@ -48,41 +48,42 @@ function get_column!( C.surface_geopotential = C.orography * planet.gravity # pressure [Pa] or [log(Pa)] - lnpₛ = D.surface.pres_grid[ij] # logarithm of surf pressure used in dynamics + lnpₛ = D.grid.pres_grid[ij] # logarithm of surf pressure used in dynamics pₛ = exp(lnpₛ) # convert back here C.ln_pres .= ln_σ_levels_full .+ lnpₛ # log pressure on every level ln(p) = ln(σ) + ln(pₛ) C.pres[1:end-1] .= σ_levels_full.*pₛ # pressure on every level p = σ*pₛ C.pres[end] = pₛ # last element is surface pressure pₛ - @inbounds for (k, layer) = enumerate(D.layers) - # read out prognostic variables on grid at previous time step - C.u[k] = layer.grid_variables.u_grid_prev[ij] - C.v[k] = layer.grid_variables.v_grid_prev[ij] + (; u_grid_prev, v_grid_prev, temp_grid_prev, humid_grid_prev) = D.grid - # add temp reference profile back in as temp_grid_prev is anomaly - C.temp[k] = layer.grid_variables.temp_grid_prev[ij] + temp_profile[k] - C.humid[k] = layer.grid_variables.humid_grid_prev[ij] + @inbounds for k in eachgrid(u_grid_prev, v_grid_prev, temp_grid_prev, humid_grid_prev) + # read out prognostic variables on grid at previous time step + # for numerical stability + C.u[k] = u_grid_prev[ij, k] + C.v[k] = v_grid_prev[ij, k] + C.temp[k] = temp_grid_prev[ij, k] + temp_profile[k] + C.humid[k] = humid_grid_prev[ij, k] end # TODO skin = surface approximation for now C.skin_temperature_sea = P.ocean.sea_surface_temperature[ij] C.skin_temperature_land = P.land.land_surface_temperature[ij] - C.soil_moisture_availability = D.surface.soil_moisture_availability[ij] + C.soil_moisture_availability = D.physics.soil_moisture_availability[ij] # RADIATION - C.cos_zenith = D.surface.cos_zenith[ij] + C.cos_zenith = D.physics.cos_zenith[ij] C.albedo = albedo.albedo[ij] end """Recalculate ring index if not provided.""" -function get_column!( C::ColumnVariables, - D::DiagnosticVariables, - P::PrognosticVariables, - ij::Int, # grid point index - model::PrimitiveEquation) - - SG = model.spectral_grid - rings = eachring(SG.Grid, SG.nlat_half) +function get_column!( + C::ColumnVariables, + D::DiagnosticVariables, + P::PrognosticVariables, + ij::Int, # grid point index + model::PrimitiveEquation +) + rings = eachring(D.grid.vor_grid) jring = whichring(ij, rings) get_column!(C, D, P, ij, jring, model) end @@ -112,32 +113,33 @@ end $(TYPEDSIGNATURES) Write the parametrization tendencies from `C::ColumnVariables` into the horizontal fields of tendencies stored in `D::DiagnosticVariables` at gridpoint index `ij`.""" -function write_column_tendencies!( D::DiagnosticVariables, - column::ColumnVariables, - planet::AbstractPlanet, - ij::Int) # grid point index - - (; nlev) = column - @boundscheck nlev == D.nlev || throw(BoundsError) - - @inbounds for (k, layer) = enumerate(D.layers) - layer.tendencies.u_tend_grid[ij] = column.u_tend[k] - layer.tendencies.v_tend_grid[ij] = column.v_tend[k] - layer.tendencies.temp_tend_grid[ij] = column.temp_tend[k] - layer.tendencies.humid_tend_grid[ij] = column.humid_tend[k] +function write_column_tendencies!( + diagn::DiagnosticVariables, + column::ColumnVariables, + planet::AbstractPlanet, + ij::Integer, # grid point index +) + (; nlayers) = column + @boundscheck nlayers == diagn.nlayers || throw(BoundsError) + + @inbounds for k in eachgrid(diagn.grid.vor_grid) + diagn.tendencies.u_tend_grid[ij, k] = column.u_tend[k] + diagn.tendencies.v_tend_grid[ij, k] = column.v_tend[k] + diagn.tendencies.temp_tend_grid[ij, k] = column.temp_tend[k] + diagn.tendencies.humid_tend_grid[ij, k] = column.humid_tend[k] end # accumulate (set back to zero when netcdf output) - D.surface.precip_large_scale[ij] += column.precip_large_scale - D.surface.precip_convection[ij] += column.precip_convection + diagn.physics.precip_large_scale[ij] += column.precip_large_scale + diagn.physics.precip_convection[ij] += column.precip_convection # Output cloud top in height [m] from geopotential height divided by gravity, # but NaN for no clouds - D.surface.cloud_top[ij] = column.cloud_top == nlev+1 ? 0 : column.geopot[column.cloud_top] - D.surface.cloud_top[ij] /= planet.gravity + diagn.physics.cloud_top[ij] = column.cloud_top == nlayers+1 ? 0 : column.geopot[column.cloud_top] + diagn.physics.cloud_top[ij] /= planet.gravity - # just use layer index 1 (top) to nlev (surface) for analysis, but 0 for no clouds - # D.surface.cloud_top[ij] = column.cloud_top == nlev+1 ? 0 : column.cloud_top + # just use layer index 1 (top) to nlayers (surface) for analysis, but 0 for no clouds + # diagn.physics.cloud_top[ij] = column.cloud_top == nlayers+1 ? 0 : column.cloud_top return nothing end @@ -164,7 +166,7 @@ function reset_column!(column::ColumnVariables{NF}) where NF column.flux_temp_downward .= 0 # Convection and precipitation - column.cloud_top = column.nlev+1 # also diagnostic from condensation + column.cloud_top = column.nlayers+1 column.precip_convection = 0 column.precip_large_scale = 0 @@ -172,5 +174,5 @@ function reset_column!(column::ColumnVariables{NF}) where NF end # iterator for convenience -eachlayer(column::ColumnVariables) = Base.OneTo(column.nlev) -Base.eachindex(column::ColumnVariables) = Base.OneTo(column.nlev) \ No newline at end of file +eachlayer(column::ColumnVariables) = eachindex(column) +Base.eachindex(column::ColumnVariables) = Base.OneTo(column.nlayers) \ No newline at end of file diff --git a/src/physics/convection.jl b/src/physics/convection.jl index 0111a72b0..758aacf95 100644 --- a/src/physics/convection.jl +++ b/src/physics/convection.jl @@ -13,9 +13,9 @@ The simplified Betts-Miller convection scheme from Frierson, 2007, https://doi.org/10.1175/JAS3935.1. This implements the qref-formulation in their paper. Fields and options are $(TYPEDFIELDS)""" -Base.@kwdef struct SimplifiedBettsMiller{NF} <: AbstractConvection - "number of vertical layers/levels" - nlev::Int +@kwdef struct SimplifiedBettsMiller{NF} <: AbstractConvection + "number of vertical layers" + nlayers::Int "[OPTION] Relaxation time for profile adjustment" time_scale::Second = Hour(4) @@ -24,7 +24,7 @@ Base.@kwdef struct SimplifiedBettsMiller{NF} <: AbstractConvection relative_humidity::NF = 0.7 end -SimplifiedBettsMiller(SG::SpectralGrid; kwargs...) = SimplifiedBettsMiller{SG.NF}(nlev=SG.nlev; kwargs...) +SimplifiedBettsMiller(SG::SpectralGrid; kwargs...) = SimplifiedBettsMiller{SG.NF}(nlayers=SG.nlayers; kwargs...) initialize!(::SimplifiedBettsMiller, ::PrimitiveEquation) = nothing # function barrier for all AbstractConvection @@ -65,7 +65,7 @@ function convection!( σ = geometry.σ_levels_full σ_half = geometry.σ_levels_half Δσ = geometry.σ_levels_thick - (; geopot, nlev, temp, temp_virt, humid, temp_tend, humid_tend) = column + (; geopot, nlayers, temp, temp_virt, humid, temp_tend, humid_tend) = column pₛ = column.pres[end] (; Lᵥ, cₚ) = clausius_clapeyron @@ -74,14 +74,14 @@ function convection!( humid_ref_profile = column.b # specific humidity [kg/kg] profile to adjust to # CONVECTIVE CRITERIA AND FIRST GUESS RELAXATION - temp_parcel = temp[nlev] - humid_parcel = humid[nlev] + temp_parcel = temp[nlayers] + humid_parcel = humid[nlayers] level_zero_buoyancy = pseudo_adiabat!(temp_ref_profile, temp_parcel, humid_parcel, temp_virt, geopot, pₛ, σ, clausius_clapeyron) - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers qsat = saturation_humidity(temp_ref_profile[k], pₛ*σ[k], clausius_clapeyron) humid_ref_profile[k] = qsat*SBM.relative_humidity end @@ -92,7 +92,7 @@ function convection!( local Qref::NF = 0 # = ∫_pₛ^p_LZB -humid_ref_profile dp # skip constants compared to Frierson 2007, i.e. no /τ, /gravity, *cₚ/Lᵥ - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers # Frierson's equation (1) # δq = -(humid[k] - humid_ref_profile[k])/SBM.time_scale.value # Pq -= δq*Δσ[k]/gravity @@ -114,13 +114,13 @@ function convection!( no_convection && return nothing # height of zero buoyancy level in σ coordinates - Δσ_lzb = σ_half[nlev+1] - σ_half[level_zero_buoyancy] + Δσ_lzb = σ_half[nlayers+1] - σ_half[level_zero_buoyancy] if deep_convection ΔT = (PT - Pq*Lᵥ/cₚ)/Δσ_lzb # eq (5) but reusing PT, Pq, and /cₚ already included - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers temp_ref_profile[k] -= ΔT # equation (6) end @@ -130,20 +130,20 @@ function convection!( # "changing the reference profiles for both temperature and humidity so the # precipitation is zero. - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers Qref -= humid_ref_profile[k]*Δσ[k] # eq (11) but in σ coordinates end fq = 1 - Pq/Qref # = 1 - Δq/Qref in eq (12) but we reuse Pq ΔT = PT/Δσ_lzb # equation (14), reuse PT and in σ coordinates - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers humid_ref_profile[k] *= fq # update humidity profile, eq (13) temp_ref_profile[k] -= ΔT # update temperature profile, eq (15) end end # GET TENDENCIES FROM ADJUSTED PROFILES - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers temp_tend[k] -= (temp[k] - temp_ref_profile[k]) / SBM.time_scale.value δq = (humid[k] - humid_ref_profile[k]) / SBM.time_scale.value humid_tend[k] -= δq @@ -185,13 +185,13 @@ function pseudo_adiabat!( @boundscheck length(temp_ref_profile) == length(geopot) == length(σ) == length(temp_virt_environment) || throw(BoundsError) - nlev = length(temp_ref_profile) # number of vertical levels + nlayers = length(temp_ref_profile) # number of vertical levels temp_ref_profile .= NaN # reset profile from any previous calculation - temp_ref_profile[nlev] = temp_parcel # start profile with parcel temperature + temp_ref_profile[nlayers] = temp_parcel # start profile with parcel temperature local saturated::Bool = false # did the parcel reach saturation yet? local buoyant::Bool = true # is the parcel still buoyant? - local k::Int = nlev # layer index top to surface + local k::Int = nlayers # layer index top to surface local temp_virt_parcel::NF = temp_parcel * (1 + μ*humid_parcel) while buoyant && k > 1 # calculate moist adiabat while buoyant till top @@ -248,15 +248,15 @@ The simplified Betts-Miller convection scheme from Frierson, 2007, https://doi.org/10.1175/JAS3935.1 but with humidity set to zero. Fields and options are $(TYPEDFIELDS)""" -Base.@kwdef struct DryBettsMiller{NF} <: AbstractConvection +@kwdef struct DryBettsMiller{NF} <: AbstractConvection "number of vertical layers/levels" - nlev::Int + nlayers::Int "[OPTION] Relaxation time for profile adjustment" time_scale::Second = Hour(4) end -DryBettsMiller(SG::SpectralGrid; kwargs...) = DryBettsMiller{SG.NF}(nlev=SG.nlev; kwargs...) +DryBettsMiller(SG::SpectralGrid; kwargs...) = DryBettsMiller{SG.NF}(nlayers=SG.nlayers; kwargs...) initialize!(::DryBettsMiller, ::PrimitiveEquation) = nothing # function barrier to unpack model @@ -286,7 +286,7 @@ function convection!( σ = geometry.σ_levels_full σ_half = geometry.σ_levels_half Δσ = geometry.σ_levels_thick - (; nlev, temp, temp_tend) = column + (; nlayers, temp, temp_tend) = column # use work arrays for temp_ref_profile temp_ref_profile = column.a # temperature [K] reference profile to adjust to @@ -300,7 +300,7 @@ function convection!( local ΔT::NF = 0 # vertically uniform temperature profile adjustment # skip constants compared to Frierson 2007, i.e. no /τ, /gravity, *cₚ/Lᵥ - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers # Frierson's equation (1) # δT = -(temp[k] - temp_ref_profile[k])/SBM.time_scale.value # PT += δT*Δσ[k]/gravity*cₚ/Lᵥ @@ -314,9 +314,9 @@ function convection!( convection || return nothing # escape immediately for no convection # height of zero buoyancy level in σ coordinates - Δσ_lzb = σ_half[nlev+1] - σ_half[level_zero_buoyancy] + Δσ_lzb = σ_half[nlayers+1] - σ_half[level_zero_buoyancy] ΔT = PT/Δσ_lzb # eq (5) or (14) but reusing PT - for k in level_zero_buoyancy:nlev + for k in level_zero_buoyancy:nlayers temp_ref_profile[k] -= ΔT # equation (6) or equation (15) temp_tend[k] -= (temp[k] - temp_ref_profile[k]) / DBM.time_scale.value end @@ -341,13 +341,13 @@ function dry_adiabat!( @boundscheck length(temp_ref_profile) == length(σ) == length(temp_environment) || throw(BoundsError) - nlev = length(temp_ref_profile) # number of vertical levels - temp_parcel = temp_environment[nlev] # parcel is at lowermost layer temperature + nlayers = length(temp_ref_profile) # number of vertical levels + temp_parcel = temp_environment[nlayers] # parcel is at lowermost layer temperature temp_ref_profile .= NaN # reset profile from any previous calculation - temp_ref_profile[nlev] = temp_parcel # start profile with parcel temperature + temp_ref_profile[nlayers] = temp_parcel # start profile with parcel temperature local buoyant::Bool = true # is the parcel still buoyant? - local k::Int = nlev # layer index top to surface + local k::Int = nlayers # layer index top to surface while buoyant && k > 1 # calculate moist adiabat while buoyant till top k -= 1 # one level up diff --git a/src/physics/define_column.jl b/src/physics/define_column.jl index a93869041..6b0ccf70a 100644 --- a/src/physics/define_column.jl +++ b/src/physics/define_column.jl @@ -4,13 +4,13 @@ export ColumnVariables """ 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 +to iterate over horizontal grid points. Most column vectors have `nlayers` entries, from [1] at the top to [end] at the lowermost model level at the planetary boundary layer. $(TYPEDFIELDS)""" -Base.@kwdef mutable struct ColumnVariables{NF<:AbstractFloat} <: AbstractColumnVariables +@kwdef mutable struct ColumnVariables{NF<:AbstractFloat} <: AbstractColumnVariables # DIMENSIONS - const nlev::Int = 0 # number of vertical levels + const nlayers::Int = 0 # number of vertical levels # COORDINATES ij::Int = 0 # grid-point index @@ -21,33 +21,33 @@ Base.@kwdef mutable struct ColumnVariables{NF<:AbstractFloat} <: AbstractColumnV orography::NF = 0 # orography height [m] # PROGNOSTIC VARIABLES - const u::Vector{NF} = zeros(NF, nlev) # zonal velocity [m/s] - const v::Vector{NF} = zeros(NF, nlev) # meridional velocity [m/s] - const temp::Vector{NF} = zeros(NF, nlev) # absolute temperature [K] - const humid::Vector{NF} = zeros(NF, nlev) # specific humidity [kg/kg] + const u::Vector{NF} = zeros(NF, nlayers) # zonal velocity [m/s] + const v::Vector{NF} = zeros(NF, nlayers) # meridional velocity [m/s] + const temp::Vector{NF} = zeros(NF, nlayers) # absolute temperature [K] + const humid::Vector{NF} = zeros(NF, nlayers) # specific humidity [kg/kg] # (log) pressure per layer, surface is prognostic, last element here, but precompute other layers too - const ln_pres::Vector{NF} = zeros(NF, nlev+1) # logarithm of pressure [log(Pa)] - const pres::Vector{NF} = zeros(NF, nlev+1) # pressure [Pa] + const ln_pres::Vector{NF} = zeros(NF, nlayers+1) # logarithm of pressure [log(Pa)] + const pres::Vector{NF} = zeros(NF, nlayers+1) # pressure [Pa] # TENDENCIES to accumulate the parametrizations into - const u_tend::Vector{NF} = zeros(NF, nlev) # zonal velocity [m/s²] - const v_tend::Vector{NF} = zeros(NF, nlev) # meridional velocity [m/s²] - const temp_tend::Vector{NF} = zeros(NF, nlev) # absolute temperature [K/s] - const humid_tend::Vector{NF} = zeros(NF, nlev) # specific humidity [kg/kg/s] + const u_tend::Vector{NF} = zeros(NF, nlayers) # zonal velocity [m/s²] + const v_tend::Vector{NF} = zeros(NF, nlayers) # meridional velocity [m/s²] + const temp_tend::Vector{NF} = zeros(NF, nlayers) # absolute temperature [K/s] + const humid_tend::Vector{NF} = zeros(NF, nlayers) # specific humidity [kg/kg/s] # FLUXES, arrays to be used for various parameterizations, on half levels incl top and bottom - const flux_u_upward::Vector{NF} = zeros(NF, nlev+1) - const flux_u_downward::Vector{NF} = zeros(NF, nlev+1) + const flux_u_upward::Vector{NF} = zeros(NF, nlayers+1) + const flux_u_downward::Vector{NF} = zeros(NF, nlayers+1) - const flux_v_upward::Vector{NF} = zeros(NF, nlev+1) - const flux_v_downward::Vector{NF} = zeros(NF, nlev+1) + const flux_v_upward::Vector{NF} = zeros(NF, nlayers+1) + const flux_v_downward::Vector{NF} = zeros(NF, nlayers+1) - const flux_temp_upward::Vector{NF} = zeros(NF, nlev+1) - const flux_temp_downward::Vector{NF} = zeros(NF, nlev+1) + const flux_temp_upward::Vector{NF} = zeros(NF, nlayers+1) + const flux_temp_downward::Vector{NF} = zeros(NF, nlayers+1) - const flux_humid_upward::Vector{NF} = zeros(NF, nlev+1) - const flux_humid_downward::Vector{NF} = zeros(NF, nlev+1) + const flux_humid_upward::Vector{NF} = zeros(NF, nlayers+1) + const flux_humid_downward::Vector{NF} = zeros(NF, nlayers+1) # boundary layer boundary_layer_depth::Int = 0 @@ -64,13 +64,13 @@ Base.@kwdef mutable struct ColumnVariables{NF<:AbstractFloat} <: AbstractColumnV # THERMODYNAMICS surface_air_density::NF = 0 - const sat_humid::Vector{NF} = zeros(NF, nlev) # Saturation specific humidity [kg/kg] - const dry_static_energy::Vector{NF} = zeros(NF, nlev) # Dry static energy [J/kg] - const temp_virt::Vector{NF} = zeros(NF, nlev) # virtual temperature [K] - const geopot::Vector{NF} = zeros(NF, nlev) # gepotential height [m] + const sat_humid::Vector{NF} = zeros(NF, nlayers) # Saturation specific humidity [kg/kg] + const dry_static_energy::Vector{NF} = zeros(NF, nlayers) # Dry static energy [J/kg] + const temp_virt::Vector{NF} = zeros(NF, nlayers) # virtual temperature [K] + const geopot::Vector{NF} = zeros(NF, nlayers) # gepotential height [m] # CONVECTION AND PRECIPITATION - cloud_top::Int = nlev + 1 # layer index k of top-most layer with clouds + cloud_top::Int = nlayers+1 # layer index k of top-most layer with clouds precip_convection::NF = 0 # Precipitation due to convection [m] precip_large_scale::NF = 0 # precipitation due to large-scale condensation [m] @@ -79,8 +79,8 @@ Base.@kwdef mutable struct ColumnVariables{NF<:AbstractFloat} <: AbstractColumnV albedo::NF = 0 # surface albedo # WORK ARRAYS - const a::Vector{NF} = zeros(NF, nlev) - const b::Vector{NF} = zeros(NF, nlev) - const c::Vector{NF} = zeros(NF, nlev) - const d::Vector{NF} = zeros(NF, nlev) + const a::Vector{NF} = zeros(NF, nlayers) + const b::Vector{NF} = zeros(NF, nlayers) + const c::Vector{NF} = zeros(NF, nlayers) + const d::Vector{NF} = zeros(NF, nlayers) end \ No newline at end of file diff --git a/src/physics/land.jl b/src/physics/land.jl index ff8e59ca4..21074134b 100644 --- a/src/physics/land.jl +++ b/src/physics/land.jl @@ -8,7 +8,7 @@ function Base.show(io::IO, L::AbstractLand) end export SeasonalLandTemperature -Base.@kwdef struct SeasonalLandTemperature{NF, Grid<:AbstractGrid{NF}} <: AbstractLand{NF, Grid} +@kwdef struct SeasonalLandTemperature{NF, Grid<:AbstractGrid{NF}} <: AbstractLand{NF, Grid} "number of latitudes on one hemisphere, Equator included" nlat_half::Int @@ -47,44 +47,51 @@ function initialize!(land::SeasonalLandTemperature, model::PrimitiveEquation) load_monthly_climatology!(land.monthly_temperature, land) end -function initialize!( land::PrognosticVariablesLand, - time::DateTime, +function initialize!( land::PrognosticVariablesLand, # just for dispatch + progn::PrognosticVariables, + diagn::DiagnosticVariables, model::PrimitiveEquation) - land_timestep!(land, time, model.land, initialize=true) - model isa PrimitiveWet && soil_timestep!(land, time, model.soil) + land_timestep!(progn, diagn, model.land, model, initialize=true) + model isa PrimitiveWet && soil_timestep!(progn, diagn, model.soil, model) end # function barrier -function land_timestep!(land::PrognosticVariablesLand, - time::DateTime, - model::PrimitiveEquation; - initialize::Bool = false) +function land_timestep!( + progn::PrognosticVariables, + diagn::DiagnosticVariables, + model::PrimitiveEquation; + initialize::Bool = false, +) # the time step is dictated by the land "model" - executed = land_timestep!(land, time, model.land; initialize) - executed && model isa PrimitiveWet && soil_timestep!(land, time, model.soil) + executed = land_timestep!(progn, diagn, model.land, model; initialize) + executed && model isa PrimitiveWet && soil_timestep!(progn, diagn, model.soil, model) end -function land_timestep!(land::PrognosticVariablesLand{NF}, - time::DateTime, - land_model::SeasonalLandTemperature; - initialize::Bool = false) where NF - - # escape immediately if Δt of land model hasn't passed yet - # unless the land hasn't been initialized yet - initialize || (time - land.time) < land_model.Δt && return false # = executed - - # otherwise update land prognostic variables: - land.time = time +function land_timestep!( + progn::PrognosticVariables, + diagn::DiagnosticVariables, + land_model::SeasonalLandTemperature, + model::PrimitiveEquation; + initialize::Bool = false +) + (; time) = progn.clock + # # escape immediately if Δt of land model hasn't passed yet + # # unless the land hasn't been initialized yet + # initialize || (time - land.time) < land_model.Δt && return false # = executed + + # # otherwise update land prognostic variables: + # land.time = time this_month = Dates.month(time) next_month = (this_month % 12) + 1 # mod for dec 12 -> jan 1 # linear interpolation weight between the two months # TODO check whether this shifts the climatology by 1/2 a month + NF = eltype(progn.land.land_surface_temperature) weight = convert(NF, Dates.days(time-Dates.firstdayofmonth(time))/Dates.daysinmonth(time)) (; monthly_temperature) = land_model - @. land.land_surface_temperature = (1-weight) * monthly_temperature[this_month] + - weight * monthly_temperature[next_month] + @. progn.land.land_surface_temperature = (1-weight) * monthly_temperature[this_month] + + weight * monthly_temperature[next_month] return true # = executed end @@ -99,23 +106,23 @@ function Base.show(io::IO, S::AbstractSoil) end export SeasonalSoilMoisture -Base.@kwdef struct SeasonalSoilMoisture{NF, Grid<:AbstractGrid{NF}} <: AbstractSoil{NF, Grid} +@kwdef struct SeasonalSoilMoisture{NF, Grid<:AbstractGrid{NF}} <: AbstractSoil{NF, Grid} "number of latitudes on one hemisphere, Equator included" nlat_half::Int # OPTIONS "Depth of top soil layer [m]" - D_top::Float64 = 0.07 + D_top::NF = 0.07 "Depth of root layer [m]" - D_root::Float64 = 0.21 + D_root::NF = 0.21 "Soil wetness at field capacity [volume fraction]" - W_cap::Float64 = 0.3 + W_cap::NF = 0.3 "Soil wetness at wilting point [volume fraction]" - W_wilt::Float64 = 0.17 + W_wilt::NF = 0.17 # READ CLIMATOLOGY FROM FILE "path to the folder containing the soil moisture file, pkg path default" @@ -153,21 +160,26 @@ function initialize!(soil::SeasonalSoilMoisture, model::PrimitiveWet) load_monthly_climatology!(soil.monthly_soil_moisture_layer2, soil, varname=soil.varname_layer2) end -function soil_timestep!(land::PrognosticVariablesLand{NF}, - time::DateTime, - soil_model::SeasonalSoilMoisture) where NF +function soil_timestep!( + progn::PrognosticVariables, + diagn::DiagnosticVariables, + soil_model::SeasonalSoilMoisture, + model::PrimitiveEquation, +) + (; time) = progn.clock this_month = Dates.month(time) next_month = (this_month % 12) + 1 # mod for dec 12 -> jan 1 # linear interpolation weight between the two months # TODO check whether this shifts the climatology by 1/2 a month + NF = eltype(progn.land.soil_moisture_layer1) weight = convert(NF, Dates.days(time-Dates.firstdayofmonth(time))/Dates.daysinmonth(time)) (; monthly_soil_moisture_layer1, monthly_soil_moisture_layer2) = soil_model - @. land.soil_moisture_layer1 = (1-weight) * monthly_soil_moisture_layer1[this_month] + + @. progn.land.soil_moisture_layer1 = (1-weight) * monthly_soil_moisture_layer1[this_month] + weight * monthly_soil_moisture_layer1[next_month] - @. land.soil_moisture_layer2 = (1-weight) * monthly_soil_moisture_layer2[this_month] + + @. progn.land.soil_moisture_layer2 = (1-weight) * monthly_soil_moisture_layer2[this_month] + weight * monthly_soil_moisture_layer2[next_month] return nothing @@ -177,14 +189,14 @@ end abstract type AbstractVegetation{NF, Grid} <: AbstractParameterization end export VegetationClimatology -Base.@kwdef struct VegetationClimatology{NF, Grid<:AbstractGrid{NF}} <: AbstractVegetation{NF, Grid} +@kwdef struct VegetationClimatology{NF, Grid<:AbstractGrid{NF}} <: AbstractVegetation{NF, Grid} "number of latitudes on one hemisphere, Equator included" nlat_half::Int # OPTIONS "Combine high and low vegetation factor, a in high + a*low [1]" - low_veg_factor::Float64 = 0.8 + low_veg_factor::NF = 0.8 "path to the folder containing the soil moisture file, pkg path default" path::String = "SpeedyWeather.jl/input_data" @@ -240,34 +252,29 @@ end # function barrier function soil_moisture_availability!( - diagn::SurfaceVariables, - land::PrognosticVariablesLand, + diagn::DiagnosticVariables, + progn::PrognosticVariables, model::PrimitiveWet) - soil_moisture_availability!(diagn, land, model.soil, model.vegetation) + soil_moisture_availability!(diagn, progn, model.soil, model.vegetation) end function soil_moisture_availability!( - ::SurfaceVariables, - ::PrognosticVariablesLand, + ::DiagnosticVariables, + ::PrognosticVariables, ::PrimitiveDry) return nothing end function soil_moisture_availability!( - diagn::SurfaceVariables{NF}, - land::PrognosticVariablesLand, + diagn::DiagnosticVariables, + progn::PrognosticVariables, soil::AbstractSoil, - vegetation::AbstractVegetation) where NF - - (; soil_moisture_availability) = diagn - (; soil_moisture_layer1, soil_moisture_layer2) = land - (; high_cover, low_cover) = vegetation - - D_top = convert(NF, soil.D_top) - D_root = convert(NF, soil.D_root) - W_cap = convert(NF, soil.W_cap) - W_wilt = convert(NF, soil.W_wilt) - low_veg_factor = convert(NF, vegetation.low_veg_factor) + vegetation::AbstractVegetation, +) + (; soil_moisture_availability) = diagn.physics + (; soil_moisture_layer1, soil_moisture_layer2) = progn.land + (; high_cover, low_cover, low_veg_factor) = vegetation + (; D_top, D_root, W_cap, W_wilt) = soil # precalculate r = 1/(D_top*W_cap + D_root*(W_cap - W_wilt)) diff --git a/src/physics/longwave_radiation.jl b/src/physics/longwave_radiation.jl index 988f44e3f..9ece7a0ab 100644 --- a/src/physics/longwave_radiation.jl +++ b/src/physics/longwave_radiation.jl @@ -25,7 +25,7 @@ relaxation term with `time_scale_stratosphere` towards `temp_stratosphere` is ap Fields are $(TYPEDFIELDS)""" -Base.@kwdef struct UniformCooling{NF} <: AbstractLongwave +@kwdef struct UniformCooling{NF} <: AbstractLongwave "[OPTION] time scale of cooling, default = -1.5K/day = -1K/16hrs" time_scale::Second = Hour(16) @@ -83,7 +83,7 @@ layer towards the tropopause temperature `T_t` with time scale `τ = 24h` (Seeley and Wordsworth, 2023 use 6h, which is unstable a low resolutions here). Fields are $(TYPEDFIELDS)""" -Base.@kwdef struct JeevanjeeRadiation{NF} <: AbstractLongwave +@kwdef struct JeevanjeeRadiation{NF} <: AbstractLongwave "Radiative forcing constant (W/m²/K²)" α::NF = 0.025 @@ -111,7 +111,7 @@ function longwave_radiation!( scheme::JeevanjeeRadiation, ) where NF - (; nlev, temp_tend) = column + (; nlayers, temp_tend) = column T = column.temp # to match Seeley, 2023 notation F = column.flux_temp_upward (; α, time_scale) = scheme @@ -119,8 +119,8 @@ function longwave_radiation!( Fₖ::NF = 0 # flux into lowermost layer from below - # integrate from surface up, skipping surface (k=nlev+1) and top-of-atmosphere flux (k=1) - @inbounds for k in nlev:-1:2 + # integrate from surface up, skipping surface (k=nlayers+1) and top-of-atmosphere flux (k=1) + @inbounds for k in nlayers:-1:2 # Seeley and Wordsworth, 2023 eq (1) Fₖ += (T[k-1] - T[k]) * α * (Tₜ - T[k]) # upward flux from layer k into k-1 F[k] += Fₖ # accumulate fluxes diff --git a/src/physics/ocean.jl b/src/physics/ocean.jl index 31f76d7b5..1bfa674c3 100644 --- a/src/physics/ocean.jl +++ b/src/physics/ocean.jl @@ -57,15 +57,14 @@ function initialize!( ocean::PrognosticVariablesOcean, end # function barrier for all oceans -function ocean_timestep!( ocean::PrognosticVariablesOcean, - time::DateTime, +function ocean_timestep!( progn::PrognosticVariables, + diagn::DiagnosticVariables, model::PrimitiveEquation) - ocean_timestep!(ocean, time, model.ocean) + ocean_timestep!(progn, diagn, model.ocean, model) end ## SEASONAL OCEAN CLIMATOLOGY - export SeasonalOceanClimatology """ @@ -74,7 +73,7 @@ fields from file, and interpolates them in time regularly (default every 3 days) to be stored in the prognostic variables. Fields and options are $(TYPEDFIELDS)""" -Base.@kwdef struct SeasonalOceanClimatology{NF, Grid<:AbstractGrid{NF}} <: AbstractOcean +@kwdef struct SeasonalOceanClimatology{NF, Grid<:AbstractGrid{NF}} <: AbstractOcean "number of latitudes on one hemisphere, Equator included" nlat_half::Int @@ -155,24 +154,25 @@ function initialize!( ocean_model::SeasonalOceanClimatology, model::PrimitiveEquation, ) - ocean.time = time # set initial time + # ocean.time = time # set initial time interpolate_monthly!( ocean.sea_surface_temperature, ocean_model.monthly_temperature, time) end -function ocean_timestep!( ocean::PrognosticVariablesOcean, - time::DateTime, - ocean_model::SeasonalOceanClimatology) +function ocean_timestep!( progn::PrognosticVariables, + diagn::DiagnosticVariables, + ocean_model::SeasonalOceanClimatology, + model::PrimitiveEquation) - # escape immediately if Δt of ocean model hasn't passed yet - (time - ocean.time) < ocean_model.Δt && return nothing + # # escape immediately if Δt of ocean model hasn't passed yet + # (time - ocean.time) < ocean_model.Δt && return nothing - # otherwise update ocean prognostic variables: - ocean.time = time - interpolate_monthly!( ocean.sea_surface_temperature, + # # otherwise update ocean prognostic variables: + # ocean.time = time + interpolate_monthly!( progn.ocean.sea_surface_temperature, ocean_model.monthly_temperature, - time) + progn.clock.time) return nothing end @@ -210,7 +210,7 @@ To be created like and the ocean time is set with initialize!(model, time=time). Fields and options are $(TYPEDFIELDS)""" -Base.@kwdef struct ConstantOceanClimatology <: AbstractOcean +@kwdef struct ConstantOceanClimatology <: AbstractOcean "[OPTION] path to the folder containing the land-sea mask file, pkg path default" path::String = "SpeedyWeather.jl/input_data" @@ -250,9 +250,10 @@ function initialize!( end function ocean_timestep!( - ocean::PrognosticVariablesOcean, - time::DateTime, - ocean_model::ConstantOceanClimatology + progn::PrognosticVariables, + diagn::DiagnosticVariables, + ocean_model::ConstantOceanClimatology, + model::PrimitiveEquation, ) return nothing end @@ -268,7 +269,7 @@ but vary in latitude following a coslat². To be created like Fields and options are $(TYPEDFIELDS)""" -Base.@kwdef struct AquaPlanet{NF} <: AbstractOcean +@kwdef struct AquaPlanet{NF} <: AbstractOcean "Number of latitude rings" nlat::Int @@ -311,9 +312,10 @@ function initialize!( end function ocean_timestep!( - ocean::PrognosticVariablesOcean, - time::DateTime, + progn::PrognosticVariables, + diagn::DiagnosticVariables, ocean_model::AquaPlanet, + model::PrimitiveEquation, ) return nothing end \ No newline at end of file diff --git a/src/physics/surface_fluxes.jl b/src/physics/surface_fluxes.jl index 5809013c6..d487f064c 100644 --- a/src/physics/surface_fluxes.jl +++ b/src/physics/surface_fluxes.jl @@ -34,7 +34,7 @@ function surface_thermodynamics!( column::ColumnVariables, # surface air density via virtual temperature (; R_dry) = model.atmosphere - Tᵥ = column.temp_virt[column.nlev] + Tᵥ = column.temp_virt[column.nlayers] column.surface_air_density = column.pres[end]/(R_dry*Tᵥ) end diff --git a/src/physics/temperature_relaxation.jl b/src/physics/temperature_relaxation.jl index b3eb6443f..6d92b71e3 100644 --- a/src/physics/temperature_relaxation.jl +++ b/src/physics/temperature_relaxation.jl @@ -22,7 +22,7 @@ Base.@kwdef struct HeldSuarez{NF<:AbstractFloat} <: AbstractTemperatureRelaxatio nlat::Int "number of vertical levels" - nlev::Int + nlayers::Int # OPTIONS "sigma coordinate below which faster surface relaxation is applied" @@ -50,7 +50,7 @@ Base.@kwdef struct HeldSuarez{NF<:AbstractFloat} <: AbstractTemperatureRelaxatio κ::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_relax_freq::Matrix{NF} = zeros(NF, nlayers, 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 @@ -59,8 +59,8 @@ end $(TYPEDSIGNATURES) create a HeldSuarez temperature relaxation with arrays allocated given `spectral_grid`""" function HeldSuarez(SG::SpectralGrid; kwargs...) - (; NF, nlat, nlev) = SG - return HeldSuarez{NF}(; nlev, nlat, kwargs...) + (; NF, nlat, nlayers) = SG + return HeldSuarez{NF}(; nlayers, nlat, kwargs...) end """$(TYPEDSIGNATURES) @@ -136,7 +136,7 @@ Base.@kwdef mutable struct JablonowskiRelaxation{NF<:AbstractFloat} <: AbstractT # DIMENSIONS nlat::Int - nlev::Int + nlayers::Int # OPTIONS "sigma coordinate below which relax_time_fast is applied [1]" @@ -164,16 +164,16 @@ Base.@kwdef mutable struct JablonowskiRelaxation{NF<:AbstractFloat} <: AbstractT relax_time_fast::Second = Day(4) # 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 + temp_relax_freq::Matrix{NF} = zeros(NF, nlayers, nlat) # (inverse) relax time scale per layer and lat + temp_equil::Matrix{NF} = zeros(NF, nlayers, nlat) # terms to calc equilibrium temperature as func end """ $(TYPEDSIGNATURES) create a JablonowskiRelaxation temperature relaxation with arrays allocated given `spectral_grid`""" function JablonowskiRelaxation(SG::SpectralGrid; kwargs...) - (; NF, nlat, nlev) = SG - return JablonowskiRelaxation{NF}(; nlev, nlat, kwargs...) + (; NF, nlat, nlayers) = SG + return JablonowskiRelaxation{NF}(; nlayers, nlat, kwargs...) end """$(TYPEDSIGNATURES) diff --git a/src/physics/tendencies.jl b/src/physics/tendencies.jl index 308819e61..f025418ea 100644 --- a/src/physics/tendencies.jl +++ b/src/physics/tendencies.jl @@ -18,7 +18,7 @@ function parameterization_tendencies!( G = model.geometry rings = eachring(G.Grid, G.nlat_half) - @floop for ij in eachgridpoint(diagn) # loop over all horizontal grid points + for ij in eachgridpoint(diagn) # loop over all horizontal grid points thread_id = Threads.threadid() # not two threads should use the same ColumnVariables column = diagn.columns[thread_id] @@ -41,9 +41,8 @@ Calls for `column` one physics parameterization after another and convert fluxes to tendencies.""" function parameterization_tendencies!( column::ColumnVariables, - model::PrimitiveEquation - ) - + model::PrimitiveEquation, +) get_thermodynamics!(column, model) temperature_relaxation!(column, model) boundary_layer_drag!(column, model) @@ -69,8 +68,8 @@ function fluxes_to_tendencies!( atmosphere::AbstractAtmosphere, ) - (; nlev, u_tend, flux_u_upward, flux_u_downward) = column - (; v_tend, flux_v_upward, flux_v_downward) = column + (; 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 @@ -78,12 +77,12 @@ function fluxes_to_tendencies!( pₛ = column.pres[end] # surface pressure (; radius) = geometry # used for scaling - # for g/Δp and g/(Δp*cₚ), see Fortran SPEEDY documentation eq. (3, 5) + # for g/Δp and g/(Δp*c_p), see Fortran SPEEDY documentation eq. (3, 5) g_pₛ = planet.gravity/pₛ cₚ = atmosphere.heat_capacity - # fluxes are defined on half levels including top k=1/2 and surface k=nlev+1/2 - @inbounds for k in 1:nlev + # fluxes are defined on half levels including top k=1/2 and surface k=nlayers+1/2 + @inbounds for k in eachindex(u_tend, v_tend, humid_tend, temp_tend) # Absorbed flux in a given layer, i.e. flux in minus flux out from above and below # Fortran SPEEDY documentation eq. (2) diff --git a/src/physics/vertical_diffusion.jl b/src/physics/vertical_diffusion.jl index ecb2e7aee..2c36c9284 100644 --- a/src/physics/vertical_diffusion.jl +++ b/src/physics/vertical_diffusion.jl @@ -17,7 +17,7 @@ vertical_diffusion!(::ColumnVariables, ::NoVerticalDiffusion, ::PrimitiveEquatio export BulkRichardsonDiffusion @kwdef struct BulkRichardsonDiffusion{NF} <: AbstractVerticalDiffusion - nlev::Int + nlayers::Int "[OPTION] von Kármán constant [1]" κ::NF = 0.4 @@ -45,16 +45,16 @@ export BulkRichardsonDiffusion sqrtC_max::Base.RefValue{NF} = Ref(zero(NF)) # precomputed operators - ∇²_above::Vector{NF} = zeros(NF, nlev) - ∇²_below::Vector{NF} = zeros(NF, nlev) + ∇²_above::Vector{NF} = zeros(NF, nlayers) + ∇²_below::Vector{NF} = zeros(NF, nlayers) end -BulkRichardsonDiffusion(SG::SpectralGrid; kwargs...) = BulkRichardsonDiffusion{SG.NF}(; nlev=SG.nlev, kwargs...) +BulkRichardsonDiffusion(SG::SpectralGrid; kwargs...) = BulkRichardsonDiffusion{SG.NF}(; nlayers=SG.nlayers, kwargs...) function initialize!(scheme::BulkRichardsonDiffusion, model::PrimitiveEquation) - (; nlev) = model.geometry - nlev == 1 && return nothing # no diffusion for 1-layer model + (; nlayers) = model.geometry + nlayers == 1 && return nothing # no diffusion for 1-layer model # ∇² operator on σ levels like 1/Δσ² but for variable Δσ # also includes a 1/2 so that the diffusion coefficients on full levels can be added @@ -63,9 +63,9 @@ function initialize!(scheme::BulkRichardsonDiffusion, model::PrimitiveEquation) σ = model.geometry.σ_levels_full σ_half = model.geometry.σ_levels_half - for k in 1:nlev + for k in 1:nlayers σ₋ = k <= 1 ? -Inf : σ[k-1] # sets the gradient across surface and top to 0 - σ₊ = k >= nlev ? Inf : σ[k+1] # = no flux boundary conditions + σ₊ = k >= nlayers ? Inf : σ[k+1] # = no flux boundary conditions scheme.∇²_above[k] = inv(2*(σ[k] - σ₋) * (σ_half[k+1] - σ_half[k])) scheme.∇²_below[k] = inv(2*(σ₊ - σ[k]) * (σ_half[k+1] - σ_half[k])) end @@ -118,31 +118,31 @@ function get_diffusion_coefficients!( K = column.b # reuse work array for diffusion coefficients (; Ri_c, κ, z₀, fb) = scheme logZ_z₀ = scheme.logZ_z₀[] - (; nlev, u, v, geopot, orography) = column + (; nlayers, u, v, geopot, orography) = column gravity⁻¹ = inv(planet.gravity) # Boundary layer depth is highest layer for which Ri < Ri_c (the "critical" threshold) # as well as all layers below Ri = bulk_richardson!(column, atmosphere) - kₕ::Int = nlev + kₕ::Int = nlayers while kₕ > 0 && Ri[kₕ] < Ri_c kₕ -= 1 end kₕ += 1 # uppermost layer where Ri < Ri_c column.boundary_layer_depth = kₕ - if kₕ <= nlev # boundary layer depth is at least 1 layer thick (calculate diffusion) + if kₕ <= nlayers # boundary layer depth is at least 1 layer thick (calculate diffusion) # Calculate diffusion coefficients following Frierson 2006, eq. 16-20 h = max(geopot[kₕ]*gravity⁻¹ - orography, 0)# always positive to avoid error in log - Ri_N = Ri[nlev] # surface bulk Richardson number + Ri_N = Ri[nlayers] # surface bulk Richardson number Ri_N = clamp(Ri_N, 0, Ri_c) # cases of eq. 12-14 sqrtC = scheme.sqrtC_max[]*(1-Ri_N/Ri_c) # sqrt of eq. 12-14 - surface_speed = sqrt(u[nlev]^2 + v[nlev]^2) + surface_speed = sqrt(u[nlayers]^2 + v[nlayers]^2) K0 = κ * surface_speed * sqrtC # height-independent K eq. 19, 20 K[1:kₕ-1] .= 0 # diffusion above boundary layer 0 - for k in kₕ:nlev + for k in kₕ:nlayers z = max(geopot[k]*gravity⁻¹ - orography, z₀) # height [m] above surface zmin = min(z, fb*h) # height [m] to evaluate Kb(z) at K_k = K0 * zmin # = κ*u_N*√Cz in eq. (19, 20) @@ -187,17 +187,17 @@ function vertical_diffusion!( scheme::BulkRichardsonDiffusion, ) (; ∇²_above, ∇²_below) = scheme - nlev = length(tend) - nlev == 1 && return nothing # escape immediately for single-layer (no diffusion) + nlayers = length(tend) + nlayers == 1 && return nothing # escape immediately for single-layer (no diffusion) - @boundscheck nlev == length(var) == length(K) || throw(BoundsError) - @boundscheck nlev == length(∇²_above) == length(∇²_below) || throw(BoundsError) + @boundscheck nlayers == length(var) == length(K) || throw(BoundsError) + @boundscheck nlayers == length(∇²_above) == length(∇²_below) || throw(BoundsError) - @inbounds for k in kₕ:nlev # diffusion only in surface boundary layer of thickness h + @inbounds for k in kₕ:nlayers # diffusion only in surface boundary layer of thickness h # sets the gradient across surface and top to 0 = no flux boundary conditions - k₋ = max(k, 1) # index above (- in σ direction which is 0 at top and 1 at surface) - k₊ = min(k, nlev) # index below (+ in σ direction which is 0 at top and 1 at surface) + k₋ = max(k, 1) # index above (- in σ direction which is 0 at top and 1 at surface) + k₊ = min(k, nlayers) # index below (+ in σ direction which is 0 at top and 1 at surface) K_∂var_below = (var[k₊] - var[k]) * (K[k₊] + K[k]) # average diffusion coefficient K here K_∂var_above = (var[k] - var[k₋]) * (K[k] + K[k₋]) # but 1/2 is already baked into the ∇² operators @@ -215,16 +215,17 @@ function bulk_richardson!( atmosphere::AbstractAtmosphere, ) cₚ = atmosphere.heat_capacity - (; u, v, geopot, temp_virt, nlev) = column - bulk_richardson = column.d # reuse work array + (; u, v, geopot, temp_virt, nlayers) = column + surface = column.nlayers # surface index = nlayers + bulk_richardson = column.a # reuse work array # surface layer - V² = u[nlev]^2 + v[nlev]^2 - Θ₀ = cₚ*temp_virt[nlev] - Θ₁ = Θ₀ + geopot[nlev] - bulk_richardson[nlev] = geopot[nlev]*(Θ₁ - Θ₀) / (Θ₀*V²) + V² = u[surface]^2 + v[surface]^2 + Θ₀ = cₚ*temp_virt[surface] + Θ₁ = Θ₀ + geopot[surface] + bulk_richardson[surface] = geopot[surface]*(Θ₁ - Θ₀) / (Θ₀*V²) - @inbounds for k in 1:nlev-1 + @inbounds for k in 1:nlayers-1 V² = u[k]^2 + v[k]^2 virtual_dry_static_energy = cₚ * temp_virt[k] + geopot[k] bulk_richardson[k] = geopot[k]*(virtual_dry_static_energy - Θ₁) / (Θ₁*V²) diff --git a/src/physics/zenith.jl b/src/physics/zenith.jl index 443213012..54d0299a8 100644 --- a/src/physics/zenith.jl +++ b/src/physics/zenith.jl @@ -126,7 +126,7 @@ end # function barrier function cos_zenith!(diagn::DiagnosticVariables, time::DateTime, model::PrimitiveEquation) (; solar_zenith, geometry) = model - (; cos_zenith) = diagn.surface + (; cos_zenith) = diagn.physics cos_zenith!(cos_zenith, solar_zenith, time, geometry) end @@ -157,7 +157,7 @@ end function initialize!( S::AbstractZenith, initial_time::DateTime, - model::ModelSetup + model::AbstractModel ) S.initial_time[] = initial_time # to fix the season if no seasonal cycle end diff --git a/src/utility_functions.jl b/src/utility_functions.jl index 1428ad05d..da4c4ca69 100644 --- a/src/utility_functions.jl +++ b/src/utility_functions.jl @@ -87,7 +87,7 @@ function print_fields(io::IO, A, keys;arrays::Bool=false) ~last ? println(io, "├ $key::$(typeof(val)) = $val") : print(io, "└ $key::$(typeof(val)) = $val") end - if filtered # add the names of arrays + if filtered # add the names of arrays s = "└── arrays: " for key in keys if ~(key in keys_filtered) @@ -105,6 +105,7 @@ Dates.Second(x::AbstractFloat) = convert(Second, x) Dates.Minute(x::AbstractFloat) = Second(60x) Dates.Hour( x::AbstractFloat) = Minute(60x) Dates.Day( x::AbstractFloat) = Hour(24x) +Dates.Week( x::AbstractFloat) = Day(7x) # use Dates.second to round to integer seconds Dates.second(x::Dates.Nanosecond) = round(Int, x.value*1e-9) diff --git a/test/callbacks.jl b/test/callbacks.jl index d1afe4b1d..0e6e4343d 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -33,14 +33,15 @@ end callback::StormChaser, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) # allocate recorder: number of time steps (incl initial conditions) in simulation callback.maximum_surface_wind_speed = zeros(progn.clock.n_timesteps + 1) # where surface (=lowermost model layer) u, v on the grid are stored - (; u_grid, v_grid) = diagn.layers[diagn.nlev].grid_variables - + u_grid = diagn.grid.u_grid[:, diagn.nlayers] + v_grid = diagn.grid.u_grid[:, diagn.nlayers] + # maximum wind speed of initial conditions callback.maximum_surface_wind_speed[1] = max_2norm(u_grid, v_grid) @@ -61,7 +62,7 @@ end callback::StormChaser, progn::PrognosticVariables, diagn::DiagnosticVariables, - model::ModelSetup, + model::AbstractModel, ) # increase counter @@ -69,8 +70,9 @@ end i = callback.timestep_counter # where surface (=lowermost model layer) u, v on the grid are stored - (; u_grid, v_grid) = diagn.layers[diagn.nlev].grid_variables - + u_grid = diagn.grid.u_grid[:,diagn.nlayers] + v_grid = diagn.grid.u_grid[:,diagn.nlayers] + # maximum wind speed at current time step callback.maximum_surface_wind_speed[i] = max_2norm(u_grid, v_grid) end @@ -85,7 +87,7 @@ end key = :storm_chaser add!(model.callbacks, key => storm_chaser) # with :storm_chaser key add!(model.callbacks, NoCallback()) # add dummy too - add!(model, NoCallback()) # add dummy with ::ModelSetup interface + add!(model, NoCallback()) # add dummy with ::AbstractModel interface simulation = initialize!(model) run!(simulation, period=Day(1)) diff --git a/test/column_variables.jl b/test/column_variables.jl index 8d657e9c7..87d0ffc5d 100644 --- a/test/column_variables.jl +++ b/test/column_variables.jl @@ -1,6 +1,6 @@ @testset "ColumnVariables initialisation" begin @testset for NF in (Float16, Float32, Float64) - column = ColumnVariables{NF}(nlev=8) + column = ColumnVariables{NF}(nlayers=8) @test eltype(column.temp) == NF SpeedyWeather.reset_column!(column) @@ -11,7 +11,7 @@ @test all(column.v_tend .=== zero(NF)) # Convection - @test column.cloud_top === column.nlev+1 + @test column.cloud_top === column.nlayers+1 # Large-scale condensation @test column.precip_large_scale === zero(NF) @@ -22,28 +22,28 @@ end @testset "ColumnVariables initialisation" begin @testset for NF in (Float32, Float64) - nlev = 8 - spectral_grid = SpectralGrid(NF; nlev) + nlayers = 8 + spectral_grid = SpectralGrid(; NF, nlayers) model = PrimitiveDryModel(; spectral_grid) simulation = initialize!(model) diagn = simulation.diagnostic_variables progn = simulation.prognostic_variables - column = ColumnVariables{NF}(; nlev) + column = ColumnVariables{NF}(; nlayers) SpeedyWeather.reset_column!(column) SpeedyWeather.get_column!(column, diagn, progn, 1, model) # set a tendency to something - humid_tend = rand(NF, nlev) + humid_tend = rand(NF, nlayers) column.humid_tend .= humid_tend # copy into diagn SpeedyWeather.write_column_tendencies!(diagn, column, model.planet, 1) # and check that that worked - for (k, layer) in enumerate(diagn.layers) - @test humid_tend[k] == layer.tendencies.humid_tend_grid[1] + for k in eachgrid(diagn.tendencies.humid_tend_grid) + @test humid_tend[k] == diagn.tendencies.humid_tend_grid[1, k] end end end diff --git a/test/convection.jl b/test/convection.jl index 284145f6d..ef663b762 100644 --- a/test/convection.jl +++ b/test/convection.jl @@ -1,6 +1,6 @@ @testset "Convection" begin - spectral_grid = SpectralGrid(trunc=31, nlev=8) + spectral_grid = SpectralGrid(trunc=31, nlayers=8) for Convection in ( NoConvection, SimplifiedBettsMiller, diff --git a/test/dates.jl b/test/dates.jl index 22ce3b5d5..35d580b84 100644 --- a/test/dates.jl +++ b/test/dates.jl @@ -1,5 +1,5 @@ @testset "Sec, min, hrs arguments" begin - SG = SpectralGrid(trunc=42, nlev=1) + SG = SpectralGrid(trunc=42, nlayers=1) L1 = Leapfrog(SG, Δt_at_T31=30) L2 = Leapfrog(SG, Δt_at_T31=Second(30)) @test L1.Δt == L2.Δt @@ -27,7 +27,7 @@ # clock tests c1 = SpeedyWeather.Clock() - SG2 = SpectralGrid(trunc=31, nlev=1) + SG2 = SpectralGrid(trunc=31, nlayers=1) L6 = Leapfrog(SG2, Δt_at_T31=Hour(1), adjust_with_output=false) SpeedyWeather.set_period!(c1, Hour(10)) diff --git a/test/diffusion.jl b/test/diffusion.jl index bedfaacbc..472791800 100644 --- a/test/diffusion.jl +++ b/test/diffusion.jl @@ -1,42 +1,43 @@ @testset "Horizontal diffusion of random" begin for NF in (Float32, Float64) - spectral_grid = SpectralGrid(NF) - m = PrimitiveDryModel(; spectral_grid) - simulation = initialize!(m) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables + spectral_grid = SpectralGrid(; NF) + model = PrimitiveDryModel(; spectral_grid) + simulation = initialize!(model) + progn = simulation.prognostic_variables + diagn = simulation.diagnostic_variables - (; vor) = p.layers[1].timesteps[1] - (; vor_tend) = d.layers[1].tendencies - (; ∇²ⁿ, ∇²ⁿ_implicit) = m.horizontal_diffusion + (; ∇²ⁿ, ∇²ⁿ_implicit) = model.horizontal_diffusion + vor = progn.vor[1] + (; vor_tend) = diagn.tendencies - vor = randn(LowerTriangularArray{eltype(vor)}, size(vor, as=Matrix)...) + vor .= randn(LowerTriangularArray{eltype(vor)}, size(vor, as=Matrix)...) - # ∇²ⁿ, ∇²ⁿ_implicit are nlev-vectors, one array per layer, pick surface - SpeedyWeather.horizontal_diffusion!(vor_tend, vor, ∇²ⁿ[end], ∇²ⁿ_implicit[end]) + # ∇²ⁿ, ∇²ⁿ_implicit are nlayers-vectors, one array per layer, pick surface + SpeedyWeather.horizontal_diffusion!(vor_tend, vor, ∇²ⁿ, ∇²ⁿ_implicit) # diffusion tendency has opposite sign (real/imag respectively) # than prognostic variable to act as a dissipation - (; lmax, mmax) = m.spectral_transform - for m in 1:mmax+1 - for l in max(2, m):lmax - @test -sign(real(vor[l, m])) == sign(real(vor_tend[l, m])) - @test -sign(imag(vor[l, m])) == sign(imag(vor_tend[l, m])) + (; lmax, mmax) = model.spectral_transform + for k in eachmatrix(vor, vor_tend) + for m in 1:mmax+1 + for l in max(2, m):lmax + @test -sign(real(vor[l, m, k])) == sign(real(vor_tend[l, m, k])) + @test -sign(imag(vor[l, m, k])) == sign(imag(vor_tend[l, m, k])) + end end end - vor0 = copy(vor) - vor1 = copy(vor) - SpeedyWeather.leapfrog!(vor0, vor1, vor_tend, m.time_stepping.Δt, 1, m.time_stepping) + vor2 = progn.vor[2] + SpeedyWeather.leapfrog!(vor, vor2, vor_tend, model.time_stepping.Δt, 1, model.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 + @test any(vor .!= vor2) # check that at least some coefficients are different + @test any(vor .== vor2) # check that at least some coefficients are identical # damping should not increase real or imaginary part of variable - for lm in SpeedyWeather.eachharmonic(vor0, vor) - @test abs(real(vor1[lm])) <= abs(real(vor[lm])) - @test abs(imag(vor1[lm])) <= abs(imag(vor[lm])) + for lmk in eachindex(vor, vor2) + @test abs(real(vor2[lmk])) <= abs(real(vor[lmk])) + @test abs(imag(vor2[lmk])) <= abs(imag(vor[lmk])) end end end \ No newline at end of file diff --git a/test/extending.jl b/test/extending.jl index 4caaa9799..a01b65868 100644 --- a/test/extending.jl +++ b/test/extending.jl @@ -1,5 +1,5 @@ @testset "Extending forcing and drag" begin - Base.@kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag + @kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag # DIMENSIONS from SpectralGrid "Spectral resolution as max degree of spherical harmonics" @@ -10,13 +10,13 @@ time_scale::Second = Day(6) "Jet strength [m/s]" - u₀::Float64 = 20 + u₀::NF = 20 "latitude of Gaussian jet [˚N]" - latitude::Float64 = 30 + latitude::NF = 30 "Width of Gaussian jet [˚]" - width::Float64 = 6 + width::NF = 6 # TO BE INITIALISED "Relaxation back to reference vorticity" @@ -28,7 +28,7 @@ end function SpeedyWeather.initialize!( drag::JetDrag, - model::ModelSetup) + model::AbstractModel) (; spectral_grid, geometry) = model (; Grid, NF, nlat_half) = spectral_grid @@ -40,30 +40,34 @@ u[ij] = drag.u₀ * exp(-(lat[ij]-drag.latitude)^2/(2*drag.width^2)) end - û = SpeedyTransforms.spectral(u, one_more_degree=true) + û = SpeedyTransforms.transform(u, model.spectral_transform) v̂ = zero(û) SpeedyTransforms.curl!(drag.ζ₀, û, v̂, model.spectral_transform) return nothing end - function SpeedyWeather.drag!( diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - drag::JetDrag, - time::DateTime, - model::ModelSetup) + function SpeedyWeather.drag!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + drag::JetDrag, + model::AbstractModel, + lf::Integer, + ) - (; vor) = progn + vor = progn.vor[lf] (; vor_tend) = diagn.tendencies (; ζ₀) = drag (; radius) = model.spectral_grid r = radius/drag.time_scale.value - for lm in eachindex(vor, vor_tend, ζ₀) - vor_tend[lm] -= r*(vor[lm] - ζ₀[lm]) + + k = diagn.nlayers # drag only on surface layer + for lm in eachharmonic(vor_tend) + vor_tend[lm, k] -= r*(vor[lm, k] - ζ₀[lm]) end end - Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing + @kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing # DIMENSIONS from SpectralGrid "Spectral resolution as max degree of spherical harmonics" @@ -78,13 +82,13 @@ decorrelation_time::Second = Day(2) "Stirring strength A [1/s²]" - strength::Float64 = 1e-11 + strength::NF = 1e-11 "Stirring latitude [˚N]" - latitude::Float64 = 45 + latitude::NF = 45 "Stirring width [˚]" - width::Float64 = 24 + width::NF = 24 "Minimum degree of spherical harmonics to force" lmin::Int = 8 @@ -118,7 +122,7 @@ end function SpeedyWeather.initialize!( forcing::StochasticStirring, - model::ModelSetup) + model::AbstractModel) # precompute forcing strength, scale with radius^2 as is the vorticity equation (; radius) = model.spectral_grid @@ -142,24 +146,28 @@ return nothing end - function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer, - progn::PrognosticVariablesLayer, - forcing::StochasticStirring, - time::DateTime, - model::ModelSetup) + function SpeedyWeather.forcing!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + forcing::StochasticStirring, + model::AbstractModel, + lf::Integer, + ) SpeedyWeather.forcing!(diagn, forcing, model.spectral_transform) end - function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer, - forcing::StochasticStirring{NF}, - spectral_transform::SpectralTransform) where NF - + function SpeedyWeather.forcing!( + diagn::DiagnosticVariables, + forcing::StochasticStirring{NF}, + spectral_transform::SpectralTransform, + ) where NF # noise and auto-regressive factors a = forcing.a[] # = sqrt(1 - exp(-2dt/τ)) b = forcing.b[] # = exp(-dt/τ) (; S) = forcing lmax, mmax = size(S, as=Matrix) + @inbounds for m in 1:mmax for l in m:lmax if (forcing.mmin <= m <= forcing.mmax) && @@ -172,21 +180,21 @@ end # to grid-point space - S_grid = diagn.dynamics_variables.a_grid - SpeedyTransforms.gridded!(S_grid, S, spectral_transform) + S_grid = diagn.dynamics.a_2D_grid + transform!(S_grid, S, spectral_transform) # mask everything but mid-latitudes RingGrids._scale_lat!(S_grid, forcing.lat_mask) # back to spectral space - (; vor_tend) = diagn.tendencies - SpeedyTransforms.spectral!(vor_tend, S_grid, spectral_transform) - SpeedyTransforms.spectral_truncation!(vor_tend) # set lmax+1 to zero - + S_masked = diagn.dynamics.a_2D + transform!(S_masked, S_grid, spectral_transform) + k = diagn.nlayers # only force surface layer + diagn.tendencies.vor_tend[:, k] .+= S_masked return nothing end - spectral_grid = SpectralGrid(trunc=31, nlev=1) + spectral_grid = SpectralGrid(trunc=31, nlayers=1) drag = JetDrag(spectral_grid, time_scale=Day(6)) forcing = StochasticStirring(spectral_grid) diff --git a/test/geopotential.jl b/test/geopotential.jl index b2b429ade..572467ddc 100644 --- a/test/geopotential.jl +++ b/test/geopotential.jl @@ -1,90 +1,87 @@ @testset "Geopotential reasonable" begin for NF in (Float32, Float64) - nlev = 8 - spectral_grid = SpectralGrid(; NF, nlev, Grid=FullGaussianGrid) + nlayers = 8 + spectral_grid = SpectralGrid(; NF, nlayers, Grid=FullGaussianGrid) model = PrimitiveWetModel(; spectral_grid) simulation = initialize!(model) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables + progn = simulation.prognostic_variables + diagn = simulation.diagnostic_variables + temp = progn.temp[1] + humid = progn.humid[1] - # 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_lf = progn_layer.timesteps[lf] - progn_layer_lf.temp[1] = temp*model.spectral_transform.norm_sphere - fill!(progn_layer_lf.humid, 0) # dry core - 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, model.geopotential, model.orography) + # give every layer some constant temperature by setting the l=m=0 mode (index 1) for all k + temp0 = 280 # in Kelvin + temp[1, :] .= temp0 * model.spectral_transform.norm_sphere + humid .= 0 + lf = 1 # leapfrog time step + SpeedyWeather.transform!(diagn, progn, lf, model) + SpeedyWeather.linear_virtual_temperature!(diagn, progn, lf, model) + SpeedyWeather.geopotential!(diagn, model.geopotential, model.orography) + geopot_grid = Array(transform(diagn.dynamics.geopot, model.spectral_transform)) + # 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]/model.planet.gravity # middle of pacific - @test heights[k] ≈ height_over_ocean rtol=0.5 # very large error allowed + height_over_ocean = geopot_grid[48, 24, :]/model.planet.gravity # middle of pacific + for k in eachmatrix(temp) + @test heights[k] ≈ height_over_ocean[k] rtol=0.5 # very large error allowed end end end @testset "Add geopotential and kinetic energy, compute -∇²B term, no errors" begin for NF in (Float32, Float64) - nlev = 8 - spectral_grid = SpectralGrid(; NF, nlev, Grid=FullGaussianGrid) - m = PrimitiveWetModel(; spectral_grid) - simulation = initialize!(m) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables - - # give every layer some constant temperature - temp = 280 # in Kelvin - for k in 1:nlev - p.layers[k].timesteps[1].temp[1] = temp*m.spectral_transform.norm_sphere - end + spectral_grid = SpectralGrid(; NF, nlayers=8, Grid=FullGaussianGrid) + model = PrimitiveWetModel(; spectral_grid) + simulation = initialize!(model) + progn = simulation.prognostic_variables + diagn = simulation.diagnostic_variables + temp = progn.temp[1] + humid = progn.humid[1] - SpeedyWeather.geopotential!(d, m.geopotential, m.orography) + # give every layer some constant temperature by setting the l=m=0 mode (index 1) for all k + temp0 = 280 # in Kelvin + temp[1, :] .= temp0 * model.spectral_transform.norm_sphere + humid .= 0 - lf = 1 - for (progn_layer, diagn_layer) in zip(p.layers, d.layers) - progn_layer_lf = progn_layer.timesteps[lf] - SpeedyWeather.gridded!(diagn_layer, progn_layer_lf, m) # propagate spectral state to grid - SpeedyWeather.bernoulli_potential!(diagn_layer, m.spectral_transform) - end + lf = 1 # leapfrog time step + SpeedyWeather.transform!(diagn, progn, lf, model) + SpeedyWeather.linear_virtual_temperature!(diagn, progn, lf, model) + SpeedyWeather.geopotential!(diagn, model.geopotential, model.orography) + SpeedyWeather.bernoulli_potential!(diagn, model.spectral_transform) end end @testset "Virtual temperature calculation" begin for NF in (Float32, Float64) - nlev = 8 - spectral_grid = SpectralGrid(; NF, nlev, Grid=FullGaussianGrid) - m = PrimitiveWetModel(; spectral_grid) - simulation = initialize!(m) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables + spectral_grid = SpectralGrid(; NF, nlayers=8, Grid=FullGaussianGrid) + model = PrimitiveWetModel(; spectral_grid) + simulation = initialize!(model) + progn = simulation.prognostic_variables + diagn = simulation.diagnostic_variables + + temp = progn.temp[1] + humid = progn.humid[1] + + # give every layer some constant temperature by setting the l=m=0 mode (index 1) for all k + temp0 = 280 # in Kelvin + temp .= 0 + temp[1, :] .= temp0 * model.spectral_transform.norm_sphere + + humid .= 0 + humid[1, :] .= 1e-2*(rand(spectral_grid.nlayers) .+ 0.1) * model.spectral_transform.norm_sphere - # 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_lf = progn_layer.timesteps[lf] - progn_layer_lf.temp[1] = temp*m.spectral_transform.norm_sphere - fill!(progn_layer_lf.humid, 0) - progn_layer_lf.humid[1] = rand(NF) - SpeedyWeather.gridded!(diagn_layer, progn_layer_lf, m) # propagate spectral state to grid - - temp_lf = progn_layer.timesteps[lf].temp # always passed on for compat with DryCore - SpeedyWeather.virtual_temperature!(diagn_layer, temp_lf, m) + lf = 1 # leapfrog time step + SpeedyWeather.transform!(diagn, progn, lf, model) + SpeedyWeather.virtual_temperature!(diagn, model) - T = diagn_layer.grid_variables.temp_grid - Tv = diagn_layer.grid_variables.temp_virt_grid + T = diagn.grid.temp_grid + Tv = diagn.grid.temp_virt_grid - for ij in SpeedyWeather.eachgridpoint(T, Tv) - @test T[ij] < Tv[ij] # virtual temperature has to be higher as q > 0 - end + for ijk in eachindex(T, Tv) + # virtual temperature has to be higher as q > 0 + @test T[ijk] < Tv[ijk] end end end \ No newline at end of file diff --git a/test/grids.jl b/test/grids.jl index 9fe1289ce..77911cb12 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -64,7 +64,7 @@ end J2 = OctaHEALPixGrid(randn(NF, 4096)) # J32 grid K2 = FullOctaHEALPixGrid(randn(NF, 128*63)) # K32 grid - for (grid1, grid2) in zip([L, F, O, C, H, J, K], [L2, F2, O2, C2, H2, J2, K2]) + for (grid1, grid2) in zip((L, F, O, C, H, J, K), (L2, F2, O2, C2, H2, J2, K2)) @test size(grid1) == size(grid2) end @@ -111,7 +111,7 @@ end n = 4 # resolution parameter nlat_half G1 = zeros(G{NF}, n) G2 = zero(G1) - G3 = G(G2) + G3 = G(G2.data) @test G1 == G2 @test G1 == G3 @@ -279,6 +279,25 @@ end end end +@testset "Ring indices" begin + + g1 = zeros(OctahedralGaussianGrid, 2) + g2 = zeros(OctahedralGaussianGrid, 2, 1) # matches above + g3 = zeros(OctahedralGaussianGrid, 2, 2) # matches horizontally only + g4 = zeros(OctahedralClenshawGrid, 2) # does not match above + + @test eachring(g1) == eachring(g1, g2) == eachring(g1, g2, g2, g1) + @test eachring(g1) == eachring(g2, g3) + @test_throws DimensionMismatch eachring(g1, g4) + @test_throws DimensionMismatch eachring(g2, g4) + @test_throws DimensionMismatch eachring(g3, g4) + + @test RingGrids.grids_match(g1, g3) == false + @test RingGrids.grids_match(g2, g3) == false + @test RingGrids.grids_match(g1, g3, horizontal_only=true) + @test RingGrids.grids_match(g2, g3, horizontal_only=true) +end + @testset "Grid broadcasting" begin n = 2 @testset for G in ( FullClenshawArray, @@ -404,6 +423,14 @@ end end end +# needed when extension is not loaded (manual testing) +# RingGrids.nonparametric_type(::Type{<:JLArray}) = JLArray + +# define for Julia 1.9 +if VERSION < v"1.10" + JLArrays.JLDeviceArray{T, N}(A::JLArrays.JLDeviceArray{T, N}) where {T,N} = A +end + @testset "AbstractGridArray: GPU (JLArrays)" begin NF = Float32 @testset for Grid in ( diff --git a/test/interpolation.jl b/test/interpolation.jl index fb31a24c7..aa4498999 100644 --- a/test/interpolation.jl +++ b/test/interpolation.jl @@ -167,7 +167,7 @@ end trunc = 10 alms = randn(LowerTriangularMatrix{Complex{NF}}, 5, 5) alms = spectral_truncation(alms, trunc+2, trunc+1) - A = gridded(alms; Grid) + A = transform(alms; Grid) # interpolate to FullGaussianGrid and back and compare nlat_half = 32 diff --git a/test/kernelabstractions.jl b/test/kernelabstractions.jl index 633baf8ca..08212171f 100644 --- a/test/kernelabstractions.jl +++ b/test/kernelabstractions.jl @@ -3,7 +3,7 @@ using KernelAbstractions @testset "KernelAbstractions tests" begin # only on CPU (currently) - device_setup = SpeedyWeather.DeviceSetup(SpeedyWeather.CPUDevice()) + device_setup = SpeedyWeather.DeviceSetup(SpeedyWeather.CPU()) @kernel function mul_test!(A, @Const(B), @Const(C)) i, j = @index(Global, NTuple) diff --git a/test/land_sea_mask.jl b/test/land_sea_mask.jl index 62a135a4e..f3a5806c0 100644 --- a/test/land_sea_mask.jl +++ b/test/land_sea_mask.jl @@ -1,9 +1,10 @@ @testset "Land sea masks" begin @testset for Mask in (LandSeaMask, AquaPlanetMask) - spectral_grid = SpectralGrid(trunc=31, nlev=8) + spectral_grid = SpectralGrid(trunc=31, nlayers=8) mask = Mask(spectral_grid) model = PrimitiveWetModel(; spectral_grid, land_sea_mask=mask) simulation = initialize!(model) + model.feedback.verbose = false run!(simulation, period=Day(5)) @test simulation.model.feedback.nars_detected == false end diff --git a/test/lower_triangular_matrix.jl b/test/lower_triangular_matrix.jl index 10f7f6f06..828eeb3f5 100644 --- a/test/lower_triangular_matrix.jl +++ b/test/lower_triangular_matrix.jl @@ -57,7 +57,6 @@ end @test size(L)[2:end] == size(A)[3:end] @test size(L)[1] == SpeedyWeather.LowerTriangularMatrices.nonzeros(size(A,1), size(A,2)) - for m in 1:mmax for l in 1:lmax @test A[l, m, [Colon() for i=1:length(idims)]...] == L[l, m, [Colon() for i=1:length(idims)]...] @@ -290,7 +289,7 @@ end end end -@testset "LowerTriangularArray: fill, copy, randn, convert" begin +@testset "LowerTriangularArray: fill, copy, randn, convert, repeat" begin @testset for NF in (Float32, Float64) mmax = 32 @testset for idims = ((), (5,), (5,5)) @@ -324,6 +323,10 @@ end for lm in SpeedyWeather.eachharmonic(L, L3) @test Float16(L[lm, [1 for i=1:length(idims)]...]) == L3[lm, [1 for i=1:length(idims)]...] end + + L_rep = repeat(L, 1, [2 for i=1:length(idims)]...) + @test typeof(L_rep) <: LowerTriangularArray + @test size(L_rep, as=Vector) == (size(L,1), [10 for i=1:length(idims)]...) # 10 = 2 * 5 = 2*idims[i] end end end @@ -377,9 +380,10 @@ end @test size(similar(L)) == size(L) @test eltype(L) == eltype(similar(L, eltype(L))) - + @test (5, 7) == size(similar(L, 5, 7), as=Matrix) @test (5, 7) == size(similar(L, (5, 7)), as=Matrix) + @test similar(L) isa LowerTriangularMatrix @test similar(L, Float64) isa LowerTriangularMatrix{Float64} end @@ -403,8 +407,9 @@ end @test size(similar(L)) == size(L) @test eltype(L) == eltype(similar(L, eltype(L))) - @test (5, 7, idims...) == size(similar(L, 5, 7, idims...); as=Matrix) - @test (5, 7, idims...) == size(similar(L, (5, 7, idims...)); as=Matrix) + @test (5, 7, idims...) == size(similar(L, 5, 7, idims...), as=Matrix) + @test (5, 7, idims...) == size(similar(L, (5, 7, idims...)), as=Matrix) + @test similar(L) isa LowerTriangularArray @test similar(L, Float64) isa LowerTriangularArray{Float64} end @@ -435,7 +440,7 @@ end # with ranges L1 = zeros(LowerTriangularMatrix{NF}, 33, 32); L2 = randn(LowerTriangularMatrix{NF}, 65, 64); - L2T = spectral_truncation(L2,(size(L1; as=Matrix) .- 1)...) + L2T = spectral_truncation(L2, size(L1, ZeroBased, as=Matrix)...) copyto!(L1, L2, 1:33, 1:32) # size of smaller matrix @test L1 == L2T @@ -498,7 +503,7 @@ end # with ranges L1 = zeros(LowerTriangularArray{NF}, 33, 32, idims...); L2 = randn(LowerTriangularArray{NF}, 65, 64, idims...); - L2T = spectral_truncation(L2,(size(L1; as=Matrix)[1:2] .- 1)...) + L2T = spectral_truncation(L2, (size(L1, ZeroBased, as=Matrix)[1:2])...) copyto!(L1, L2, 1:33, 1:32) # size of smaller matrix @test L1 == L2T @@ -591,6 +596,7 @@ end @test (5, 7, 5) == size(similar(L, 5, 7, idims...), as=Matrix) @test (5, 7, 5) == size(similar(L, (5, 7, idims...)), as=Matrix) + @test similar(L) isa LowerTriangularArray # copyto! same size @@ -607,10 +613,11 @@ end # So, we do this with regular Array but with the _copyto_core! function that implements # the core of this copyto! in a GPU compatible way, and is called by copyto! with CuArrays - L1 = zeros(LowerTriangularArray{NF}, 33, 32, idims...); - L2 = randn(LowerTriangularArray{NF}, 65, 64, idims...); + L1 = zeros(LowerTriangularArray{NF}, 33, 32, idims...) + L2 = randn(LowerTriangularArray{NF}, 65, 64, idims...) + L2T = spectral_truncation(L2, (size(L1, ZeroBased; as=Matrix)[1:2])...) - L3 = zeros(LowerTriangularArray{NF}, 33, 32, idims...); + L3 = zeros(LowerTriangularArray{NF}, 33, 32, idims...) SpeedyWeather.LowerTriangularMatrices._copyto_core!(L1, L2, 1:33, 1:32) # size of smaller matrix @test L1 == L2T @@ -632,7 +639,6 @@ end @test L3 == L1 end - @testset "LowerTriangularArray: broadcast" begin @testset for idims = ((), (5,), (5,5)) @testset for NF in (Float16, Float32, Float64) diff --git a/test/netcdf_output.jl b/test/netcdf_output.jl index ab841f9a3..d68fabcf4 100644 --- a/test/netcdf_output.jl +++ b/test/netcdf_output.jl @@ -1,108 +1,159 @@ using NCDatasets, Dates -@testset "Output on various grids" begin +@testset "Output for BarotropicModel" begin tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits period = Day(1) - # default grid, Float64, ShallowWater - spectral_grid = SpectralGrid(; NF=Float64, nlev=1) - output = OutputWriter(spectral_grid, ShallowWater, path=tmp_output_path) - model = ShallowWaterModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false + for Grid in (FullGaussianGrid, FullClenshawGrid, OctahedralGaussianGrid, OctahedralClenshawGrid, HEALPixGrid, OctaHEALPixGrid) + spectral_grid = SpectralGrid(nlayers=1) + output = NetCDFOutput(spectral_grid, path=tmp_output_path) + model = BarotropicModel(; spectral_grid, output) + simulation = initialize!(model) + run!(simulation, output=true; period) + @test simulation.model.feedback.nars_detected == false + + # read netcdf file and check that all variables exist + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + for key in keys(output.variables) + @test haskey(ds, key) + + # test dimensions + nx, ny, nz, nt = size(ds[key]) + @test (nx, ny) == RingGrids.matrix_size(output.grid2D) + @test nz == spectral_grid.nlayers + @test nt == Int(period / output.output_dt) + 1 + end + end +end - # default grid, Float32, ShallowWater - spectral_grid = SpectralGrid(; NF=Float32, nlev=1) - output = OutputWriter(spectral_grid, ShallowWater, path=tmp_output_path) - model = ShallowWaterModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false +@testset "Output for ShallowWaterModel" begin + tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits + period = Day(1) - # FullClenshawGrid, Float32, ShallowWater - spectral_grid = SpectralGrid(; NF=Float32, Grid=FullClenshawGrid, nlev=1) - output = OutputWriter(spectral_grid, ShallowWater, path=tmp_output_path) - model = ShallowWaterModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false + for output_NF in (Float32, Float64) + spectral_grid = SpectralGrid(nlayers=1) + output = NetCDFOutput(spectral_grid, ShallowWater, path=tmp_output_path; output_NF) + model = ShallowWaterModel(; spectral_grid, output) + simulation = initialize!(model) + run!(simulation, output=true; period) + @test simulation.model.feedback.nars_detected == false + + # read netcdf file and check that all variables exist + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + for key in keys(output.variables) + @test haskey(ds, key) + + # test output number format, NCDatasets returns Union{Missing, Float32} for Float32 so do <: to check + @test output_NF <: eltype(ds[key][:]) + end - # OctahedralClenshawGrid, Float32, ShallowWater - spectral_grid = SpectralGrid(; NF=Float32, Grid=OctahedralClenshawGrid, nlev=1) - output = OutputWriter(spectral_grid, ShallowWater, path=tmp_output_path) - model = ShallowWaterModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false + # add divergence output + div_output = SpeedyWeather.DivergenceOutput() + add!(output.variables, div_output) + run!(simulation, output=true; period) + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + @test haskey(ds, div_output.name) + + # add orography output + orog_output = SpeedyWeather.OrographyOutput() + add!(output.variables, orog_output) + run!(simulation, output=true; period) + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + @test haskey(ds, orog_output.name) + + nx, ny = size(ds[orog_output.name]) + @test (nx, ny) == RingGrids.matrix_size(output.grid2D) + + # delete divergence output + delete!(output, div_output.name) + run!(simulation, output=true; period) + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + @test ~haskey(ds, div_output.name) + end +end - # HEALPixGrid, Float32, ShallowWater - spectral_grid = SpectralGrid(; NF=Float32, Grid=HEALPixGrid, nlev=1) - output = OutputWriter(spectral_grid, ShallowWater, path=tmp_output_path) - model = ShallowWaterModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false +@testset "Output for PrimitiveDryModel" begin + tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits + period = Day(1) - # OctaHEALPixGrid, Float32, ShallowWater - spectral_grid = SpectralGrid(; NF=Float32, Grid=OctaHEALPixGrid, nlev=1) - output = OutputWriter(spectral_grid, ShallowWater, path=tmp_output_path) - model = ShallowWaterModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false + # test also output at various resolutions + for nlat_half in (24, 32, 48, 64) + spectral_grid = SpectralGrid(nlayers=8) + output = NetCDFOutput(spectral_grid, ShallowWater, path=tmp_output_path; nlat_half) + model = PrimitiveDryModel(; spectral_grid, output) + simulation = initialize!(model) + run!(simulation, output=true; period) + @test simulation.model.feedback.nars_detected == false + + # read netcdf file and check that all variables exist + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + for key in keys(output.variables) + @test haskey(ds, key) + end - # OctahedralClenshawGrid, as matrix, Float32, ShallowWater - spectral_grid = SpectralGrid(; NF=Float32, Grid=OctahedralClenshawGrid, nlev=1) - output = OutputWriter(spectral_grid, ShallowWater, path=tmp_output_path, as_matrix=true) - model = ShallowWaterModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false + nx, ny, nz, nt = size(ds[:vor]) + @test nx == 2ny + @test ny == 2nlat_half + end +end - # OctaHEALPixGrid, as matrix, Float32, PrimitiveDry - # using T42 as T31 has stability issues in the first day - spectral_grid = SpectralGrid(; NF=Float32, trunc=42, Grid=OctaHEALPixGrid) - output = OutputWriter(spectral_grid, PrimitiveDry, path=tmp_output_path, as_matrix=true) - model = PrimitiveDryModel(; spectral_grid, output) - simulation = initialize!(model) - run!(simulation, output=true; period) - @test simulation.model.feedback.nars_detected == false +@testset "Output for PrimitiveWetModel" begin + tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits + period = Day(1) - # OctaHEALPixGrid, as matrix, Float32, but output Float64 PrimitiveDry - # using T42 as T31 has stability issues in the first day - spectral_grid = SpectralGrid(; NF=Float32, trunc=42, Grid=OctaHEALPixGrid) - output = OutputWriter(spectral_grid, PrimitiveDry, path=tmp_output_path, as_matrix=true, NF=Float64) - model = PrimitiveDryModel(; spectral_grid, output) + # test also output at various resolutions + spectral_grid = SpectralGrid(nlayers=8) + for output_dt in (Hour(1), Minute(120), Hour(3), Hour(6), Day(1)) + output = NetCDFOutput(spectral_grid, PrimitiveWet, path=tmp_output_path; output_dt) + model = PrimitiveWetModel(; spectral_grid, output) + simulation = initialize!(model) + run!(simulation, output=true; period) + @test simulation.model.feedback.nars_detected == false + + # read netcdf file and check that all variables exist + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + for key in keys(output.variables) + @test haskey(ds, key) + + # test time + nt = size(ds[key])[end] + @test nt == Int(period / output.output_dt) + 1 + end + end + + # test outputting other model defaults + output = NetCDFOutput(spectral_grid, Barotropic, path=tmp_output_path) + model = PrimitiveWetModel(; spectral_grid, output) simulation = initialize!(model) run!(simulation, output=true; period) @test simulation.model.feedback.nars_detected == false + ds = NCDataset(joinpath(model.output.run_path, model.output.filename)) + @test ~haskey(ds, "temp") + @test ~haskey(ds, "humid") + @test ~haskey(ds, "pres") end @testset "Restart from output file" begin tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits spectral_grid = SpectralGrid() - output = OutputWriter(spectral_grid, PrimitiveDry, path=tmp_output_path, id="restart-test") + output = NetCDFOutput(spectral_grid, PrimitiveDry, path=tmp_output_path, id="restart-test") model = PrimitiveDryModel(; spectral_grid, output) simulation = initialize!(model) run!(simulation, output=true; period=Day(1)) initial_conditions = StartFromFile(path=tmp_output_path, id="restart-test") - model2 = PrimitiveDryModel(; spectral_grid, initial_conditions) - simulation2 = initialize!(model2) + model_new = PrimitiveDryModel(; spectral_grid, initial_conditions) + simulation_new = initialize!(model_new) - p1 = simulation.prognostic_variables - p2 = simulation2.prognostic_variables + progn_old = simulation.prognostic_variables + progn_new = simulation_new.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 + for varname in (:vor, :div, :temp, :pres) + var_old = getfield(progn_old, varname)[1] + var_new = getfield(progn_new, varname)[1] + @test all(var_old .== var_new) end - @test all(SpeedyWeather.get_pressure(p1) .== SpeedyWeather.get_pressure(p2)) end @testset "Time axis" begin @@ -118,7 +169,7 @@ end tmp_output_path = mktempdir(pwd(), prefix = "tmp_testruns_") # Cleaned up when the process exits spectral_grid = SpectralGrid() - output = OutputWriter(spectral_grid, PrimitiveDry, path=tmp_output_path, id="dense-output-test", output_dt=Hour(0)) + output = NetCDFOutput(spectral_grid, PrimitiveDry, path=tmp_output_path, id="dense-output-test", output_dt=Hour(0)) model = PrimitiveDryModel(; spectral_grid, output) simulation = initialize!(model) run!(simulation, output=true; period=Day(1)) @@ -129,7 +180,7 @@ end @test t == manual_time_axis(model.output.startdate, model.time_stepping.Δt_millisec, progn.clock.n_timesteps) # do a simulation with the adjust_Δt_with_output turned on - output = OutputWriter(spectral_grid, PrimitiveDry, path=tmp_output_path, id="adjust_dt_with_output-test", output_dt=Minute(70)) + output = NetCDFOutput(spectral_grid, PrimitiveDry, path=tmp_output_path, id="adjust_dt_with_output-test", output_dt=Minute(70)) time_stepping = Leapfrog(spectral_grid, adjust_with_output=true) model = PrimitiveDryModel(; spectral_grid, output, time_stepping) simulation = initialize!(model) @@ -144,7 +195,7 @@ end # 1kyrs simulation spectral_grid = SpectralGrid() time_stepping = Leapfrog(spectral_grid, Δt_at_T31=Day(3650)) - output = OutputWriter(spectral_grid, PrimitiveDry, path=tmp_output_path, id="long-output-test", output_dt=Day(3650)) + output = NetCDFOutput(spectral_grid, PrimitiveDry, path=tmp_output_path, id="long-output-test", output_dt=Day(3650)) model = PrimitiveDryModel(; spectral_grid, output, time_stepping) simulation = initialize!(model) run!(simulation, output=true, period=Day(365000)) diff --git a/test/ocean.jl b/test/ocean.jl index 0dbcb71de..a1279a508 100644 --- a/test/ocean.jl +++ b/test/ocean.jl @@ -1,6 +1,6 @@ @testset "Ocean models" begin - spectral_grid = SpectralGrid(trunc=31, nlev=5) + spectral_grid = SpectralGrid(trunc=31, nlayers=5) for OceanModel in ( SeasonalOceanClimatology, ConstantOceanClimatology, diff --git a/test/orography.jl b/test/orography.jl index d1ff6127a..ef82de34d 100644 --- a/test/orography.jl +++ b/test/orography.jl @@ -1,6 +1,6 @@ @testset "Orographies" begin @testset for Orography in (EarthOrography, ZonalRidge) - spectral_grid = SpectralGrid(trunc=31, nlev=8) + spectral_grid = SpectralGrid(trunc=31, nlayers=8) orography = Orography(spectral_grid) model = PrimitiveWetModel(; spectral_grid, orography) simulation = initialize!(model) diff --git a/test/particle_advection.jl b/test/particle_advection.jl index f3a00f421..ed2e97f2e 100644 --- a/test/particle_advection.jl +++ b/test/particle_advection.jl @@ -1,19 +1,19 @@ @testset "Particle advection" begin - for Model in (BarotropicModel, + for Model in ( BarotropicModel, ShallowWaterModel, PrimitiveDryModel, PrimitiveWetModel) if Model <: PrimitiveEquation - nlev = 8 + nlayers = 8 else - nlev = 1 + nlayers = 1 end - spectral_grid = SpectralGrid(trunc=31, nlev=1, n_particles=100) + spectral_grid = SpectralGrid(trunc=31, nlayers=nlayers, nparticles=100) particle_advection = ParticleAdvection2D(spectral_grid) - model = Model(;spectral_grid,particle_advection) + model = Model(;spectral_grid, particle_advection) add!(model.callbacks, ParticleTracker(spectral_grid)) simulation = initialize!(model) diff --git a/test/radiation.jl b/test/radiation.jl index 87ada46cd..55f63bb54 100644 --- a/test/radiation.jl +++ b/test/radiation.jl @@ -1,6 +1,6 @@ @testset "Longwave radiation" begin - spectral_grid = SpectralGrid(trunc=31, nlev=5) + spectral_grid = SpectralGrid(trunc=31, nlayers=5) for Radiation in (NoLongwave, UniformCooling, diff --git a/test/run_speedy.jl b/test/run_speedy.jl index 590d4f461..b22b2ae12 100644 --- a/test/run_speedy.jl +++ b/test/run_speedy.jl @@ -1,13 +1,13 @@ @testset "run_speedy no errors, no blowup" begin # Barotropic - spectral_grid = SpectralGrid(nlev=1) + spectral_grid = SpectralGrid(nlayers=1) model = BarotropicModel(; spectral_grid) simulation = initialize!(model) run!(simulation, period=Day(10)) @test simulation.model.feedback.nars_detected == false - # ShallowWater - spectral_grid = SpectralGrid(nlev=1) + # ShallowWater + spectral_grid = SpectralGrid(nlayers=1) model = ShallowWaterModel(; spectral_grid) simulation = initialize!(model) run!(simulation, period=Day(10)) diff --git a/test/runtests.jl b/test/runtests.jl index cff6463de..556bf37a0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,7 +7,7 @@ include("dates.jl") include("lower_triangular_matrix.jl") include("grids.jl") include("interpolation.jl") -#include("set_vars.jl") +include("set.jl") # GPU/KERNELABSTRACTIONS include("kernelabstractions.jl") @@ -24,7 +24,7 @@ include("particles.jl") include("particle_advection.jl") # VERTICAL LEVELS -include("vertical_levels.jl") +include("vertical_coordinates.jl") include("geopotential.jl") # PHYSICS diff --git a/test/schedule.jl b/test/schedule.jl index 050a77128..63acac882 100644 --- a/test/schedule.jl +++ b/test/schedule.jl @@ -1,8 +1,8 @@ @testset "Periodic schedule" begin - # especially T170 uses 337500 milliseconds time steps + # especially T170 uses 337500 milliseconds time steps # not representable as seconds for trunc in (31,42,63,85,127,170,255,341) - spectral_grid = SpectralGrid(trunc=trunc, nlev=1) + spectral_grid = SpectralGrid(trunc=trunc, nlayers=1) time_stepping = Leapfrog(spectral_grid) clock = Clock() @@ -14,7 +14,7 @@ schedule = Schedule(every=hour) initialize!(schedule, clock) - # adapted schedule time step should be within 20% + # adapted schedule time step should be within 20% @test schedule.every.value ≈ Second(hour).value rtol=2e-1 @test schedule.steps ≈ period/hour rtol=2e-1 end diff --git a/test/set.jl b/test/set.jl new file mode 100644 index 000000000..8830fb12c --- /dev/null +++ b/test/set.jl @@ -0,0 +1,132 @@ +@testset "Test PrognosticVariables set!" begin + + nlayers = 8 + trunc = 31 + NF = Float64 + complex_NF = Complex{NF} + spectral_grid = SpectralGrid(; NF, trunc, nlayers) # define resolution + model = PrimitiveWetModel(; spectral_grid) # construct model + simulation = initialize!(model) # initialize all model components + + lmax = model.spectral_transform.lmax + mmax = model.spectral_transform.mmax + lf = 2 + + # test data + L = rand(spectral_grid.SpectralVariable3D, trunc+2, trunc+1, nlayers) + L_grid = transform(L, model.spectral_transform) + + L2 = rand(spectral_grid.SpectralVariable3D, trunc-5, trunc-6, nlayers) # smaller + L2_trunc = spectral_truncation(L2, size(L, 1, ZeroBased, as=Matrix), size(L, 2, ZeroBased, as=Matrix)) + L3 = rand(spectral_grid.SpectralVariable3D, trunc+6, trunc+5, nlayers) # bigger + L3_trunc = spectral_truncation(L3, size(L, 1, ZeroBased, as=Matrix), size(L, 2, ZeroBased, as=Matrix)) + + A = rand(spectral_grid.Grid{NF}, spectral_grid.nlat_half, nlayers) # same grid + A_spec = transform(A, model.spectral_transform) + B = rand(OctaHEALPixGrid{NF}, spectral_grid.nlat_half, nlayers) # different grid + + f(lon, lat, sig) = sind(lon)*cosd(lat)*(1 - sig) + + prog_new = simulation.prognostic_variables + prog_old = deepcopy(prog_new) + + # set things ... + + # LTA + set!(simulation, vor=L, lf = lf) + @test prog_new.vor[lf] == L + + set!(simulation, div=L, lf = lf; add=true) + @test prog_new.div[lf] == (prog_old.div[lf] .+ L) + + set!(simulation, temp=L2, lf=lf) + @test prog_new.temp[lf] ≈ L2_trunc + + set!(simulation, humid=L3, lf=lf) + @test prog_new.humid[lf] ≈ L3_trunc + + set!(simulation, pres=L[:,1], lf=lf) + @test prog_new.pres[lf] == L[:,1] + + set!(simulation, vor=A, lf=lf) + @test prog_new.vor[lf] == A_spec + + set!(simulation, vor=A, lf=lf, add=true) + @test prog_new.vor[lf] == (2 .* A_spec) + + # grids + set!(simulation, sea_surface_temperature=A[:,1], lf=lf) + @test prog_new.ocean.sea_surface_temperature == A[:,1] + + set!(simulation, sea_ice_concentration=B[:,1], lf=lf, add=true) + C = similar(A[:,1]) + SpeedyWeather.RingGrids.interpolate!(C, B[:,1]) + @test prog_new.ocean.sea_ice_concentration ≈ (prog_old.ocean.sea_ice_concentration .+ C) + + set!(simulation, land_surface_temperature=L[:,1], lf=lf) + @test prog_new.land.land_surface_temperature ≈ L_grid[:,1] + + set!(simulation, soil_moisture_layer2=L[:,1], lf=lf) + @test prog_new.land.soil_moisture_layer2 ≈ L_grid[:,1] + + # numbers + set!(simulation, vor=Float32(3.), lf=lf) + M3 = zeros(spectral_grid.Grid{NF}, spectral_grid.nlat_half, nlayers) .+ 3 # same grid + M3_spec = transform(M3, model.spectral_transform) + @test prog_new.vor[lf] ≈ M3_spec + + set!(simulation, vor=Float32(3.), lf=lf, add=true) + @test prog_new.vor[lf] ≈ (2 .* M3_spec) + + set!(simulation, sea_surface_temperature=Float16(3.)) + @test all(prog_new.ocean.sea_surface_temperature .≈ 3.) + + set!(simulation, sea_surface_temperature=Float16(3.), add=true) + @test all(prog_new.ocean.sea_surface_temperature .≈ 6.) + + # vor_div, create u,v first in spectral space + u = randn(spectral_grid.SpectralVariable3D, trunc+2, trunc+1, nlayers) + v = randn(spectral_grid.SpectralVariable3D, trunc+2, trunc+1, nlayers) + + # set imaginary component of m=0 to 0 as the rotation of zonal modes is arbitrary + SpeedyTransforms.zero_imaginary_zonal_modes!(u) + SpeedyTransforms.zero_imaginary_zonal_modes!(v) + + spectral_truncation!(u, 25) # truncate to lowest 11 wavenumbers + spectral_truncation!(v, 25) + + u_grid = transform(u, model.spectral_transform) + v_grid = transform(v, model.spectral_transform) + + set!(simulation, u=u_grid, v=v_grid, coslat_scaling_included=false, lf=lf) + + # now obtain U, V (scaled with coslat) from vor, div + U = similar(u) + V = similar(v) + + SpeedyTransforms.UV_from_vordiv!(U, V, prog_new.vor[lf], prog_new.div[lf], model.spectral_transform) + + # back to grid and unscale on the fly + u_grid2 = transform(U, model.spectral_transform, unscale_coslat=true) + v_grid2 = transform(V, model.spectral_transform, unscale_coslat=true) + + u2 = transform(u_grid2, model.spectral_transform) + v2 = transform(v_grid2, model.spectral_transform) + + # for lm in eachindex(u, v, u2, v2) + # @test u[lm] ≈ u2[lm] atol = sqrt(sqrt(eps(spectral_grid.NF))) + # @test v[lm] ≈ v2[lm] atol = sqrt(sqrt(eps(spectral_grid.NF))) + # end + + # functions + (; londs, latds, σ_levels_full) = model.geometry + for k in SpeedyWeather.RingGrids.eachgrid(A) + for ij in SpeedyWeather.RingGrids.eachgridpoint(A) + A[ij,k] = f(londs[ij], latds[ij], σ_levels_full[k]) + end + end + transform!(A_spec, A, model.spectral_transform) + + set!(simulation, vor=f; lf) + @test prog_new.vor[lf] ≈ A_spec +end \ No newline at end of file diff --git a/test/set_vars.jl b/test/set_vars.jl deleted file mode 100644 index 2a1638cc9..000000000 --- a/test/set_vars.jl +++ /dev/null @@ -1,103 +0,0 @@ -@testset "Test PrognosticVariables set_vars! and get_var" begin - - # test setting LowerTriangularMatrices - spectral_grid = SpectralGrid() - initial_conditions = StartFromRest() - M = PrimitiveWetModel(; 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{spectral_grid.NF}, lmax+1, mmax+1) for i=1:nlev] - - SpeedyWeather.set_vorticity!(P, sph_data) - SpeedyWeather.set_divergence!(P, sph_data) - SpeedyWeather.set_temperature!(P, sph_data) - SpeedyWeather.set_humidity!(P, sph_data) - SpeedyWeather.set_pressure!(P, sph_data[1]) - - for i=1:nlev - @test P.layers[i].timesteps[lf].vor == sph_data[i] - @test P.layers[i].timesteps[lf].div == sph_data[i] - @test P.layers[i].timesteps[lf].temp == sph_data[i] - @test P.layers[i].timesteps[lf].humid == sph_data[i] - end - @test P.surface.timesteps[lf].pres == sph_data[1] - - @test SpeedyWeather.get_vorticity(P) == sph_data - @test SpeedyWeather.get_divergence(P) == sph_data - @test SpeedyWeather.get_temperature(P) == sph_data - @test SpeedyWeather.get_humidity(P) == sph_data - @test SpeedyWeather.get_pressure(P) == sph_data[1] - - SpeedyWeather.set_vorticity!(P, 0) - for i=1:nlev - @test all(P.layers[i].timesteps[lf].vor .== 0) - end - - grid_data = [gridded(sph_data[i], M.spectral_transform) for i in eachindex(sph_data)] - - SpeedyWeather.set_vorticity!(P, grid_data) - SpeedyWeather.set_divergence!(P, grid_data) - SpeedyWeather.set_temperature!(P, grid_data) - SpeedyWeather.set_humidity!(P, grid_data) - SpeedyWeather.set_pressure!(P, grid_data[1]) - - for i=1:nlev - @test P.layers[i].timesteps[lf].vor ≈ sph_data[i] - @test P.layers[i].timesteps[lf].div ≈ sph_data[i] - @test P.layers[i].timesteps[lf].temp ≈ sph_data[i] - @test P.layers[i].timesteps[lf].humid ≈ sph_data[i] - end - @test P.surface.timesteps[lf].pres ≈ sph_data[1] - - grid_data = [gridded(sph_data[i], M.spectral_transform) for i in eachindex(sph_data)] - - SpeedyWeather.set_vorticity!(P, grid_data, M) - SpeedyWeather.set_divergence!(P, grid_data, M) - SpeedyWeather.set_temperature!(P, grid_data, M) - SpeedyWeather.set_humidity!(P, grid_data, M) - SpeedyWeather.set_pressure!(P, grid_data[1], M) - - for i=1:nlev - @test P.layers[i].timesteps[lf].vor ≈ sph_data[i] - @test P.layers[i].timesteps[lf].div ≈ sph_data[i] - @test P.layers[i].timesteps[lf].temp ≈ sph_data[i] - @test P.layers[i].timesteps[lf].humid ≈ sph_data[i] - end - @test P.surface.timesteps[lf].pres ≈ sph_data[1] - - # test setting matrices - spectral_grid = SpectralGrid(Grid=FullGaussianGrid) - initial_conditions = StartFromRest() - M = PrimitiveWetModel(; 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, matrix_data) - SpeedyWeather.set_divergence!(P, matrix_data) - SpeedyWeather.set_temperature!(P, matrix_data) - SpeedyWeather.set_humidity!(P, matrix_data) - SpeedyWeather.set_pressure!(P, matrix_data[1]) - - for i=1:nlev - @test P.layers[i].timesteps[lf].vor ≈ sph_data[i] - @test P.layers[i].timesteps[lf].div ≈ sph_data[i] - @test P.layers[i].timesteps[lf].temp ≈ sph_data[i] - @test P.layers[i].timesteps[lf].humid ≈ sph_data[i] - end - @test P.surface.timesteps[lf].pres ≈ sph_data[1] - -end \ No newline at end of file diff --git a/test/spectral_gradients.jl b/test/spectral_gradients.jl index 524f3ba84..1fe4a431e 100644 --- a/test/spectral_gradients.jl +++ b/test/spectral_gradients.jl @@ -1,42 +1,41 @@ @testset "Divergence of a non-divergent flow zero?" begin @testset for NF in (Float32, Float64) - spectral_grid = SpectralGrid(; NF, nlev=1) - m = ShallowWaterModel(; spectral_grid) - simulation = initialize!(m) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables + spectral_grid = SpectralGrid(; NF, nlayers=1) + model = ShallowWaterModel(; spectral_grid) + simulation = initialize!(model) + progn = simulation.prognostic_variables + diagn = 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) + # make sure vorticity and divergence are 0 + progn.vor[1] .= 0 + progn.div[1] .= 0 + diagn.tendencies.vor_tend .= 0 # start with some vorticity only - vor0 = randn(LowerTriangularMatrix{Complex{NF}}, p.trunc+2, p.trunc+1) - p.layers[1].timesteps[1].vor .= vor0 + progn.vor[1] .= randn(Complex{NF}, size(progn.vor[1])...) + # get corresponding non-divergent u_grid, v_grid lf = 1 - SpeedyWeather.gridded!(d, p, lf, m) # get corresponding non-divergent u_grid, v_grid + transform!(diagn, progn, lf, model) - u_grid = d.layers[1].grid_variables.u_grid - v_grid = d.layers[1].grid_variables.v_grid + (; u_grid, v_grid) = diagn.grid - # check we've actually created non-zero U=u*coslat, V=v*coslat + # check we've actually created non-zero u,v excl coslat scaling @test all(u_grid .!= 0) @test all(v_grid .!= 0) - G = m.geometry - S = m.spectral_transform RingGrids.scale_coslat⁻¹!(u_grid) RingGrids.scale_coslat⁻¹!(v_grid) - - uω_coslat⁻¹ = d.layers[1].dynamics_variables.a - vω_coslat⁻¹ = d.layers[1].dynamics_variables.b - - SpeedyWeather.spectral!(uω_coslat⁻¹, u_grid, S) - SpeedyWeather.spectral!(vω_coslat⁻¹, v_grid, S) + + uω_coslat⁻¹ = diagn.dynamics.a + vω_coslat⁻¹ = diagn.dynamics.b + + S = model.spectral_transform + SpeedyWeather.transform!(uω_coslat⁻¹, u_grid, S) + SpeedyWeather.transform!(vω_coslat⁻¹, v_grid, S) - div = zero(vor0) + div = progn.div[1] SpeedyWeather.divergence!(div, uω_coslat⁻¹, vω_coslat⁻¹, S) for div_lm in div @@ -48,42 +47,39 @@ end @testset "Curl of an irrotational flow zero?" begin @testset for NF in (Float32, Float64) - spectral_grid = SpectralGrid(; NF, nlev=1) - m = ShallowWaterModel(; spectral_grid) - simulation = initialize!(m) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables + spectral_grid = SpectralGrid(; NF, nlayers=1) + model = ShallowWaterModel(; spectral_grid) + simulation = initialize!(model) + progn = simulation.prognostic_variables + diagn = 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) + # make sure vorticity and divergence are 0 + progn.vor[1] .= 0 + progn.div[1] .= 0 + diagn.tendencies.div_tend .= 0 - # start with some vorticity only - div0 = randn(LowerTriangularMatrix{Complex{NF}}, p.trunc+2, p.trunc+1) - p.layers[1].timesteps[1].div .= div0 + # start with some divergence only + progn.div[1] .= randn(Complex{NF}, size(progn.div[1])...) + # get corresponding non-divergent u_grid, v_grid lf = 1 - SpeedyWeather.gridded!(d, p, lf, m) # get corresponding non-divergent u_grid, v_grid + transform!(diagn, progn, lf, model) - u_grid = d.layers[1].grid_variables.u_grid - v_grid = d.layers[1].grid_variables.v_grid + (; u_grid, v_grid) = diagn.grid - # check we've actually created non-zero U=u*coslat, V=v*coslat + # check we've actually created non-zero u,v (excl coslat scaling) @test all(u_grid .!= 0) @test all(v_grid .!= 0) - G = m.geometry - S = m.spectral_transform - C = m.coriolis - # to evaluate ∇×(uv) use curl of vorticity fluxes (=∇×(uv(ζ+f))) with ζ=1, f=0 - fill!(d.layers[1].grid_variables.vor_grid, 1) - fill!(m.coriolis.f, 0) + progn.div[1] .= 1 + model.coriolis.f .= 0 # calculate uω, vω in spectral space - SpeedyWeather.vorticity_flux_curldiv!(d.layers[1], C, G, S, div=true) + (; coriolis, geometry, spectral_transform) = model + SpeedyWeather.vorticity_flux_curldiv!(diagn, coriolis, geometry, spectral_transform, div=true) - for div_lm in d.layers[1].tendencies.div_tend + for div_lm in diagn.tendencies.div_tend @test abs(div_lm) < sqrt(eps(NF)) end end @@ -97,7 +93,7 @@ end OctahedralClenshawGrid, HEALPixGrid) - SG = SpectralGrid(NF; Grid, nlev=1) + SG = SpectralGrid(; NF, Grid, nlayers=1) G = Geometry(SG) A = Grid(randn(NF, SG.npoints)) @@ -118,7 +114,7 @@ end @testset "Flipsign in divergence!, curl!" begin @testset for NF in (Float32, Float64) - SG = SpectralGrid(NF) + SG = SpectralGrid(; NF) S = SpectralTransform(SG) lmax, mmax = S.lmax, S.mmax @@ -140,7 +136,7 @@ end @testset "Add in divergence!, curl!" begin @testset for NF in (Float32, Float64) - SG = SpectralGrid(NF) + SG = SpectralGrid(; NF) S = SpectralTransform(SG) lmax, mmax = S.lmax, S.mmax @@ -173,73 +169,59 @@ end @testset "D, ζ -> u, v -> D, ζ" begin @testset for NF in (Float32, Float64) - spectral_grid = SpectralGrid(; NF, nlev=1) - m = ShallowWaterModel(; spectral_grid) - simulation = initialize!(m) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables + spectral_grid = SpectralGrid(; NF, nlayers=2) + model = PrimitiveDryModel(; spectral_grid) + simulation = initialize!(model) + progn = simulation.prognostic_variables + diagn = simulation.diagnostic_variables + vor = progn.vor[1] + div = progn.div[1] # make sure vorticity and divergence are 0 - fill!(p.layers[1].timesteps[1].vor, 0) - fill!(p.layers[1].timesteps[1].div, 0) - - # make sure vorticity and divergence are 0 - fill!(d.layers[1].tendencies.vor_tend, 0) - fill!(d.layers[1].tendencies.div_tend, 0) + vor .= 0 + div .= 0 # create initial conditions - (; lmax, mmax) = m.spectral_transform - vor0 = rand(LowerTriangularMatrix{Complex{NF}}, lmax+1, mmax+1) - div0 = rand(LowerTriangularMatrix{Complex{NF}}, lmax+1, mmax+1) + vor .= rand(Complex{NF}, size(vor)...) + div .= rand(Complex{NF}, size(div)...) - vor0[1, 1] = 0 # zero mean - div0[1, 1] = 0 + for k in eachmatrix(vor, div) + vor[1, k] = 0 # zero mean on every layer + div[1, k] = 0 + end # set imaginary component of m=0 to 0 as the rotation of zonal modes is arbitrary - SpeedyTransforms.zero_imaginary_zonal_modes!(vor0) - SpeedyTransforms.zero_imaginary_zonal_modes!(div0) - - spectral_truncation!(vor0) # set unusued last row (l=lmax+1) to zero - spectral_truncation!(div0) + SpeedyTransforms.zero_imaginary_zonal_modes!(vor) + SpeedyTransforms.zero_imaginary_zonal_modes!(div) - # copy into prognostic variables - p.layers[1].timesteps[1].vor .= vor0 - p.layers[1].timesteps[1].div .= div0 + spectral_truncation!(vor) # set unusued last row (l=lmax+1) to zero + spectral_truncation!(div) - vor1 = zero(vor0) - div1 = zero(div0) - - # get corresponding irrotational u_grid, v_grid (incl *coslat scaling) + # get corresponding u_grid, v_grid (excl *coslat scaling) lf = 1 - SpeedyWeather.gridded!(d, p, lf, m) + SpeedyWeather.transform!(diagn, progn, lf, model) # check we've actually created non-zero u, v - @test all(d.layers[1].grid_variables.u_grid .!= 0) - @test all(d.layers[1].grid_variables.v_grid .!= 0) - - u = d.layers[1].grid_variables.u_grid - v = d.layers[1].grid_variables.v_grid + (; u_grid, v_grid) = diagn.grid + @test all(u_grid .!= 0) + @test all(v_grid .!= 0) # times coslat² in grid space - G = m.geometry - RingGrids.scale_coslat⁻¹!(u) - RingGrids.scale_coslat⁻¹!(v) + RingGrids.scale_coslat⁻¹!(u_grid) + RingGrids.scale_coslat⁻¹!(v_grid) # transform back - S = m.spectral_transform - u_coslat⁻¹ = zero(vor0) - v_coslat⁻¹ = zero(vor0) - SpeedyWeather.spectral!(u_coslat⁻¹, u, S) - SpeedyWeather.spectral!(v_coslat⁻¹, v, S) + u_coslat⁻¹ = transform(u_grid, model.spectral_transform) + v_coslat⁻¹ = transform(v_grid, model.spectral_transform) # curl and div in spectral space - SpeedyWeather.curl!(vor1, u_coslat⁻¹, v_coslat⁻¹, S) - SpeedyWeather.divergence!(div1, u_coslat⁻¹, v_coslat⁻¹, S) + vor2 = curl(u_coslat⁻¹, v_coslat⁻¹, model.spectral_transform) + div2 = divergence(u_coslat⁻¹, v_coslat⁻¹, model.spectral_transform) - for lm in SpeedyWeather.eachharmonic(vor0, vor1, div0, div1) + for lm in eachindex(vor, div, vor2, div2) # increased to 30 as 10, 20 caused single fails every now and then - @test vor0[lm] ≈ vor1[lm] rtol=30*sqrt(eps(NF)) - @test div0[lm] ≈ div1[lm] rtol=30*sqrt(eps(NF)) + @test vor[lm] ≈ vor2[lm] rtol=30*sqrt(eps(NF)) + @test div[lm] ≈ div2[lm] rtol=30*sqrt(eps(NF)) end end end @@ -294,48 +276,44 @@ end end @testset "∇×∇=0 and ∇⋅∇=∇²" begin - for NF in (Float32, Float64) - - trunc = 31 - spectral_grid = SpectralGrid(; NF, trunc, Grid=FullGaussianGrid, nlev=1) - m = ShallowWaterModel(; spectral_grid) - simulation = initialize!(m) - p = simulation.prognostic_variables - d = simulation.diagnostic_variables - - a = randn(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) - spectral_truncation!(a) - SpeedyTransforms.zero_imaginary_zonal_modes!(a) - - dadx = zero(a) - dady = zero(a) - SpeedyWeather.∇!(dadx, dady, a, m.spectral_transform) - - dadx_grid = gridded(dadx, m.spectral_transform) - dady_grid = gridded(dady, m.spectral_transform) - - RingGrids.scale_coslat⁻²!(dadx_grid) - RingGrids.scale_coslat⁻²!(dady_grid) - - SpeedyWeather.spectral!(dadx, dadx_grid, m.spectral_transform) - SpeedyWeather.spectral!(dady, dady_grid, m.spectral_transform) - - # CURL(GRAD(A)) = 0 - ∇x∇a = zero(a) - SpeedyWeather.curl!(∇x∇a, dadx, dady, m.spectral_transform) - - for lm in SpeedyWeather.eachharmonic(∇x∇a) - @test ∇x∇a[lm] ≈ 0 atol=5*sqrt(eps(NF)) - end - - # DIV(GRAD(A)) = LAPLACE(A) - div_∇a = zero(a) - SpeedyWeather.divergence!(div_∇a, dadx, dady, m.spectral_transform) - ∇²a = zero(a) - SpeedyWeather.∇²!(∇²a, a, m.spectral_transform) - - for lm in SpeedyWeather.eachharmonic(div_∇a, ∇²a) - @test div_∇a[lm] ≈ ∇²a[lm] atol=5*sqrt(eps(NF)) rtol=5*sqrt(eps(NF)) + for nlayers in (1, 2) + for NF in (Float32, Float64) + + trunc = 31 + spectral_grid = SpectralGrid(; NF, trunc, Grid=FullGaussianGrid, nlayers) + model = PrimitiveDryModel(; spectral_grid) + simulation = initialize!(model) + progn = simulation.prognostic_variables + diagn = simulation.diagnostic_variables + + a = randn(LowerTriangularArray{Complex{NF}}, trunc+2, trunc+1, nlayers) + spectral_truncation!(a) + SpeedyTransforms.zero_imaginary_zonal_modes!(a) + + dadx, dady = ∇(a, model.spectral_transform) + dadx_grid = transform(dadx, model.spectral_transform) + dady_grid = transform(dady, model.spectral_transform) + + RingGrids.scale_coslat⁻²!(dadx_grid) + RingGrids.scale_coslat⁻²!(dady_grid) + + transform!(dadx, dadx_grid, model.spectral_transform) + transform!(dady, dady_grid, model.spectral_transform) + + # CURL(GRAD(A)) = 0 + ∇x∇a = curl(dadx, dady, model.spectral_transform) + + for lm in eachharmonic(∇x∇a) + @test ∇x∇a[lm] ≈ 0 atol=5*sqrt(eps(NF)) + end + + # DIV(GRAD(A)) = LAPLACE(A) + ∇dot∇a = divergence(dadx, dady, model.spectral_transform) + ∇²a = ∇²(a, model.spectral_transform) + + for lm in eachindex(∇dot∇a, ∇²a) + @test ∇dot∇a[lm] ≈ ∇²a[lm] atol=5*sqrt(eps(NF)) rtol=5*sqrt(eps(NF)) + end end end end \ No newline at end of file diff --git a/test/spectral_transform.jl b/test/spectral_transform.jl index ee1ce843b..8e084eeff 100644 --- a/test/spectral_transform.jl +++ b/test/spectral_transform.jl @@ -50,14 +50,14 @@ spectral_resolutions_inexact = (127, 255) FullHEALPixGrid, FullOctaHEALPixGrid) - SG = SpectralGrid(NF; trunc, Grid) + SG = SpectralGrid(; NF, trunc, Grid) S = SpectralTransform(SG) alms = zeros(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) fill!(alms, 0) alms[1, 1] = 1 - map = gridded(alms, S) + map = transform(alms, S) for ij in SpeedyWeather.eachgridpoint(map) @test map[ij] ≈ map[1] > zero(NF) @@ -71,14 +71,14 @@ end for trunc in spectral_resolutions for NF in (Float32, Float64) - SG = SpectralGrid(NF; trunc) + SG = SpectralGrid(; NF, trunc) S1 = SpectralTransform(SG, recompute_legendre=true) S2 = SpectralTransform(SG, recompute_legendre=false) alms = randn(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) - map1 = gridded(alms, S1) - map2 = gridded(alms, S2) + map1 = transform(alms, S1) + map2 = transform(alms, S2) # is only approx as recompute_legendre may use a different precision @test map1 ≈ map2 @@ -94,7 +94,7 @@ end OctahedralGaussianGrid, OctahedralClenshawGrid) - SG = SpectralGrid(NF; trunc, Grid) + SG = SpectralGrid(; NF, trunc, Grid) S = SpectralTransform(SG, recompute_legendre=true) lmax = 3 @@ -103,8 +103,8 @@ end alms = zeros(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) alms[l, m] = 1 - map = gridded(alms, S) - alms2 = spectral(map, S) + map = transform(alms, S) + alms2 = transform(map, S) for lm in SpeedyWeather.eachharmonic(alms, alms2) @test alms[lm] ≈ alms2[lm] atol=100*eps(NF) @@ -116,6 +116,139 @@ end end end +@testset "Transform: Singleton dimensions" begin + @testset for trunc in spectral_resolutions + for NF in (Float32, Float64) + for Grid in ( FullGaussianGrid, + FullClenshawGrid, + OctahedralGaussianGrid, + OctahedralClenshawGrid) + + 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}}, trunc+2, trunc+1) + alms[l, m] = 1 + + map = transform(alms, S) + + # add singleton dimension for lower triangular matrix + alms2 = zeros(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1, 1) + alms2[:, 1] = alms + map2 = deepcopy(map) + transform!(map2, alms2, S) + + for ij in eachindex(map, map2) + @test map[ij] == map2[ij] + end + + # add singleton dimension for grid + grid = randn(SG.GridVariable2D, SG.nlat_half) + alms = transform(grid, S) + alms2 = deepcopy(alms) + + grid3D = zeros(SG.GridVariable3D, SG.nlat_half, 1) + grid3D[:, 1] = grid + + transform!(alms2, grid3D, S) + + for lm in eachindex(alms, alms2) + @test alms[lm] == alms2[lm] + end + end + end + end + end + end +end + +@testset "Transform: Real to real transform" begin + for NF in (Float32, Float64) + # test Float64 -> Float32 + spectral_grid = SpectralGrid(; NF) + + # create real and complex otherwise identical + Lreal = randn(LowerTriangularMatrix{NF}, spectral_grid.trunc+2, spectral_grid.trunc+1) + Lcomplex = complex.(Lreal) + + grid1 = zeros(spectral_grid.GridVariable2D, spectral_grid.nlat_half) + grid2 = zeros(spectral_grid.GridVariable2D, spectral_grid.nlat_half) + + S = SpectralTransform(spectral_grid) + + transform!(grid1, Lreal, S) + transform!(grid2, Lcomplex, S) + + for ij in eachindex(grid1, grid2) + @test grid1[ij] ≈ grid2[ij] + end + end +end + +@testset "Transform: NF flexibility spectral inputs" begin + + # test Float64 -> Float32 + spectral_grid = SpectralGrid() + L_f32 = randn(LowerTriangularMatrix{ComplexF32}, spectral_grid.trunc+2, spectral_grid.trunc+1) + L_f64 = ComplexF64.(L_f32) + + grid1 = zeros(spectral_grid.GridVariable2D, spectral_grid.nlat_half) + grid2 = zeros(spectral_grid.GridVariable2D, spectral_grid.nlat_half) + + S = SpectralTransform(spectral_grid) + + transform!(grid1, L_f64, S) # Float64 -> Float32 + transform!(grid2, L_f32, S) # Float32 -> Float32 + + for ij in eachindex(grid1, grid2) + @test grid1[ij] ≈ grid2[ij] atol=sqrt(eps(Float32)) + end + + # TODO NF flexibility for grid inputs currently not supported for a good reason: + # the fourier transform is pre-planned for number type NF so you can't use another one + # above in spectral->grid this doesn't matter as the fourier transform is applied after + # the Legendre transform (which isn't pre-planned and hence NF flexible) + # consequently, the following is commented out + + # L1 = deepcopy(L_f32) + # L2 = deepcopy(L_f32) + # fill!(L1, 0) + # fill!(L2, 0) + + # grid3_f32 = randn(spectral_grid.GridVariable2D, spectral_grid.nlat_half) + # grid3_f64 = Float64.(grid3_f32) + + # transform!(L1, grid3_f32, S) + # transform!(L2, grid3_f64, S) # throws an error +end + +@testset "Transform: NF flexibility for spectral outputs" begin + + # test Float32 -> Float64 + spectral_grid = SpectralGrid() + NF = spectral_grid.NF + grid = randn(spectral_grid.Grid{NF}, spectral_grid.nlat_half) + + L1 = zeros(LowerTriangularMatrix{Complex{Float32}}, spectral_grid.trunc+2, spectral_grid.trunc+1) + L2 = zeros(LowerTriangularMatrix{Complex{Float64}}, spectral_grid.trunc+2, spectral_grid.trunc+1) + + S = SpectralTransform(spectral_grid) + + transform!(L1, grid, S) # Float32 -> Float32 + transform!(L2, grid, S) # Float32 -> Float64 + + for lm in eachindex(L1, L2) + @test L1[lm] ≈ L2[lm] + end + + # TODO similar to the above, NF flexibility for grid outputs currently not supported because + # the fourier transform is pre-planned and inplace with LinearAlgebra.mul! which cannot write into + # another eltype than the FFT plan +end + @testset "Transform: Individual Legendre polynomials (inexact transforms)" begin @testset for trunc in spectral_resolutions_inexact @testset for NF in (Float32, Float64) @@ -124,7 +257,7 @@ end FullHEALPixGrid, FullOctaHEALPixGrid) - SG = SpectralGrid(NF; trunc, Grid) + SG = SpectralGrid(; NF, trunc, Grid) S = SpectralTransform(SG, recompute_legendre=true) lmax = 3 @@ -133,8 +266,8 @@ end alms = zeros(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) alms[l, m] = 1 - map = gridded(alms, S) - alms2 = spectral(map, S) + map = transform(alms, S) + alms2 = transform(map, S) tol = 1e-3 @@ -161,19 +294,19 @@ end # clenshaw-curtis grids are only exact for cubic truncation dealiasing = Grid in (FullGaussianGrid, OctahedralGaussianGrid) ? 2 : 3 - SG = SpectralGrid(NF; trunc, Grid, dealiasing) + SG = SpectralGrid(; NF, trunc, Grid, dealiasing) S = SpectralTransform(SG, recompute_legendre=false) O = EarthOrography(SG, smoothing=true) E = Earth(SG) initialize!(O, E, S) oro_grid = O.orography - oro_spec = spectral(oro_grid, S) + oro_spec = transform(oro_grid, S) - oro_grid1 = gridded(oro_spec, S) - oro_spec1 = spectral(oro_grid1, S) - oro_grid2 = gridded(oro_spec1, S) - oro_spec2 = spectral(oro_grid2, S) + oro_grid1 = transform(oro_spec, S) + oro_spec1 = transform(oro_grid1, S) + oro_grid2 = transform(oro_spec1, S) + oro_spec2 = transform(oro_grid2, S) tol = 1e-3 diff --git a/test/time_stepping.jl b/test/time_stepping.jl index 74baf0b9f..63a924636 100644 --- a/test/time_stepping.jl +++ b/test/time_stepping.jl @@ -19,7 +19,7 @@ end # loop over different precisions @testset for NF in (Float16, Float32, Float64) - spectral_grid = SpectralGrid(NF) + spectral_grid = SpectralGrid(; NF) L = Leapfrog(spectral_grid) # INITIAL CONDITIONS @@ -59,7 +59,7 @@ end # loop over different precisions @testset for NF in (Float16, Float32, Float64) - spectral_grid = SpectralGrid(NF) + spectral_grid = SpectralGrid(; NF) L = Leapfrog(spectral_grid) # INITIAL CONDITIONS @@ -87,7 +87,7 @@ end @test M_RAW < 1 # CHECK THAT NO WILLIAMS FILTER IS WORSE - spectral_grid = SpectralGrid(NF) + spectral_grid = SpectralGrid(; NF) L = Leapfrog(spectral_grid, williams_filter=1) # INITIAL CONDITIONS diff --git a/test/vertical_advection.jl b/test/vertical_advection.jl index e57c29164..506a25989 100644 --- a/test/vertical_advection.jl +++ b/test/vertical_advection.jl @@ -1,3 +1,20 @@ +@testset "Vertical advection stencils" begin + @test (1, 1, 2) == SpeedyWeather.SpeedyWeather.retrieve_stencil(1, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + @test (1, 2, 3) == SpeedyWeather.retrieve_stencil(2, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + @test (2, 3, 4) == SpeedyWeather.retrieve_stencil(3, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + @test (7, 8, 8) == SpeedyWeather.retrieve_stencil(8, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + + @test (1, 1, 2) == SpeedyWeather.retrieve_stencil(1, 5, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + @test (1, 2, 3) == SpeedyWeather.retrieve_stencil(2, 5, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + @test (2, 3, 4) == SpeedyWeather.retrieve_stencil(3, 5, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + @test (4, 5, 5) == SpeedyWeather.retrieve_stencil(5, 5, SpeedyWeather.CenteredVerticalAdvection{Float32, 1}()) + + @test (1, 1, 1, 2, 3) == SpeedyWeather.retrieve_stencil(1, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 2}()) + @test (1, 1, 2, 3, 4) == SpeedyWeather.retrieve_stencil(2, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 2}()) + @test (1, 2, 3, 4, 5) == SpeedyWeather.retrieve_stencil(3, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 2}()) + @test (6, 7, 8, 8, 8) == SpeedyWeather.retrieve_stencil(8, 8, SpeedyWeather.CenteredVerticalAdvection{Float32, 2}()) +end + @testset "Vertical advection runs" begin spectral_grid = SpectralGrid() model_types = (PrimitiveDryModel, PrimitiveWetModel) diff --git a/test/vertical_levels.jl b/test/vertical_coordinates.jl similarity index 60% rename from test/vertical_levels.jl rename to test/vertical_coordinates.jl index e0890376e..9f3a421f9 100644 --- a/test/vertical_levels.jl +++ b/test/vertical_coordinates.jl @@ -1,23 +1,23 @@ -@testset "Initialize sigma levels manually" begin +@testset "Initialize sigma layers manually" begin # automatic levels - spectral_grid = SpectralGrid(nlev=4) + spectral_grid = SpectralGrid(nlayers=4) G = Geometry(spectral_grid) @test length(G.σ_levels_half) == 5 @test length(G.σ_levels_full) == 4 # manual levels σ = SigmaCoordinates([0, 0.4, 0.6, 1]) - spectral_grid = SpectralGrid(vertical_coordinates=σ) + spectral_grid = SpectralGrid(nlayers=σ.nlayers, vertical_coordinates=σ) G = Geometry(spectral_grid) - @test spectral_grid.nlev == 3 + @test spectral_grid.nlayers == 3 @test length(G.σ_levels_half) == 4 @test length(G.σ_levels_full) == 3 # specify both σ = SigmaCoordinates([0, 0.4, 0.6, 1]) - spectral_grid = SpectralGrid(nlev=3, vertical_coordinates=σ) + spectral_grid = SpectralGrid(nlayers=3, vertical_coordinates=σ) G = Geometry(spectral_grid) - @test spectral_grid.nlev == 3 + @test spectral_grid.nlayers == 3 @test length(G.σ_levels_half) == 4 @test length(G.σ_levels_full) == 3 end \ No newline at end of file