Verifying conserved quantities

[mini-DONG]

Hi, @milankl, I and @miniufo calculated some conserved quantities, such as:

• total circulation:
  $$C = \iint q h dxdy = \iint {\zeta} + f dxdy$$
  we used julia code like:

function total_C(ζ, η, model)
    AV = zero(ζ)
    f = coriolis(ζ)
    H = model.atmosphere.layer_thickness
    Hb = model.orography.orography
    @. AV = (ζ+f)/(η + H - Hb)*(η + H - Hb)
    real(spectral(AV)[1])/model.spectral_transform.norm_sphere
end
function total_C2(ζ, model)
    AV = zero(ζ)
    f = coriolis(ζ)
    @. AV = (ζ+f)
    real(spectral(AV)[1])/model.spectral_transform.norm_sphere
end
η = simulation.diagnostic_variables.surface.pres_grid;
ζ  = simulation.diagnostic_variables.layers[1].grid_variables.vor_grid;

but we got the different values:

println(total_C(ζ, η, model))
println(total_C2(ζ, model))

1.8663078e-13
-3.757976e-14
we think maybe it's a matter of numerical precision, so how do we change SpeedyWeather to float64?

[milankl]

Yeah with 1e-13 and lower it's a question of precision. Changing to Float64

spectral_grid = SpectralGrid(NF=Float64, trunc=31, nlev=1)
model = ShallowWaterModel(;spectral_grid)
simulation = initialize!(model)
run!(simulation, period=Day(10))

η = simulation.diagnostic_variables.surface.pres_grid
ζ  = simulation.diagnostic_variables.layers[1].grid_variables.vor_grid

I get with your functions O(1e-22)

julia> total_C2(ζ, model)
5.421747509835887e-22

julia> total_C(ζ, η, model)
4.376369334730129e-22

I don't know why one would calculate

$$ \iint qh dxdy$$

