Add tracer advection

[milankl]

fixes #564 by

• define Tracer (name, id, active)
• add model.tracers for a list of tracers
• based on length(model.tracers) add additional progn/diagn variables
• advect tracer variables
• time stepping for tracers
• set! for tracers
• 3D advection
• output for tracers (preliminary)
• tests
• documentation

[milankl]

new syntax is

add!(model, Tracer(:ch4))      # before initialize! or
add!(simulation, Tracer(:ch4)) # after

tracers are now just identified with their Symbol (not an id anymore) as their just live in a dictionary in the prognostic/diagnostic variables

[milankl]

Tracers can be added (add!(simulation, Tracer(:newtracer)) (and technically also be deleted) anytime, but with model.tracer[:co2].active = false also deactivated (=frozen in time). Implemented the set! function for tracers too,
so

julia> simulation.prognostic_variables.tracers
Dict{Symbol, Tuple{LowerTriangularArray{ComplexF32, 2, Matrix{ComplexF32}}, LowerTriangularArray{ComplexF32, 2, Matrix{ComplexF32}}}} with 2 entries:
  :ch4 => ([0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;], [0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;])
  :co2 => ([0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;], [0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;])

can be set! via

julia> set!(simulation, ch4=1)

julia> simulation.prognostic_variables.tracers
Dict{Symbol, Tuple{LowerTriangularArray{ComplexF32, 2, Matrix{ComplexF32}}, LowerTriangularArray{ComplexF32, 2, Matrix{ComplexF32}}}} with 2 entries:
  :ch4 => ([3.54491+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;], [0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;])
  :co2 => ([0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;], [0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;])

julia> set!(simulation, co2=2)

julia> simulation.prognostic_variables.tracers
Dict{Symbol, Tuple{LowerTriangularArray{ComplexF32, 2, Matrix{ComplexF32}}, LowerTriangularArray{ComplexF32, 2, Matrix{ComplexF32}}}} with 2 entries:
  :ch4 => ([3.54491+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;], [0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;])
  :co2 => ([7.08982+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;], [0.0+0.0im; 0.0+0.0im; … ; 0.0+0.0im; 0.0+0.0im;;])

@maximilian-gelbrecht which I implemented simply by checking whether additional kwargs are in the keys for progn.tracers. This happens if not

julia> set!(simulation, co3=2)
ERROR: UndefVarError: `co3` not defined

Technically we could do the same to avoid all the predefined vor = nothing, ... with isnothing(vor) checks. One could simply do :vor in keys(kwargs) && set!(..., kwargs[:vor], ...)

[milankl]

This seems to work!

out.mp4

done via

spectral_grid = SpectralGrid(trunc=63, nlayers=1)
model = ShallowWaterModel(spectral_grid)

# add tracer
add!(model, Tracer(:co2))
add!(model, SpeedyWeather.TracerOutput())

simulation = initialize!(model)

# set tracer and run a simulation
set!(simulation, co2 = (λ, φ, σ) -> 0 < λ < 10 ? 1 : 0)
run!(simulation, period=Day(10))

# reset tracer but continue simulation
set!(simulation, co2 = (λ, φ, σ) -> exp(-(λ-180)^2/10^2))
run!(simulation, period=Day(10), output=true)

[milankl]

And T170 resolution and 20days long

out.mp4

[milankl]

If everything passes this is ready to go. Tracer output needs some overhaul, but I want #633 merged first so will simply do this in another pull request

[maximilian-gelbrecht]

Great! I can do a review tomorrow or Friday.

[milankl]

Okay, docs built, see here https://speedyweather.github.io/SpeedyWeather.jl/previews/PR579/tracers/ hope you like the interface to tracers!!

[milankl]

I also think this is a good test case for the semi-Lagrangian advection I've been dreaming about! Because that should allow us to get to maybe 3-4k SYPD at T31 by allowing for larger time steps and cutting down on the number of transforms, especially for all the tracers. Every tracer adds 4 transforms (1x spec-> grid, 1x grid tendencies -> spec, 2x grid -> spec for the fluxes with u,v). So we could

1. start using semi-Lagrange for tracers only, compare speed, compare to Eulerian advection in terms of diffusion/accuracy/conservation
2. use semi-Lagrange for humidity (active tracer but maybe just straight forward otherwise), virtual temperature?
3. use semi-Lagrange for temperature (even more active tracer)
4. use for momentum too

The reason I say "dreaming" is that the semi-implicit semi-Lagrangian scheme is at the heart of what makes ECMWF's model so fast, allowing really big time steps and making the model very stable.