Shallow water run

[mini-DONG]

Hi, It's awesome work here! I want to run a global shallow water model, and I'm reading the documentation. Now, I can run the example like this:

using SpeedyWeather
spectral_grid = SpectralGrid(trunc=63,nlev=1)
orography = NoOrography(spectral_grid)
initial_conditions = ZonalJet()
model = ShallowWaterModel(;spectral_grid, orography, initial_conditions)
simulation = initialize!(model)
run!(simulation,n_days=6)

Next, I want to use myself initial field, such as the u,v in 500hpa. How can I add this customized data?

Also, I have a few minor problems in other example of documentation - sometimes errors are reported:

• Example1

using SpeedyWeather
spectral_grid = SpectralGrid(trunc=63,nlev=1)
forcing = JetStreamForcing(spectral_grid,latitude=60)
drag = QuadraticDrag(spectral_grid)
output = OutputWriter(spectral_grid,ShallowWater,output_dt=6,output_vars=[:u,:v,:pres,:orography])
model = ShallowWaterModel(;spectral_grid,output,drag,forcing)
simulation = initialize!(model)
run!(simulation,n_days=20)   # discard first 20 days
run!(simulation,n_days=20,output=true)

with error:

UndefVarError: `JetStreamForcing` not defined

• Example2

using SpeedyWeather
spectral_grid = SpectralGrid(trunc=127,nlev=1)
time_stepping = SpeedyWeather.Leapfrog(spectral_grid,Δt_at_T31=30)
implicit = SpeedyWeather.ImplicitShallowWater(spectral_grid,α=0.5)
orography = EarthOrography(spectral_grid,smoothing=false)
initial_conditions = SpeedyWeather.RandomWaves()
output = OutputWriter(spectral_grid,ShallowWater,output_dt=12,output_vars=[:u,:pres,:div,:orography])
model = ShallowWaterModel(;spectral_grid,orography,output,initial_conditions,implicit,time_stepping)
simulation = initialize!(model)
run!(simulation,n_days=2,output=true)

with error:

MethodError: no method matching SpeedyWeather.ImplicitShallowWater(::SpectralGrid; α::Float64)

Ps: I use the v0.6.0. And I am pretty new to both numerical models and speaking English. So if I'm not clear, please let me know.

[milankl]

Thanks for using SpeedyWeather! No these are absolutely reasonable questions. It seems you have read the dev docs, so the documentation that is based on the current #main branch, not the latest release. You should be able to solve the errors you face by installing the #main branch instead of the tagged release v0.6

julia>] add https://github.com/SpeedyWeather/SpeedyWeather.jl#main

instead of julia>] add SpeedyWeather (which will install v0.6) you should then see in the package manager

julia>] status
...
  [9e226e20] SpeedyWeather v0.6.0 `https://github.com/SpeedyWeather/SpeedyWeather.jl#main`
...

Sorry! The examples in the documentation are only a few weeks old. Can you post here if you keep having problems running the examples?

[mini-DONG]

That works fine! Thanks! I installed the #main branch and errors was gone.
Next, I want to customize it. First, I want to add the initial filed, such as a specified velocity field (u,v component) at 500hPa. What should I do or what should I change?

[milankl]

The shallow water model does not have a vertical coordinate. Velocities here are assumed to be layer averages, whereby the layer starts at the surface (incl some orography) and has a thickness of $h$ (default 8.5km). If you want to set the initial conditions with velocity, note that the model uses vorticity and divergence, so you'll need to compute those first (I thought about an easier interface whereby you can directly set u,v and you don't have to worry about the initial steps)

1. Get your initial conditions onto one of the grids of our supported grids, e.g.

julia> u = randn(FullGaussianGrid,24)
julia> v = randn(FullGaussianGrid,24)

2. Compute the curl and divergence in spherical coordinates and spectral space. You have to divide both fields with the cos of latitude first, then convert to spectral. As the model uses one more degree of the spherical harmonics internally, you'll want those too

julia> RingGrids.scale_coslat⁻¹!(u)
julia> RingGrids.scale_coslat⁻¹!(v)
julia> us = spectral(u,one_more_degree=true)
julia> vs = spectral(u,one_more_degree=true)
julia> div = SpeedyTransforms.divergence(us,vs)
julia> curl = SpeedyTransforms.curl(us,vs)

Both div and curl are now divergence and curl in spectral space but scaled with the radius of the Earth.
3. Create a model and simulation as before

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

