Modified dynamics

[miniufo]

From the JOSS paper, I see that "operators used inside SpeedyWeather.jl are exposed to the user, facilitating analysis of the simulation data". Also, the slogan is "Play atmospheric modelling like it's LEGO".

Now I would like to try to modified the dynamics of shallow-water equations, hoping the design of SpeedyWeather could help on my own purpose. The paper by Theodore Shepherd provide a method of finding extremal states given an initial one. For shallow-water case [Eq. 3.26a, b],

the modified dynamics is Eq. (3.26a'): one would get extremal state through integrating Eq. (3.26a').

To implement this using speedy, I guess I need to:

• get the gradient of any variable simply by calling $\psi_x, \psi_y = \nabla\psi$;
• get inversion of Laplacian by $\psi = \nabla^{-2}(\zeta)$;
• get the curl of two variables in spectral space by \zeta_spec = curl(u_spec, v_spec);
• get product of two variables in grid space and then transform to spectral space by prod_spec = prod(A_spec, B_spec);

Not sure if this is enough but we can add more. I know that for spectral model, the first two can be done straighforwardly. I just need a doc to do this. The last two should follow the doc here. But I still need some code examples to really give a try. I could also contribute to the docs once this is done, and if you guys feel it is necessary.

[milankl]

I can look at it on Monday. For these types of projects I'd

• check whether I can implement the new term as a forcing/drag (whether it removes energy or not but in the docs you see the formulation we offer for additional terms.
• if not, try to bring the equations into a form that's as close as possible to an existing model that we have. Then we'd need to subtype ShallowWater for example so that you can reuse as much possible from the existing model

Looking at the equation, if I see this right doing some maths in my head, the Laplace of the Bernoulli potential appears when you take the divergence of that equation. That one we have too in our shallow water equations. What does the first term look like if you take curl and div? Because our models are formulated with vorticity and divergence as prognostic variables one should ideally do that here too.

[miniufo]

The modified dynamics will monotonically increase/decrease total energy but retain its PV (PV^2, PV^3... or Casimiars). I will do some math to get the equations of vorticity and divergence, and compare these two the standard equations. Maybe their differences can be implemented as the forcing terms.

[miniufo]

Sorry I made a mistake. The modified equations for shallow-water model should be Eq. (4.8) from Vallis et al. 1989 JFM:

where $\alpha$ determines the rate of energy removal. So basically, we only need to add the gradient of thickness-tendency to the momentum equation. Shepherd made another form of the last term as $\alpha\nabla\left[\nabla\cdot(\mathbf hu)\right]$ using continuity equation. Is it easy to implement this as a forcing term for the original shallow-water dynamics?

[milankl]

In vorticity-divergence formulation, $\nabla \times 4.8$ the additional term would drop out, but $\nabla \cdot 4.8$ would become

$$ \frac{\partial \mathcal{D}}{\partial t} - \nabla \times (\mathbf{u}(\zeta + f)) = F_\mathcal{D} + \nabla \cdot \mathbf{F}_\mathbf{u} -\nabla^2(\tfrac{1}{2}(u^2 + v^2) + g\eta + \alpha g \frac{\partial \eta}{\partial t}) $$

Scaling of the new term would be (continuity equation is scaled with radius $R$ so the new term needs to have a $1/R$ included)

$$ R^2 \frac{\partial \mathcal{D}}{\partial t} = \frac{\partial \mathcal{\tilde{D}}}{\partial \tilde{t}} + ... = ... -\tilde{\nabla}^2(\tfrac{1}{2}(u^2 + v^2) + g\eta + \frac{\alpha}{R} g \frac{\partial \eta}{\partial \tilde{t}}) $$

So if I would design that from scratch then I'd calculate the continuity equation $\partial \eta / \partial t - \nabla \cdot (\mathbf{u} h) = 0$ first (at the moment it's calculated last) and then reuse that to add this term to the Bernoulli potential when also $-g\eta$ is added to the Bernoulli potential (this is done in spectral space and also the tendency is available in spectral space, so that would be perfect). Given that the current order of terms is

