API Reference

The functionality of pyqg-jax is split into several modules.

Generally, the operation of each of the base models is the same as the equivalent in PyQG so its documentation can be used as a reference.

However as part of the port, some changes to the overall structure have been made. These changes both improve compatibility with JAX and separate portions of the model that research projects may wish to modify (in particular, the time-stepping).

The expected steps to construct a model and step it forward are:

  1. Select a base model

  2. Apply a parameterization (optional)

  3. Select a time stepper

  4. Combine into a SteppedModel

New states can then be randomly initialized, and the model can be stepped forward in time.

Contents