4. Now set the vor and div in the model (there's only one layer, and you always set the first time step) and unscale the radius

simulation.prognostic_variables.layers[1].timesteps[1].vor .= vor/spectral_grid.radius
simulation.prognostic_variables.layers[1].timesteps[1].div .= div/spectral_grid.radius

5. run

julia> run!(simulation,n_days=10)

Which yields with my random initial conditions

[mini-DONG]

It's very convenient! Thanks a lots! I wrote some codes to do my job by following yours:

1. Data loading: I want to use wind speed as initial condition. So, I interpolated the wind data into the lat/lon grid of the SpeedyWeather output. And this is my u/v800.nc in the following codes cell.

using NCDatasets
ds_u = NCDataset("G:/SpeedyWeather/test/test_IC/u800.nc")
ds_v = NCDataset("G:/SpeedyWeather/test/test_IC/v800.nc")
u = Matrix{Float32}(ds_u["ucomp"])
v = Matrix{Float32}(ds_v["vcomp"])

2. Data processing:

using SpeedyWeather
u = FullGaussianGrid(u)
v = FullGaussianGrid(v)
RingGrids.scale_coslat⁻¹!(u)
RingGrids.scale_coslat⁻¹!(v)
us = spectral(u,one_more_degree=true)
vs = spectral(v,one_more_degree=true) 
div = SpeedyTransforms.divergence(us,vs)
vor = SpeedyTransforms.curl(us,vs)

3. Model running:

spectral_grid = SpectralGrid(trunc=31,nlev=1)
initial_conditions = StartFromRest()
output = OutputWriter(spectral_grid,ShallowWater,
                      output_vars=[:u,:v])
model = ShallowWaterModel(;spectral_grid,initial_conditions,output=output)
simulation = initialize!(model)

simulation.prognostic_variables.layers[1].timesteps[1].vor .= vor/spectral_grid.radius
simulation.prognostic_variables.layers[1].timesteps[1].div .= div/spectral_grid.radius

run!(simulation,n_days=100,output=true)

My questions is:

• Whether my code is correct, including:

1. interpolating the wind data into the lat/lon grid of the SpeedyWeather output;
2. in u = Matrix{Float32}(ds_u["ucomp"]), should I use {Float32} or {Float64}, or it doesn't matter.
3. and any other mistakes.

• How do I add initial thickness $h$ like adding $u,v$, I want to match $h$ with $u,v$.

[milankl]

Btw thanks for doing this! Based on what users need and functionality they expect I'm planning to provide more convenient functions that will shortcut some of the things I've explained above.

So, I interpolated the wind data into the lat/lon grid of the SpeedyWeather output

What grid does your original data come on? SpeedyWeather is grid-flexible, in the documentation you can see what grids we support and in the end you just have to define what grid your data comes on, like FullGaussianGrid(u) what you did there. But you can start on any of our grids! If your original data comes on a grid we don't support you'd need to interpolate it first that is correct. But I reckon you probably had data on a regular lon-lat grid? Because in that cause you'd want to use the FullClenshawGrid instead.

in u = Matrix{Float32}(ds_u["ucomp"]), should I use {Float32} or {Float64}, or it doesn't matter.

In generally doesn't matter. SpeedyWeather uses Float32 as default, but you can use Float64 data too and it'll just use that. You set the number format that's actually used for the simulation in SpectralGrid(NF=Float32,...) if you start with Float64 data then you'll eventually convert it down. We should be fully number format flexible, if you do end up with an error that's related please post it.

How do I add initial thickness $h$ like adding $u,v$, I want to match $h$ with $u,v$.

I assume you mean by matching that it's in geostrophic balance? Technically there is no matching because u,v,h are all independent variables and so you can start with any combination of fields for these three. You can start with $h = H$, i.e. $\eta = 0$ an undisturbed interface displacement $\eta$ which would be happening if you follow the code from above. However, if you want geostrophic balance you'll need to solve for $\eta$ in

$$-fv = -g\partial_\lambda \eta, \quad fu = -g\partial_\theta \eta$$

which is conveniently if you take the divergence of these two equations (did I miss a minus?)

$$\zeta = \frac{g}{f}\nabla^2 \eta$$

So once you know vorticity you'll need to invert the Laplace operator to get $\eta$, I'll write down tomorrow how to do that, unless you figure it out till then ;)

[miniufo]

Hi @milankl, I guess I have the same need as @mini-DONG. It is like we need to start the SWM from some random pickups of 500 hPa wind and geopotential fields (e.g., from NCEP or ECMWF, regular lat/lon grids). The wind and mass fields are generally not in an exact geostrophic balance, as you suggested (also, it is problematic at equator because $f=0$). They should be close to this balance but ageostrophic motions are also presented.