1. Forcing/drag
2. Vorticity flux $\nabla \cdot (\mathbf{u}(\zeta + f)), \nabla \times (\mathbf{u}(\zeta + f))$
3. Geopotential $\Phi = g\eta$
4. Bernoulli potential $-\nabla^2(1/2(u^2 + v^2) + \Phi)$
5. Volume flux divergence $-\nabla \cdot (\mathbf{u} h)$

We need to calculate the forcing before the volume flux divergene has been calculated, which is also after the Bernoulli potential is added to the divergence tendency. So not ideal for our situation here. I therefore see 3 (actually, currently only 1) options

1. Implement as forcing, calculate volume flux divergence twice

(while trying to implement this I realised that our forcing/drag interface currently cannot depend on $\eta$, so the $h_t$ term here cannot be calculated as a forcing right now. This is because we're still working on an overhaul of the variables with #493 which depends on #503. After that this will be possible, sorry for this. While we could implement an intermediate fix before those PRs are merged, I'd rather do this after to make my life a bit easier)

2. Generalise forcing, drag to have a callsite option

Maybe in general our forcing/drag implementation should have the option to define a callsite, e.g. BeginningOfTimeStep <: AbstractCallsite, EndOfTimestep <: AbstractCallsite so that you can define a forcing also by when it's supposed to be executed. However, I'm not sure how many options to give, in this example one would ideally not just calculate the forcing after the geopotential but before the bernoulli potential, but also change the order of these executions (where we do have some freedom), hence the currently only possible solution would be as follows:

3. Hacky solution: Redefine the shallow water RHS

The right-hand side (RHS) is currently defined here (after the import steps you can just copy paste that back into your session, REPL, notebook, to go back to the original). We modify (and overwrite! hence hacky) this as follows

using SpeedyWeather

# avoids a bunch of `SpeedyWeather.` for non-exported functions and types
import SpeedyWeather: SurfaceVariables, forcing!, drag!, vorticity_flux!, geopotential!,
    bernoulli_potential!, volume_flux_divergence!

