Single-Layer (Barotropic) Model

An implementation of pyqg.BTModel.

class pyqg_jax.bt_model.BTModel(*, nx=64, ny=None, L=1000000.0, W=None, rek=5.787e-07, filterfac=23.6, f=None, g=9.81, beta=0.0, rd=0.0, H=1.0, U=0.0, precision=Precision.SINGLE)

Single-layer (barotropic) quasigeostrophic model.

See also pyqg.BTModel.

Added in version 0.6.0.

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{m}^{-1}\ \mathrm{sec}^{-1}\).

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

• H (float, optional)

• U (float, optional) – Upper 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

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_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_grid()

Retrieve information on the model grid.

Added in version 0.8.0.

Returns:

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

Return type:

Grid

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