[milankl]

If you have $u,v,\eta$ you can set them as described above. If you only have $u,v$ and want to provide a geostrophically balanced $\eta$ you can do the following (note that $f=0$ isn't actually a problem here because you go from $u,v$ to obtain $\eta$ not vice versa! So you don't divide by zero on the RHS but multiply by 0 on the left)

(The following requires PR #392 to be merged, or use mk/coriolis branch)

1. Create some data

julia> spectral_grid = SpectralGrid(trunc=31,nlev=1)
julia> forcing = SpeedyWeather.JetStreamForcing(spectral_grid)
julia> drag = QuadraticDrag(spectral_grid)
julia> model = ShallowWaterModel(;spectral_grid,forcing,drag)
julia> simulation = initialize!(model);
julia> run!(simulation,n_days=30)

2. Pretend you only have $u,v$ to get vorticity

julia> u = simulation.diagnostic_variables.layers[1].grid_variables.u_grid
julia> v = simulation.diagnostic_variables.layers[1].grid_variables.v_grid
julia> vor = SpeedyTransforms.curl(u,v) / spectral_grid.radius

Note that the operators in SpeedyTransforms all act on the unit sphere, so we divide by the radius

3. Get interface displacement $\eta$

julia> vor_grid = gridded(vor,Grid=spectral_grid.Grid);
julia> f = SpeedyWeather.coriolis(vor_grid);
julia> fζ_g = spectral_grid.Grid(vor_grid .* f ./ model.planet.gravity)

The last line is currently a bit awkward because .* operations with grids escape the grid and just return the underlying data, so we wrap that back into a grid here. I need to think about how to make something like this reliably more convenient. Apologies.

Now we invert the Laplace operator in spectral space and don't forget to account for the omitted radius scaling of that operation. You can now use $\eta$ to set the initial conditions (reminder to self: we should implement a general set! function like Oceananigans does it).

julia> fζ_g_spectral = spectral(fζ_g,one_more_degree=true);
julia> η = SpeedyTransforms.∇⁻²(fζ_g_spectral) * spectral_grid.radius^2
julia> η_grid = gridded(η,Grid=spectral_grid.Grid)

4. Verification

The velocity field we started from looks like

with an actual corresponding $\eta$ field, but we also now calculated $\eta$ using geostrophy, so how's that field different?

Strikingly similar! The remaining differences are the ageostrophic motions but also note that the mean is off! This is because geostrophy only use/defines the gradient of $\eta$ not the absolute values itself. Our geostrophy - $\eta$ has by construction a mean of zero (that's how we define the inverse Laplace operator) but the actual $\eta$ is some 1400m higher.

[miniufo]

If you only have and want to provide a geostrophically balanced you can do the following (note that isn't actually a problem here because you go from to obtain not vice versa! So you don't divide by zero on the RHS but multiply by 0 on the left)

Thanks @milankl for your detailed explanations. I made a mistake on this.

IIUC, you are talking about how to calculate $\eta$ from $u,v$, using existing functions from Speedy.

But sometimes we need $\eta$ as a initial condition. For example, we need to start the model from rest ($\vec u=0$) with a step function of $\eta$. This is a classical geostrophic adjustment problem that a jet will develop quickly from rest. Then in this idealized case how to turn the prescribed step-function $\eta$ into model state-variables (should be divergence and vorticity in the spectral model here)?

[mini-DONG]

@milankl thanks a lot! I followed your codes but an error was encountered at:

vor = SpeedyTransforms.curl(u,v)

with error:

MethodError: no method matching curl(::OctahedralGaussianGrid{Float32}, ::OctahedralGaussianGrid{Float32})

Also, I tried the given $u,v,\eta$ as initial condition. First, I download the 500 hPa ERA5 data about $u,v,\eta$. Then, I run a case:

using SpeedyWeather
spectral_grid = SpectralGrid(trunc=63,nlev=1)
orography = NoOrography(spectral_grid)
initial_conditions = ZonalJet()
model = ShallowWaterModel(;spectral_grid, orography, initial_conditions)
simulation = initialize!(model)
run!(simulation,n_days=6,output=true)

The purpose is to get the output's lon/lat grid. After running, I interpolated ERA5 data at SpeedyWeather's output lon/lat.
The interpolated $u,\eta$ is like:

Then, I processed the interpolated data like last talking:

using SpeedyWeather
u = FullGaussianGrid(u)
v = FullGaussianGrid(v)
eta = FullGaussianGrid(eta)
RingGrids.scale_coslat⁻¹!(u)
RingGrids.scale_coslat⁻¹!(v)
RingGrids.scale_coslat⁻¹!(eta)
us = spectral(u,one_more_degree=true)
vs = spectral(v,one_more_degree=true) 
etas = spectral(eta,one_more_degree=true) 

div = SpeedyTransforms.divergence(us,vs)
vor = SpeedyTransforms.curl(us,vs)

spectral_grid = SpectralGrid(trunc=63,nlev=1)
initial_conditions = StartFromRest()
output = OutputWriter(spectral_grid,ShallowWater,
                      output_vars=[:u,:v,:pres,:vor])
model = ShallowWaterModel(;spectral_grid,initial_conditions,output=output)
simulation = initialize!(model)

simulation.prognostic_variables.layers[1].timesteps[1].vor .= vor/spectral_grid.radius
simulation.prognostic_variables.layers[1].timesteps[1].div .= div/spectral_grid.radius
simulation.prognostic_variables.surface.timesteps[1].pres .= etas/spectral_grid.radius
run!(simulation,n_days=100,output=true)

Finally, I plotted the output first time of $u,\eta$:

It seem like $u$ is correct, but the $\eta$ is wrong. I want to know where I went wrong.

[milankl]

But sometimes we need η as a initial condition. For example, we need to start the model from rest (u→=0) with a step function of η. This is a classical geostrophic adjustment problem that a jet will develop quickly from rest. Then in this idealized case how to turn the prescribed step-function η into model state-variables (should be divergence and vorticity in the spectral model here)?

The you start from rest

initial_conditions = StartFromRest()

and set $\eta$ after model initialization

simulation = initialize!(model)

# say you have some data on a grid
eta = randn(FullClenshawGrid,24)

# transform to spectral, depending on input data resolution or model resolution you may need to
# use  the `spectral_truncation` function to get the LowerTriangularMatrix` into the right size
eta_spectral = spectral(eta)
# trunc = spectral_grid.trunc
# eta_spectral = spectral_truncation(eta_spectral,trunc+1,trunc)

simulation.prognostic_variables.surface.timesteps[1].pres .= eta_spectral
run!(simulation)

Reminder to self (or if anyone else would like to contribute something to SpeedyWeather!) that this might be easier with a

set!(simulation.prognostic_variables,η=my_eta_as_grid_or_spectral)

function.

[milankl]

MethodError: no method matching curl(::OctahedralGaussianGrid{Float32}, ::OctahedralGaussianGrid{Float32})

This is because you didn't use the mk/coriolis branch as mentioned above ☝🏼, I've merged #392 so you can also update #main branch now and it implements that method.

After running, I interpolated ERA5 data at SpeedyWeather's output lon/lat.

I believe ERA5 data comes on a regular lon-lat grid, which is almost identical to the FullClenshawGrid used here. You may need to throw away the points at the poles. But I don't think you actually need to interpolate.

RingGrids.scale_coslat⁻¹!(eta)

Only vectors need to be scaled by 1/coslat before taking a divergence or curl. Comment out this line.

div = SpeedyTransforms.divergence(us,vs)

With #392 merged you can also directly do div = SpeedyTransforms.divergence(u,v) with u,v gridded.

vor = SpeedyTransforms.curl(us,vs)

Same here.

simulation.prognostic_variables.layers[1].timesteps[1].vor .= vor/spectral_grid.radius
simulation.prognostic_variables.layers[1].timesteps[1].div .= div/spectral_grid.radius
simulation.prognostic_variables.surface.timesteps[1].pres .= etas/spectral_grid.radius

Don't divide $\eta$ by the radius, only divergence and vorticity have to be divided by the radius as the SpeedyTransform div,grad,curl operators act on the unit sphere. See

https://speedyweather.github.io/SpeedyWeather.jl/dev/barotropic/#scaling
https://speedyweather.github.io/SpeedyWeather.jl/dev/shallowwater/#scaled_swm
https://speedyweather.github.io/SpeedyWeather.jl/dev/spectral_transform/#Divergence-and-curl-in-spherical-harmonics

Maybe we should have a single clearer section on how to use SpeedyTransforms with the various aspects of scaling which are scattered across the documentation. I feel this question would come up often. If anyone here would like to give this section a go and create a pull request after what you learned, I'd be happy to review that! Would make it easier for anyone else who has similar questions!

It seem like u is correct, but the η is wrong. I want to know where I went wrong.

Don't scale eta by 1/coslat, don't divide by the radius and let me know if it's still off!

[milankl]

The example here is also added to the documentation with #397, I'll close this issue if that PR is merged. However, feel always free to add more comments, questions or open a new issue!