I've seen people doing this on staggered grids, but spectral models don't have staggered grids, and so divided by $h$ and then multiply by $h$ doesn't make sense to me. But instead of your total_C2 I'd even go one step further and split the integral. You can then also use $\zeta$ in spectral space where its mean is always exactly zero (in the unforced system) because its tendency is a divergence of a flux (similar argument to as before with the total mass in #400) see https://speedyweather.github.io/SpeedyWeather.jl/stable/shallowwater/#Shallow-water-equations

julia> ζ = simulation.prognostic_variables.layers[1].timesteps[1].vor;

julia> ζ[1]
0.0 + 0.0im

Then I'd calculate the total circulation as

function total_C3(ζ::LowerTriangularMatrix, model)    # ζ is spectral here!
    (; NF, Grid, nlat_half) = model.spectral_grid
    f = coriolis(Grid{NF}, nlat_half)
    ζ_mean = real(ζ[1]) / model.spectral_transform.norm_sphere
    f_mean = real(spectral(f)[1]) / model.spectral_transform.norm_sphere
    return f_mean + ζ_mean
end

where the ζ_mean is basically just a sanity check (or in cases where you force the vorticity equation). Then I get exactly zero in the standard T31 Float32 OctahedralGaussianGrid setup

julia> total_C3(ζ, model)
0.0f0

which boils down to the question of whether the coriolis parameter integrated over the sphere is zero (it should, but depending on grid, number format and resolution it has sometimes rounding errors! You can play around with this:

julia> spectral(coriolis(OctaHEALPixGrid{Float32}, 32))[1]
0.0f0 + 0.0f0im

julia> spectral(coriolis(FullGaussianGrid{Float64}, 64))[1]
2.960999582118573e-20 + 0.0im

julia> spectral(coriolis(OctahedralClenshawGrid{Float32}, 128))[1]
0.0f0 + 0.0f0im

julia> spectral(coriolis(OctahedralGaussianGrid{Float64}, 24))[1]
1.8763299776796212e-20 + 0.0im

the syntax is coriolis(Grid{NF}, nlat_half), Grid is one of our grids

FullClenshawGrid
FullGaussianGrid
FullHEALPixGrid
FullOctaHEALPixGrid
HEALPixGrid
OctaHEALPixGrid
OctahedralClenshawGrid
OctahedralGaussianGrid

NF is Float32 or Float64 (with full grids you can also use Float16 or BFloat16) and nlat_half is the resolution parameter of all grids, the number of latitude rings on one hemisphere, Equator included. Looks like Float16 is more accurate than Float64 here 😆

julia> spectral(coriolis(FullGaussianGrid{Float16},24))[1]
Float16(0.0) + Float16(0.0)im

julia> spectral(coriolis(FullGaussianGrid{Float64},24))[1]
1.8727327922476907e-20 + 0.0im

[miniufo]

Don't worry. We do know that the global integrals of $\zeta$ and $f$ are exactly zero. We have defined a function:

function ∬dA(v, h)
        h_weighted = zero(h)
        @. h_weighted = v * h
    
        return real(spectral(h_weighted)[1]) / norm
end

So we can check a series of Casimir, i.e., functions of PV, like total mass $v=q^0$, total circulation $v=q^1$, total enstrophy $v=q^2/2$ etc. which are expected to be conservative. In practice, higher order Casimir are fragile and less conservative in the presence of small dissipations. So glad to see the first two do conserve to an almost perfect sense.

Just happen to see that Float32 is not enough for validating the total circulation. We can now move on to the rest quantites.

[miniufo]

Now we need to get these values for every time step. Do we need to explicitly loop over each step or is there an easy way to do so? Then we could see their temporal variations (like this plot).

[milankl]

Don't worry. We do know that the global integrals of ζ and f are exactly zero.

Sorry, of course you do, that's not what I meant. I was referring to the numerical integral. Getting that one exactly to zero is anything but trivial. And so in practice many would just consider small numbers to be zero. For me that's <1e-7 where I get handwavy, others are a bit more picky. That's why I'm sometimes actually surprised when an integral turns out to be exactly zero. Like woaaaa, all numerical integration and roundings errors cancel 🎉

Just happen to see that Float32 is not enough for validating the total circulation.

This one for example. I know that many people consider only 0 to be 0, but when using Float32 I don't consider it "not enough", I consider 0 to be anything between -1e-7 and 1e-7. It's an uncertainty where you couldn't distinguish a number anymore from 0. Any error that's on the order of the rounding error that's fantastic! The smallest error you could have. Because normally all your other errors are massive. If you calculate the trajectory of a rocket to Saturn with Float64 and with Float32 it ends up crashing into Jupiter, then Float32 is not enough (although in most cases it's your algorithm that's not good enough for Float32...). Sorry you're talking to the guy who believes we should actually run all our climate models in 16 bits. Tangent over.

[milankl]

Now we need to get these values for every time step. Do we need to explicitly loop over each step or is there an easy way to do so?

You could do that, like

for i in 1:100
    run!(simulation, period=Day(1))
    M[i] = total_mass(...)
end

But that's a) start and stop which introduces its own errors because the first time steps are handled numerically differently then the others and b) slower and c) you can't say period = just one time step (because of a)).

Alternatively you can define a callback, which let's you inject any code into a simulation which is then executed after every timestep / periodically / at specific times. Read in the documentation: https://speedyweather.github.io/SpeedyWeather.jl/dev/callbacks/

To define an EnergyConservation callback for example, I'd define the total_energy function as before just pass on the spectral transform to avoid recomputation on every time step (performance only)

function total_energy(u, v, η, model)
    h = zero(u)
    E = zero(u)                             # allocate grid variable

    H = model.atmosphere.layer_thickness
    Hb = model.orography.orography
    g = model.planet.gravity

    @. h = η + H - Hb
    @. E = h/2*(u^2 + v^2) + g*h^2

    # reuse precomputed spectral transform for performance
    S = model.spectral_transform           

    # transform to spectral, take l=m=0 mode at [1] and normalize for mean
    E_mean = real(spectral(E, S)[1]) / model.spectral_transform.norm_sphere
end

function total_energy(diagn::DiagnosticVariables, model::ModelSetup)
    u = diagn.layers[1].grid_variables.u_grid
    v = diagn.layers[1].grid_variables.v_grid
    η = diagn.surface.pres_grid
    return total_energy(u, v, η, model)
end

and add a second method that takes the diagnostic variables and models and unpacks them and calls the first for convenience. Then our EnergyConservation callback needs a counter and a vector that records every time step

Base.@kwdef mutable struct EnergyConservation <: SpeedyWeather.AbstractCallback
    timestep_counter::Int = 0
    TE::Vector{Float32} = []     # total energy per time step
end

You can choose the name but it has to be <: SpeedyWeather.AbstractCallback (subtype of). You can define its fields as you like (add options, other vectors etc). Every callback has to extend the initialize!, callback! and finish! functions that determine what it does to initialize, on every time step and at the very end. In your case I'd do

function SpeedyWeather.initialize!(
    callback::EnergyConservation,
    progn::PrognosticVariables,
    diagn::DiagnosticVariables,
    model::ModelSetup,
)
    callback.TE = zeros(progn.clock.n_timesteps+1)  # replace with vector of correct length
    callback.TE[1] = total_energy(diagn, model)     # set initial conditions
    callback.timestep_counter = 1                   # (re)set counter to 1
end

function SpeedyWeather.callback!(
    callback::EnergyConservation,
    progn::PrognosticVariables,
    diagn::DiagnosticVariables,
    model::ModelSetup,
)
    callback.timestep_counter += 1  
    i = callback.timestep_counter
    callback.TE[i] = total_energy(diagn, model)
end

SpeedyWeather.finish!(::EnergyConservation, args...) = nothing

the function signatures (number of arguments and their types) have to match because this is how they are called inside the model, but you can write into the function bodies whatever you want and it'll be executed. Here, we're calling the total_energy function which isn't even defined inside the SpeedyWeather module (but here in global scope) so you've succesfully integrated your own piece of code into the model, yay!

Now we run this with

julia> spectral_grid = SpectralGrid(trunc=42, nlev=1)
julia> model = ShallowWaterModel(;spectral_grid)
julia> simulation = initialize!(model)

# create callback and add to the model
julia> energy_conservation_callback = EnergyConservation()
EnergyConservation <: AbstractCallback
├ timestep_counter::Int64 = 0
└── arrays: TE

julia> add!(model.callbacks, energy_conservation_callback)
[ Info: EnergyConservation callback added with key callback_Q3bN

# running this now will call the callback on every time step
julia> run!(simulation, period=Day(100))
Weather is speedy: 100%|██████████████████████████████| Time: 0:00:05 ( 4.43 millenia/day)

My callback got the random key :callback_Q3bN so now I can do

julia> TE = model.callbacks[:callback_Q3bN].TE
6401-element Vector{Float32}:
 9.069605f8
 9.0701824f8
 ⋮
 9.047158f8
 9.0471546f8

and with

using PythonPlot
PythonPlot.plot(TE/TE[1]*100)
xlabel("time step")
ylabel("Total energy [%]")
tight_layout()

I get

I believe for the previous plot I did record every day, not every time step. Super interesting how it's pretty well conserved but there are these smooth ups and downs!

[milankl]

If you want to track several quantities simultaneously, I'd probably just do

Base.@kwdef mutable struct CasimirRecorder <: SpeedyWeather.AbstractCallback
    timestep_counter::Int = 0
    M::Vector{Float32} = []    # mass per time step
    TE::Vector{Float32} = []   # total energy per time step
    Z::Vector{Float32} = []    # total enstrophy per time step
    # add others
end

(are all those Casimirs? I've heard about them from Rick Salmon's papers, but I'm not sure whether all quantities you mentioned are)

and then include them accordingly in the initialize! and callback! functions. Other remarks

• if you don't want to evaluate the callback every time step use a Schedule see here https://speedyweather.github.io/SpeedyWeather.jl/dev/callbacks/#Schedules
• you can define the finish! function to output a netcdf file

using NCDatasets
ds = NCDataset("total_energy.nc","c")
defVar(ds,"total energy",TE,("time step",))
close(ds)

For a more sophisticated netCDF output writer check our ParticleTracker which saves on the fly to netCDF and is implemented as a callback too. https://github.com/SpeedyWeather/SpeedyWeather.jl/blob/main/src/output/particle_tracker.jl

[miniufo]

I know that many people consider only 0 to be 0, but when using Float32 I don't consider it "not enough", I consider 0 to be anything between -1e-7 and 1e-7. It's an uncertainty where you couldn't distinguish a number anymore from 0.

I feel it is "not enough" because we calculate both circulation and enstrophy. Enstrophy should be positive-definite and turn out to be about 1e-13. But circulation is about 1e-7. This leads to a feeling that circulation is larger than a positve-definite enstrophy, which is also significantly larger than zero. Using Float64, circulation becomes 1e-22, which is much much smaller and we feel close "enough" to zero.

in most cases it's your algorithm that's not good enough for Float32...

Totally agreed. And you show us a perfect way to get an exact zero circulation.

Alternatively you can define a callback, which let's you inject any code into a simulation which is then executed after every timestep / periodically / at specific times.

The callback mechanism is great for this purpose. I am trying your code snippet:

spectral_grid = SpectralGrid(trunc=42, nlev=1)
model = ShallowWaterModel(;spectral_grid)
simulation = initialize!(model)

energy_conservation_callback = EnergyConservation()
add!(model, energy_conservation_callback)
run!(simulation, period=Day(100))

but get

MethodError: no method matching add!(::ShallowWaterModel{Float32, SpeedyWeather.DeviceSetup{SpeedyWeather.CPUDevice, DataType}, Earth{Float32}, EarthAtmosphere{Float32}, Coriolis{Float32}, EarthOrography{Float32, OctahedralGaussianGrid{Float32}}, NoForcing, NoDrag, NoParticleAdvection, InitialConditions{ZonalJet, ZeroInitially, ZeroInitially, ZeroInitially}, Leapfrog{Float32}, SpectralTransform{Float32}, ImplicitShallowWater{Float32}, HyperDiffusion{Float32}, Geometry{Float32}, OutputWriter{Float32, ShallowWater}, Feedback}, ::EnergyConservation)

Closest candidates are:
  add!(::Dict{Symbol, SpeedyWeather.AbstractCallback}, ::SpeedyWeather.AbstractCallback...)
   @ SpeedyWeather ~/.julia/packages/SpeedyWeather/rhmTB/src/output/callbacks.jl:78


Stacktrace:
 [1] top-level scope
   @ In[6]:15

Is this caused by different versions of Speedy?

[milankl]

Yes, sorry! add!(model, callback) only works on #main, you probably have to do add!(model.callbacks, callback). I've added the former only recently as a convenience function

[miniufo]

add!(model.callbacks, callback)

This works perfect.

Super interesting how it's pretty well conserved but there are these smooth ups and downs!

I do get some oscillation too. But your case could be induced by the imbalance between KE and PE, where you may forget to divide the PE by 2 at the code snippet here: @. E = h/2*(u^2 + v^2) + g*h^2

[milankl]

Oh yes, I've probably look at the Bernoulli potential for too long. How does the total energy change with the additional $1/2$?

[miniufo]

I got this plot. So both KE and PE oscillate quite significantly. These oscillations reduce greatly once you sum them up. But if you miss a factor of 1/2 in PE, then may not cancel so well as shown here.

This is initialized with ZonalJet(). Looks like the flow and mass are not in an exact balance so that the kinetic energy is continuously converted into potential energy, while their sum is almost conserved (~0.01%).

[milankl]

Awesome so it's even better conserved than I thought it would be! Honestly, I would have expected energy conservation to only be in the order of a few percent, this is 100x better than that. Makes me very happy to see. In the end we have a horizontal diffusion that takes out kinetic energy as well as the RAW filter that takes out energy during the time integration. Given the jump on the (what appears to be) first time step(s) the time integration might contribute to the energy loss significantly, one would need to play around with this. Thanks for doing this analysis, btw!

I think the Galewsky jet with ZonalJet() is only in balance without orography, you could see what happens if you use NoOrography instead ShallowWaterModel(;spectral_grid, orography = NoOrography(spectral_grid)).

[miniufo]

I think the Galewsky jet with ZonalJet() is only in balance without orography, you could see what happens if you use NoOrography instead ShallowWaterModel(;spectral_grid, orography = NoOrography(spectral_grid)).

Yes, we have used NoOrography already. So the dissipation of energy is quite weak. Glad to see this too.

For the balance, we mean gradient-wind balance. Actually, we could get a gradient wind:

$$u_{grd} = \left(\Omega^2 r^2 - \frac{r}{\sin\phi}\frac{\partial gh}{a\partial \phi}\right)^{1/2}- \Omega r$$

which is much larger than initial $u$. So the oscillation could be the imbalance between mass and wind fields which then leads to geostrophic adjustment.

[milankl]

I thought I had tested the initial conditions to be stable without the perturbation of the zonal jet that's on by default, but I see you are right, it is slightly unstable. Your plot suggests that given the initial conditions for $\eta$ one should use a higher windspeed but also the location is a tiny bit off?

julia> spectral_grid = SpectralGrid(trunc=31, nlev=1)
julia> initial_conditions = ZonalJet(perturb_height = 0)
julia> model = ShallowWaterModel(;spectral_grid, initial_conditions, orography=NoOrography(spectral_grid))
julia> simulation = initialize!(model)

Maybe I did something wrong when coding it up? It's defined here

SpeedyWeather.jl/src/dynamics/initial_conditions.jl

Lines 82 to 193 in a6898b6

	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.
	$(TYPEDFIELDS)"""
	Base.@kwdef mutable struct ZonalJet <: AbstractInitialConditions
	"jet latitude [˚N]"
	latitude::Float64 = 45
	"jet width [˚], default ≈ 19.29˚"
	width::Float64 = (1/4-1/7)*180
	"jet maximum velocity [m/s]"
	umax::Float64 = 80
	"perturbation latitude [˚N], position in jet by default"
	perturb_lat::Float64 = latitude
	"perturbation longitude [˚E]"
	perturb_lon::Float64 = 270
	"perturbation zonal extent [˚], default ≈ 19.1˚"
	perturb_xwidth::Float64 = 1/3*360/2π
	"perturbation meridinoal extent [˚], default ≈ 3.8˚"
	perturb_ywidth::Float64 = 1/15*360/2π
	"perturbation amplitude [m]"
	perturb_height::Float64 = 120
	end
	"""
	$(TYPEDSIGNATURES)
	Initial conditions from Galewsky, 2004, Tellus"""
	function initialize!( progn::PrognosticVariables,
	initial_conditions::ZonalJet,
	model::ModelSetup)
	(; latitude, width, umax) = initial_conditions # for jet
	(; perturb_lat, perturb_lon, perturb_xwidth, # for perturbation
	perturb_ywidth, perturb_height) = initial_conditions
	θ₀ = (latitude-width)/360*2π # southern boundary of jet [radians]
	θ₁ = (latitude+width)/360*2π # northern boundary of jet
	eₙ = exp(-4/(θ₁-θ₀)^2) # normalisation
	θ₂ = perturb_lat*2π/360 # perturbation latitude [radians]
	α = perturb_xwidth*2π/360 # zonal extent of interface perturbation [radians]
	β = perturb_ywidth*2π/360 # meridional extent of interface perturbation [radians]
	λ = perturb_lon*2π/360 # perturbation longitude [radians]
	(; rotation, gravity) = model.planet
	(; radius) = model.spectral_grid
	# always create on F64 grid then convert to spectral and interpolate there
	Grid = FullGaussianGrid
	nlat_half = 64
	u_grid = zeros(Grid, nlat_half)
	η_grid = zeros(Grid, nlat_half)
	colats = RingGrids.get_colat(Grid, nlat_half)
	_, lons = RingGrids.get_colatlons(Grid, nlat_half)
	weights = FastGaussQuadrature.gausslegendre(2nlat_half)[2]
	η_sum = 0
	for (j, ring) in enumerate(eachring(u_grid, η_grid))
	θ = π/2 - colats[j] # latitude in radians
	coslat⁻¹j = 1/cos(θ)
	f = 2rotation*sin(θ)
	# velocity per latitude
	if θ₀ < θ < θ₁
	u_θ = umax/eₙ*exp(1/(θ-θ₀)/(θ-θ₁)) # u as in Galewsky, 2004
	else
	u_θ = 0
	end
	# integration for layer thickness h / interface height η
	w = weights[j]
	η_sum += 2w*(radius*u_θ/gravity * (f + tan(θ)/radius*u_θ))
	# lon-constant part of perturbation
	ηθ = perturb_height*cos(θ)*exp(-((θ₂-θ)/β)^2)
	# store in all longitudes
	for ij in ring
	u_grid[ij] = u_θ/radius*coslat⁻¹j # include scaling for curl!
	# calculate perturbation (possibly shifted in lon compared to Galewsky 2004)
	ϕ = lons[ij] - λ
	η_grid[ij] = η_sum + exp(-(ϕ/α)^2)*ηθ
	end
	end
	u = spectral(u_grid)
	η = spectral(η_grid)
	# interpolate in spectral space to desired resolution
	(; lmax, mmax) = model.spectral_transform
	(; NF) = model.spectral_grid
	u = spectral_truncation(complex(NF), u, lmax, mmax)
	# 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]
	curl!(vor, u, v, model.spectral_transform)
	# transform interface height η (use pres as prognostic variable) in spectral
	(; pres) = progn.surface.timesteps[1]
	copyto!(pres, η)
	spectral_truncation!(pres)
	end

The calculation of $u$ looks pretty straight forward and correct (=80m/s by default) but maybe the meridional integral to obtain $\eta$ has a bug? Feel free to open a pull request if you find it!

[miniufo]

Just check the codes and expression in Galewsky 2004 paper. The codes look OK except the weights. You use 2w, which I am not sure if the missing $d\phi$ is included in the weights in the integration

$$h = h_0 - \int^{\phi} \frac{au}{g}\left[f+\frac{\tan\phi^{\prime}}{a}u\right]d\phi^{\prime}$$

https://github.com/SpeedyWeather/SpeedyWeather.jl/blob/a6898b6e1605b83b2854d981b75ac651a256c362/src/dynamics/initial_conditions.jl#L159C1-L161C67

Sorry I am not quite farmiliar with julia right now. But I do the numerical integration using python-xarray, which gives a balanced h different than that in ZonalJet().