Diagnostics
Functions for computing diagnostics of simulation states.
The functions in this module can be used to compute diagnostic quantities such as kinetic energy or various spectra. See Diagnostics for examples of how to use them and plot the results.
Scalar Diagnostics
These diagnostic functions produce a single scalar value for a single snapshot.
- pyqg_jax.diagnostics.total_ke(full_state, grid)[source]
Compute the total kinetic energy in a single snapshot.
The density in the KE calculation is taken such that the entire model grid space has a mass of one unit. To use a different density value, multiply the result of this calculation by the total mass of the full space.
New in version 0.8.0.
- Parameters:
full_state (FullPseudoSpectralState) –
The state for which the kinetic energy is to be computed. This argument should be retrieved from a model, for example from
get_full_state().This function only operates on a single time step. To apply it to a trajectory use
jax.vmap().grid (Grid) – Information on the spatial grid for full_state. This should be retrieved from a model, for example from
get_grid().
- Returns:
The total kinetic energy for the provided snapshot.
- Return type:
- pyqg_jax.diagnostics.cfl(full_state, grid, ubg, dt)[source]
Calculate the CFL condition value for a single snapshot.
This computes the CFL condition value at each grid point in a given state. To report the worst value across the full state, aggregate them using
jnp.max.New in version 0.8.0.
- Parameters:
full_state (FullPseudoSpectralState) –
The state for which the CFL condition is to be checked. This argument should be retrieved from a model, for example from
get_full_state().This function only operates on a single time step. To apply it to a trajectory use
jax.vmap().grid (Grid) – Information on the spatial grid for full_state. This should be retrieved from a model, for example from
get_grid().ubg (jax.Array) – The model’s background velocity. Retrieve it from the same model as full_state, for example from
Ubg.dt (float) – The time step size. This should be retrieved from the relevant time stepper, for example form
dt.
- Returns:
The CFL condition value at each spatial grid location. These may optionally be aggregated with
jnp.max.- Return type:
Spectral Diagnostics
The diagnostics in this section produce array diagnostics which should
be averaged across a trajectory’s time
dimension, and then further processed with calc_ispec().
- pyqg_jax.diagnostics.ke_spec_vals(full_state, grid)[source]
Calculate the kinetic energy spectrum values for a snapshot.
The values produced by this function should be further processed by
calc_ispec()to produce the kinetic energy spectrum.New in version 0.8.0.
- Parameters:
full_state (FullPseudoSpectralState) –
The state for which the KE spectrum values should be computed. This argument should be retrieved from a model, for example from
get_full_state().This function only operates on a single time step. To apply it to a trajectory use
jax.vmap().grid (Grid) – Information on the spatial grid for full_state. This should be retrieved from a model, for example from
get_grid().
- Returns:
The KE spectrum values for the provided time step.
- Return type:
Note
The returned array should be treated as opaque. Values should only be averaged over any vmapped time dimensions, then passed to
calc_ispec().
Computing Spectra
These functions can be used to process the averaged values or a spectral diagnostic into an isotropic spectrum.
- pyqg_jax.diagnostics.calc_ispec(spec_vals, grid)[source]
Compute the isotropic spectrum from the given values.
The array spec_vals should have been computed by one of the spectral diagnostics functions–for example
ke_spec_vals().To correctly plot or interpret the spectrum computed by this function, use the result of
ispec_grid().New in version 0.8.0.
- Parameters:
spec_vals (jax.Array) – The input values which should be processed into an isotropic spectrum. These values should be a squared modulus of the Fourier coefficients.
grid (Grid) – The spatial grid over which the base values were defined. This should be retrieved from a model, for example from
get_grid().
- Returns:
A one-dimensional array providing the isotropic spectrum of spec_vals.
- Return type:
- pyqg_jax.diagnostics.ispec_grid(grid)[source]
Information on the spacing of values in an isotropic spectrum.
This function produces to results: iso_k and keep. The values iso_k are the isotropic wavenumbers for each entry in the result of calc_ispec. The result keep is an integer which should be used to slice the result of calc_ispec. Only the first keep entries should be interpreted.
These values from this function are useful when plotting the result of
calc_ispec().New in version 0.8.0.
- Parameters:
grid (Grid) – The spatial grid over which the base values were defined. This should be retrieved from a model, for example from
get_grid().- Returns:
iso_k (jax.Array) – The isotropic wavenumbers for each spectrum entry.
keep (int) – An integer indicating how many of the first spectrum entries should be interpreted or plotted.