# MODIFIED
function SpeedyWeather.dynamics_tendencies!(  
    diagn::DiagnosticVariablesLayer,
    progn::PrognosticVariablesLayer,
    surface::SpeedyWeather.SurfaceVariables,
    pres::LowerTriangularMatrix,    # spectral pressure/η for geopotential
    time::DateTime,                 # time to evaluate the 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
    
    # = ∇×(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ), tendency for vorticity
    # = ∇⋅(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ), tendency for divergence
    vorticity_flux!(diagn, model)
                                           
    # (MODIFIED: MOVED UP) = -∇⋅(uh, vh), tendency for "pressure" η
    volume_flux_divergence!(diagn, surface, orography, atmosphere, geometry, spectral_transform)
    
    geopotential!(diagn, pres, planet)              # geopotential Φ = gη in shallow water
    
    # (MODIFIED: ADDED) add +αgη_t to geopotential
    modified_shallow_water!(diagn, surface, planet, geometry, forcing)
    
    bernoulli_potential!(diagn, spectral_transform) # = -∇²(E+Φ+αgη_t), tendency for divergence
end

import SpeedyWeather: AbstractPlanet, AbstractGeometry

function modified_shallow_water!(
    diagn::DiagnosticVariablesLayer,
    surface::SurfaceVariables,
    planet::AbstractPlanet,
    geometry::AbstractGeometry,
    forcing::ModifiedShallowWater,
)
    (; geopot) = diagn.dynamics_variables
    α = forcing.α / geometry.radius   # unscale with radius because pres_tend is scaled!
    geopot .+= α * planet.gravity * surface.pres_tend
end

I have then defined a dummy forcing that holds the $\alpha$ value

Base.@kwdef struct ModifiedShallowWater{NF} <: SpeedyWeather.AbstractForcing
    α::NF = 0.1    # should be a time scale?
end

ModifiedShallowWater(SG::SpectralGrid; kwargs...) = ModifiedShallowWater{SG.NF}(; kwargs...)
SpeedyWeather.initialize!(::ModifiedShallowWater,::ModelSetup) = nothing

# dummy forcing as it only holds the α for modified_shallow_water!
function SpeedyWeather.forcing!(
    diagn::DiagnosticVariablesLayer,
    progn::PrognosticVariablesLayer,
    forcing::ModifiedShallowWater,
    time::DateTime,
    model::ModelSetup)
    return nothing
end

So that we can do

spectral_grid = SpectralGrid(trunc=31,nlev=1)
modified_sw = ModifiedShallowWater(spectral_grid, α=0.1)
model = ShallowWaterModel(;spectral_grid, forcing=modified_sw)
simulation = initialize!(model)
run!(simulation)

This seems to work for me for a wide range of $\alpha$ values, but it eventually blows up when $\alpha$ is too large. I don't know how to test that it does what it's supposed to do. I'm still confused, $\alpha$ should be a time scale but the vallis paper only talks about $\alpha =0.1$ unless I've missed something. Could you check @miniufo ?

[milankl]

Shepherd made another form of the last term as α∇[∇⋅(hu)] using continuity equation. Is it easy to implement this as a forcing term for the original shallow-water dynamics?

I don't understand? This is the same as before, no?

[miniufo]

Hi @milankl, thank you very much for all these detailed explanations.

I don't understand? This is the same as before, no?

Yes, this is the same for $\alpha\nabla\partial_t h$ and $\alpha\nabla[\nabla\cdot(h\mathbf u)]$. Just not sure which form is preferred for an easy implementation. Since we could access surface.pres_tend, the first one would be much better.

After that this will be possible, sorry for this. While we could implement an intermediate fix before those PRs are merged, I'd rather do this after to make my life a bit easier)

Hope this does not give you too many troubles. I am also expecting the GPU version of the model.

Maybe in general our forcing/drag implementation should have the option to define a callsite, e.g. BeginningOfTimeStep <: AbstractCallsite, EndOfTimestep <: AbstractCallsite so that you can define a forcing also by when it's supposed to be executed.

I understand that the overall integration flows are twisted. My purpose here would be too special to be generalized. I think the hacky way is OK here. Thanks for all the code snippets. Right now I can make it work.

This seems to work for me for a wide range of values, but it eventually blows up when is too large. I don't know how to test that it does what it's supposed to do. I'm still confused, should be a time scale but the vallis paper only talks about unless I've missed something. Could you check @miniufo ?

I tried several $\alpha$, and found that large values would cause model to blow up. I guess this is somewhat similar to the linear damping: to increase the damping coefficient one needs to reduce dt to make the model stable (not very sure). The effect of adding this term is to increase/decrease the total energy of the flow while retain its moments of PV. I am checking these now.

If I want to output some new user-defined diagnostics, like mass streamfunction $\nabla^{*}\psi = h\mathbf u$, ($\nabla^{*}$ is a rotation of $\nabla$) or gradient of vorticity $\nabla \zeta$, is there a convient way to add these to model outputs?

[milankl]

Yes, this is the same for α∇∂th and α∇[∇⋅(hu)]. Just not sure which form is preferred for an easy implementation. Since we could access surface.pres_tend, the first one would be much better.

Yes, whether you call it $\partial_t h$ or $\nabla(\nabla \cdot (h \mathbf{u}))$ indeed the idea is to not calculate that term twice as I did above.

Hope this does not give you too many troubles.

No no it's great that you ask because these things should be easily possible. And to be fair, I also like the hacky way because it illustrates that you can indeed just redefine parts of the simulation. Sure, you need to know where to hook in (meaning I need to tell you I guess 😉) but at least you don't have to wait for the next release or a merged pull request to start experimenting with it.

Right now I can make it work.

Great. I'm also thinking about a general way that one can just change the order in which terms are calculated. Imagine you could just pass on order = (:continuity => 1, :vorticity => 2, :divergence => 3) to say in which order terms are supposed to be calculated. That'd be cool, but generally maybe also trick to generalise easily...

If I want to output some new user-defined diagnostics, like mass streamfunction ∇∗ψ=hu, (∇∗ is a rotation of ∇) or gradient of vorticity ∇ζ, is there a convient way to add these to model outputs?

The normal streamfunction $\Psi = \nabla^{-2} \zeta$ you can get like

vor = simulation.prognostic_variables.layers[1].timesteps[1].vor
R = spectral_grid.radius
Ψ = SpeedyTransforms.∇⁻²(vor) * R^2   # all gradient operators omit the radius scaling, do manually
plot(gridded(Ψ))

Do I see this right (doing some maths in my head) that you are looking for $\nabla^{-2}(\nabla \times (h\mathbf{u}))$? So like inverting vorticity but your vorticity would be calculated from the volume fluxes $h\mathbf{u}$ instead of the velocity $\mathbf{u}$ ? That would be

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
H = model.atmosphere.layer_thickness
Hb = model.orography.orography
h = @. η + H - Hb
uh = @. u*h
vh = @. v*h
vorh = SpeedyTransforms.curl(uh, vh) / R
Ψ = SpeedyTransforms.∇⁻²(vorh) * R^2
plot(gridded(Ψ))

For $\nabla \zeta$ I realise this is not as convenient as it could be, but here we go

dζdx = zero(vor)
dζdy = zero(vor)
SpeedyTransforms.∇!(dζdx, dζdy, vor, model.spectral_transform)
dζdx ./= R
dζdy ./= R
ζx = gridded(dζdx)
ζy = gridded(dζdy)
RingGrids.scale_coslat⁻¹!(ζx)
RingGrids.scale_coslat⁻¹!(ζy)

(this should be a convenience function like dζdx, dζdy = SpeedyTransforms.∇(vor) that does all the steps above)
I get (all correct?) EDIT, implemented and exported now in #523

[miniufo]

I try to follow your codes, and write everything into a callback, as previous global integrated diagnostics:

# add callbacks
function gradient(v)
    grdx = zero(v)
    grdy = zero(v)
    
    SpeedyTransforms.∇!(grdx, grdy, vor, model.spectral_transform)
    grdx ./= R
    grdy ./= R
    ∂x = gridded(grdx)
    ∂y = gridded(grdy)
    RingGrids.scale_coslat⁻¹!(∂x)
    RingGrids.scale_coslat⁻¹!(∂y)

    return ∂x, ∂y
end

function extra_diagnostics(u, v, ζ, η, model)
    uh = zero(u)
    vh = zero(u)
    ζh = zero(u)
    ψh = zero(u)
    δh = zero(u)
    χh = zero(u)
    
    # constants from model
    H  = model.atmosphere.layer_thickness
    Hb = model.orography.orography
    R  = model.spectral_grid.radius

    @. h  = η + H - Hb           # thickness
    @. uh = u*h
    @. vh = v*h
    
    ζh = SpeedyTransforms.curl(uh, vh) / R
    Ψh = SpeedyTransforms.∇⁻²(ζh) * R^2
    δh = SpeedyTransforms.divergence(uh, vh) / R
    χh = SpeedyTransforms.∇⁻²(δh) * R^2

    vhs, uhs = gradient(Ψh)
    uhd, vhd = gradient(χh)

    return Ψh, -uhs, vhs, χh, uhd, vhd
end

# unpack diagnostic variables and call global_diagnostics from above
function extra_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

    # vorticity during simulation is scaled by radius R, unscale here
    ζ = zero(ζR)
    @. ζ = ζR / diagn.scale[]

    return extra_diagnostics(u, v, ζ, η, model)
end

# define a GlobalDiagnostics callback and the fields it needs
Base.@kwdef mutable struct ExtraDiagnostics <: SpeedyWeather.AbstractCallback
    timestep_counter::Int = 0

    time::Vector{DateTime} = []
    M::Vector{Float64} = []  # mean mass per time step
    C::Vector{Float64} = []  # mean circulation per time step
    Λ::Vector{Float64} = []  # mean angular momentum per time step
    K::Vector{Float64} = []  # mean kinetic energy per time step
    P::Vector{Float64} = []  # mean potential energy per time step
    Q::Vector{Float64} = []  # mean enstrophy per time step
end

# define how to initialize a GlobalDiagnostics callback
function SpeedyWeather.initialize!(
    callback::ExtraDiagnostics,
    progn::PrognosticVariables,
    diagn::DiagnosticVariables,
    model::ModelSetup,
)
    # replace with vector of correct length
    n = progn.clock.n_timesteps + 1    # +1 for initial conditions
    callback.time = zeros(DateTime, n)

    v = diagn.layers[1].grid_variables.u_grid
    callback.ψh  = zeros(v)
    callback.χh  = zeros(v)
    callback.uhs = zeros(v)
    callback.vhs = zeros(v)
    callback.uhd = zeros(v)
    callback.vhd = zeros(v)

    Ψh, uhs, vhs, χh, uhd, vhd = extra_diagnostics(diagn, model)

    callback.time[1] = progn.clock.time
    callback.Ψh  = Ψh   # set initial conditions
    callback.uhs = uhs  # set initial conditions
    callback.vhs = vhs  # set initial conditions
    callback.χh  = χh   # set initial conditions
    callback.uhd = uhd  # set initial conditions
    callback.vhd = vhd  # set initial conditions

    callback.timestep_counter = 1  # (re)set counter to 1

    return nothing
end

# define what a GlobalDiagnostics callback does on every time step
function SpeedyWeather.callback!(
    callback::GlobalDiagnostics,
    progn::PrognosticVariables,
    diagn::DiagnosticVariables,
    model::ModelSetup,
)
    callback.timestep_counter += 1
    i = callback.timestep_counter

    Ψh, uhs, vhs, χh, uhd, vhd = extra_diagnostics(diagn, model)

    # store current time and diagnostics for timestep i
    callback.time[i] = progn.clock.time
    callback.Ψh  = Ψh
    callback.uhs = uhs
    callback.vhs = vhs
    callback.χh  = χh
    callback.uhd = uhd
    callback.vhd = vhd
end

# define how to finish a GlobalDiagnostics callback after simulation finished
function SpeedyWeather.finish!(
    callback::GlobalDiagnostics,
    progn::PrognosticVariables,
    diagn::DiagnosticVariables,
    model::ModelSetup,
)
    n_timesteps = callback.timestep_counter

    # create a netCDF file in current path
    ds = NCDataset(joinpath(pwd(), "extra_diagnostics.nc"), "c")

    # save diagnostics variables within
    defDim(ds, "time", n_timesteps)
    defVar(ds, "sfh",  callback.time,  ("time",))
    defVar(ds, "uhs",  callback.M,     ("time",))
    defVar(ds, "vhs",  callback.C,     ("time",))
    defVar(ds, "vph",  callback.Λ,     ("time",))
    defVar(ds, "uhd",  callback.K,     ("time",))
    defVar(ds, "vhd",  callback.P,     ("time",))

    close(ds)

    return nothing
end

There could be two problems:

1. How should I define the type of array in Base.@kwdef mutable struct ExtraDiagnostics <: SpeedyWeather.AbstractCallback? It is not a Vector{Float64}, but could be Matrix{Float64}?
2. When I try to output the extra diagnostics to a file, I realize that it is a 3D array of (t, y, x). When the simulation is long, one cannot hold such large array in memory. So I guess previous framework does not fit this here. Then how to hack it?

I really hope I can simply add some names of diagnostics to the output_vars of OutputWriter. Maybe pre-define some of them is possible?

[milankl]

This

I really hope I can simply add some names of diagnostics to the output_vars of OutputWriter. Maybe pre-define some of them is possible?

is the actual problem. Ideally I'd like to have a user interface that lets you add any user-defined variable easily, maybe like so

output = OutputWriter(spectral_grid)
streamfunction = StreamfunctionOutput(spectral_grid, output)
add!(output, streamfunction)

meaning that one would just need to define a field (or even a scalar) inside StreamfunctionOutput that's outputted on every output time step. How to calculate this is then defined in output!(::StreamfunctionOutput, progn, diagn, model) which is called by the output writer on every output time step. I see that eventually I would need to find the time to restructure the output to allow such an interface. With these projects it's always hard to imagine them before you ever face the need for it, so thanks for raising this and feel free to contribute any ideas on how the user interface should be structured, what features it should have etc. The more we know what's needed the quicker it is to implement this.

Second first

2. yes, you are right, for the previous examples just allocating a super long vector wasn't a problem, but you don't want to allocate a 2D field for every time step you could run your simulation for. This is where you need to open a netcdf file on disk and continuously write into it as the OutputWriter does it too, or as a callback, see how we implemented the Particle Tracker here https://github.com/SpeedyWeather/SpeedyWeather.jl/blob/main/src/output/particle_tracker.jl. However, you can only really write the full grids into file as they are representable as a matrix, the octahedral Gaussian grid for example you probably don't want to unravel as a vector, that's why we generally interpolate reduced grids to full grids equivalents for output. However, then the question is, given you calculate your extra diagnostics from $u, v, \zeta, \eta$ anyway, which are outputted as normally, why don't you just calculate these extra diagnostics offline, e.g.

# run a simulation as normal
using SpeedyWeather
spectral_grid = SpectralGrid(trunc=31, nlev=1)
model = ShallowWaterModel(;spectral_grid);
simulation = initialize!(model);
run!(simulation, output=true)

# now load in that data
using NCDatasets
ds = NCDataset("run_0001/output.nc")
vor = Float32.(ds["vor"][:,:,end])     # vorticity on last time step with "end", Float32. to convert from Union{Missing,Float32} to Float32

# wrap into FullGaussianGrid to do analysis with SpeedyWeather
vor_grid = FullGaussianGrid(vor)
vor_spectral = spectral(vor_grid)
Ψ = SpeedyTransforms.∇⁻²(vor_spectral) * spectral_grid.radius^2

yields

again (see post above). So if you load in the output as netcdf then it comes as Julia array, but you can wrap this into a grid with FullGaussianGrid (the default) and then continue as if you did online analysis. There's an additional interpolation error, but if you'd like to avoid that you can simulate on full grid directly, the shallow water model isn't considerably slower with FullGaussianGrid or FullClenshawGrid. Once you've wrapped u, v, ζ, η into a grid again you can use the functions you've defined above as if it was online - avoiding all the stuff you need to consider when writing your own netcdf output writer.

And first second

1. Yes you can put inside your custom callback (or any struct for that purpose) whatever you want, including Matrix{Float32} or Array{Float32, 3} (3 dimensional array) etc. but I hope given 2. that's not actually needed!

[milankl]

(Little feedback on best practices writing Julia functions, in this situation here

function gradient(v)
    grdx = zero(v)
    grdy = zero(v)
    
    SpeedyTransforms.∇!(grdx, grdy, vor, model.spectral_transform)
    grdx ./= R
    grdy ./= R

you are using globals. Neither vor, model, nor R are defined inside the function. In that situation Julia uses variable definitions from the outer scope. It is generally advised to avoid that as (1) it can more easily lead to unexpected outcomes (R might be redefined unless it's set as constant using const) (2) is less reproducible (you copied in only the functions not vor, model, nor R defined outside) (3) is bad for performance as Julia makes all decisions how to compile a function based in the types of the input arguments, neither vor, model, nor R are input arguments so the generated code will have to make a lot of decision at runtime depending on what vor, model, nor R are when called.

I can recommend https://docs.julialang.org/en/v1/manual/variables-and-scoping/ for more details!

Hope this helps!!)

[miniufo]

you are using globals. Neither vor, model, nor R are defined inside the function.

You are right. I should define gradient inside extra_diagnostics. I guess this would solve all the problems you mentioned.

So if you load in the output as netcdf then it comes as Julia array, but you can wrap this into a grid with FullGaussianGrid (the default) and then continue as if you did online analysis. There's an additional interpolation error, but if you'd like to avoid that you can simulate on full grid directly, the shallow water model isn't considerably slower with FullGaussianGrid or FullClenshawGrid. Once you've wrapped u, v, ζ, η into a grid again you can use the functions you've defined above as if it was online - avoiding all the stuff you need to consider when writing your own netcdf output writer.

That's great. I was a bit worry about writing data out because I thought this would lose some information. So I tried offline calculation as:

using NCDatasets

dsO = NCDataset(path * "run_0004/output.nc")
dsN = NCDataset(path * "run_0004/extra_diagnostics.nc", "c")

defDim(dsN, "time", Inf)       # unlimited time dimension
defDim(dsN, "lat", size(dsO["lat"])[1])
defDim(dsN, "lon", size(dsO["lon"])[1])

defVar(dsN, "time", Float64, ("time",),
    attrib=Dict("units"=>"day", "long_name"=>"time",
    "standard_name"=>"time", "calendar"=>"proleptic_gregorian"))

# coordinates of particles (the variables inside netCDF)
defVar(dsN, "lon", Float64, ("lon",), attrib = 
    Dict("long_name"=>"longitude", "units"=>"degrees_north"))

defVar(dsN, "lat", Float64, ("lat",), attrib = 
    Dict("long_name"=>"latitude", "units"=>"degrees_east"))

defVar(dsN, "vor", Float64, ("lon", "lat", "time"), attrib = 
    Dict("long_name"=>"vorticity", "units"=>"1/s"))

    
dsN["lon"][:]   = dsO["lon"]
dsN["lat"][:]   = dsO["lat"]

for i in 1:size(dsO["time"])[1]
    u = Float64.(dsO["u"][:,:,1,i])     # vorticity on last time step with "end", Float32. to convert from Union{Missing,Float32} to Float32
    v = Float64.(dsO["v"][:,:,1,i])
    
    # wrap into FullxxxGrid to do analysis with SpeedyWeather
    u_grid = FullClenshawGrid(u)
    v_grid = FullClenshawGrid(v)
    
    S = SpectralTransform(spectral_grid)
    
    u_spec = spectral(u_grid, S)
    v_spec = spectral(v_grid, S)
    
    R = model.spectral_grid.radius
    
    ζ_spec = SpeedyTransforms.curl(u_spec, v_spec) / R
    
    Ψ_spec = SpeedyTransforms.∇⁻²(ζ_spec) * R^2
    
    dsN["time"][i] = i
    dsN["vor"][:, :, i] = gridded(ζ_spec, S)
    NCDatasets.sync(dsN)
end

NCDatasets.close(dsN)
NCDatasets.close(dsO)

Then I can compare the online and offline vorticity. But I get this:

Maybe I missed something here.

I do need to write another file, as I still need to analyse using other python package. I also tried windspharm to do this offline calculation in python, but it turns out I cannot reproduce the online vorticity exactly.

[milankl]

I'm a bit confused why you are dealing with two nc datasets but the only thing you're missing is unscaling the coslat for the curl.

# wrap into FullxxxGrid to do analysis with SpeedyWeather
u_grid = FullClenshawGrid(u)
v_grid = FullClenshawGrid(v)

S = SpectralTransform(spectral_grid)

u_spec = spectral(u_grid, S)
v_spec = spectral(v_grid, S)

R = model.spectral_grid.radius

ζ_spec = SpeedyTransforms.curl(u_spec, v_spec) / R

needs to have

  RingGrids.scale_coslat⁻¹!(u_grid)
  RingGrids.scale_coslat⁻¹!(v_grid)

before the transform. This is explained in the docstring

help?> SpeedyWeather.curl
  curl(
      u::SpeedyWeather.RingGrids.AbstractGrid,
      v::SpeedyWeather.RingGrids.AbstractGrid
  ) -> Any


  Curl (∇×) of two vector components u, v on a grid. Applies 1/coslat scaling, transforms
  to spectral space and returns the spectral curl, which is scaled with the radius of the
  sphere. Divide by radius for unscaling.

  ───────────────────────────────────────────────────────────────────────────────────────

  curl(
      u::LowerTriangularMatrix,
      v::LowerTriangularMatrix
  ) -> Any


  Curl (∇×) of two vector components u, v of size (n+1)xn, the last row will be set to
  zero in the returned LowerTriangularMatrix. This function requires both u, v to be
  transforms of fields that are scaled with 1/cos(lat). An example usage is therefore

  RingGrids.scale_coslat⁻¹!(u_grid)
  RingGrids.scale_coslat⁻¹!(v_grid)
  u = spectral(u_grid)
  v = spectral(v_grid)
  vor = curl(u, v)
  vor_grid = gridded(div)

[miniufo]

I'm a bit confused why you are dealing with two nc datasets

The first is output from the model integration. I just want to output extra diagnostics to another nc file (hope you could refactor the output stream so this second file can be easily merged into the first one during integration). Right now I am testing if the vorticity can be reconstructed exactly. Then I will output streamfunction etc. to the second nc file.

but the only thing you're missing is unscaling the coslat for the curl.

I should have thought that, because it is so close (similar pattern).

[milankl]

I should have thought that, because it is so close (similar pattern).

Yes, you can visually see that there's a coslat scaling applied dampening the signals towards the poles 😉

[miniufo]

I am now back to the original issue of this post. Chaning the $\alpha$ really takes effects:

The energy of divergent mode is reduced! But not enough, we need to increase $\alpha$. However, a large $\alpha$ would make the model overflow, so I am trying to reduce the timestep of the model to see if this can solve the problem.

I print the model.time_stepping and get

Leapfrog{Float64} <: SpeedyWeather.AbstractTimeStepper
├ trunc::Int64 = 170
├ Δt_at_T31::Second = 1800 seconds
├ radius::Float64 = 6.371e6
├ adjust_with_output::Bool = true
├ robert_filter::Float64 = 0.05
├ williams_filter::Float64 = 0.53
├ Δt_millisec::Dates.Millisecond = 337500 milliseconds
├ Δt_sec::Float64 = 337.5
└ Δt::Float64 = 5.2974415319416104e-5

So I try to reduce dt as:

sec = 60
model.time_stepping.Δt_sec = sec
model.time_stepping.Δt_millisec = Dates.Millisecond(sec * 1000)
model.time_stepping.Δt = sec / 6371e3

I feel this is the hacky way to change $\Delta t$. It would be much better to do this from providing an argument to the constructor. Do you have any suggest if I would like to reduce $\Delta t$? It looks like $\Delta t$ is fixed for each resolution currently.

[milankl]

The time step is automatically chosen depending on resolution. At T31 it's Δt_at_T31 (default 30min) at other resolution it uses that value as a baseline and then scales the timestep. So for half the time step (at any resolution) compared to this default you'd set it to 15min, etc.

time_stepping = Leapfrog(spectral_grid,  Δt_at_T31 = Minute(15))
model = ShallowWaterModel(;spectral_grid, time_stepping)