Two-Layer QG Model

An implementation of pyqg.QGModel.

class pyqg_jax.qg_model.QGModel(*, nx=64, ny=None, L=1000000.0, W=None, rek=5.787e-07, filterfac=23.6, f=None, g=9.81, beta=1.5e-11, rd=15000.0, delta=0.25, H1=500, U1=0.025, U2=0.0, precision=Precision.SINGLE)

Two-layer quasi-geostrophic model.

See also pyqg.QGModel.

Parameters:

• nx (int, optional) – Number of grid points in the x direction.

• ny (int, optional) – Number of grid points in the y direction. Defaults to nx.

• L (float, optional) – Domain length in the x direction. Units: \(\mathrm{m}\).

• W (float, optional) – Domain length in the y direction. Defaults to L. Units: \(\mathrm{m}\).

• rek (float, optional) – Linear drag in lower layer. Units: \(\mathrm{sec}^{-1}\).

• filterfac (float, optional) – Amplitude of the spectral spherical filter.

• f (float, optional) –

• g (float, optional) –

• beta (float, optional) – Gradient of coriolis parameter. Units: \(\mathrm{sec}^{-1} \mathrm{m}^{-1}\).

• rd (float, optional) – Deformation radius. Units: \(\mathrm{m}\).

• delta (float, optional) – Layer thickness ratio \(\mathrm{H1}/\mathrm{H2}\). Unitless.

• H1 (float, optional) – Layer thickness

• U1 (float, optional) – Upper layer flow. Units: \(\mathrm{m}\ \mathrm{sec}^{-1}\).

• U2 (float, optional) – Lower layer flow. Units: \(\mathrm{m}\ \mathrm{sec}^{-1}\).

• precision (Precision, optional) – Precision of model computation. Selects dtype of state values.

Ubg

The background velocity for this model.

Type:

jax.Array

Note

This model internally uses 64-bit floating point values for part of its computation regardless of the chosen precision.

Make sure that JAX has 64-bit precision enabled.

create_initial_state(key)

Create a new initial state with random initialization.

Parameters:

key (jax.random.key) – The PRNG state used as the random key for initialization.

Returns:

The new state with random initialization.

Return type:

PseudoSpectralState

get_grid()

Retrieve information on the model grid.

New in version 0.8.0.

Returns:

A grid instance with attributes giving information on the spatial and spectral model grids.

Return type:

Grid

get_full_state(state)

Expand a partial state into a full state with all computed values.

Parameters:

state (PseudoSpectralState) – The partial state to be expanded.

Returns:

New state object with all computed fields derived from state.

Return type:

FullPseudoSpectralState

get_updates(state)

Get updates for time-stepping state.

Parameters:

state (PseudoSpectralState) – The state which will be time stepped using the computed updates.

Returns:

A new state object where each field corresponds to a time-stepping update to be applied.

Return type:

PseudoSpectralState

Note

The object returned by this function has the same type of state, but contains updates. This is so the time-stepping can be done by mapping over the states and updates as JAX pytrees with the same structure.

postprocess_state(state)

Apply fixed filtering to state.

This function should be called once on each new state after each time step.

SteppedModel handles this internally.

Parameters:

state (PseudoSpectralState) – The state to be filtered.

Returns:

The filtered state.

Return type:

PseudoSpectralState