dabench.data

Data generators and downloaders

Classes

`Barotropic`	Barotropic model data generator based on pyqg
`Data`	Base for all data generator objects.
`ENSOIndices`	Gets ENSO indices from CPC website
`GCP`	Loads ERA5 data from Google Cloud Platform
`Lorenz63`	Lorenz 63 model data generator.
`Lorenz96`	Lorenz 96 model data generator.
`PyQG`	PyQG quasi-geotropic model data generator.
`PyQGJax`	PyQGJax quasi-geotropic model data generator.
`QGS`	QGS quasi-geostrophic model data generator.
`SQGTurb`	SQGTurb model data generator.

Package Contents

class dabench.data.Barotropic(beta=0.0, rek=0.0, rd=0.0, H=1.0, nx=256, ny=None, L=2 * np.pi, x0=None, delta_t=0.001, taveint=1, ntd=1, time_dim=None, store_as_jax=False, **kwargs)

Barotropic model data generator based on pyqg

This data class is a wrapper of a “optional” pyqg package. See https://pyqg.readthedocs.io

Notes

DEPRECATED Uses default attribute values from pyqg.BTModel: https://pyqg.readthedocs.io/en/latest/api.html#pyqg.BTModel Those values originally come from Mcwilliams 1984:

J. C. Mcwilliams (1984). The emergence of isolated coherent vortices in turbulent flow. Journal of Fluid Mechanics, 146, pp 21-43 doi:10.1017/S0022112084001750.

Parameters:

system_dim – system dimension
beta (float) – Gradient of coriolis parameter. Units: meters^-1 * seconds^-1. Default is 0.
rek (float) – Linear drag in lower layer. Units: seconds^-1. Default is 0.
rd (float) – Deformation radius. Units: meters. Default is 0.
H (float) – Layer thickness. Units: meters. Default is 1.
nx (int) – Number of grid points in the x direction. Default is 256.
ny (int | None) – Number of grid points in the y direction. Default: nx.
L (float) – Domain length in x direction. Units: meters. Default is 2*pi.
x0 (ArrayLike | None) – the initial conditions. Can also be provided when initializing model object. If provided by both, the generate() arg takes precedence.
W – Domain width in y direction. Units: meters. Default: L.
filterfac (float) – amplitdue of the spectral spherical filter. Default is 23.6.
delta_t (float) – Numerical timestep. Units: seconds.
taveint (float) – Time interval for accumulation of diagnostic averages. For performance purposes, averaging does not have to occur every timestep. Units: seconds. Default is 1 (i.e. every 1000 timesteps when delta_t = 0.001)
ntd (int) – Number of threads to use. Should not exceed the number of cores on your machine.
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
time_dim (int | None)

__advance__()

Advances the QG model according to set attributes

Returns:: Array of absolute potential vorticity (relative potential vorticity + background vorticity).

forecast(n_steps=None, t_final=None, x0=None): Alias for self.generate(), except returns values as output

generate(n_steps=None, t_final=40, x0=None)

Generates values and times, saves them to the data object

Notes

Either provide n_steps or t_final in order to indicate the length of the forecast. These are used to set the values, times, and time_dim attributes.

Parameters:

n_steps (int | None) – Number of timesteps. Default is None, which sets
t_final/delta_t (n_steps to)
t_final (float) – Final time of trajectory. Default is 40, which results in n_steps = 40000
x0 (ArrayLike | None) – the initial conditions. Can also be provided when initializing model object. If provided by both, the generate() arg takes precedence.

Return type:

xarray.Dataset

class dabench.data.Data(system_dim=3, original_dim=None, random_seed=37, delta_t=0.01, store_as_jax=False, x0=None, **kwargs)

Base for all data generator objects.

Parameters:

system_dim (int) – system dimension
original_dim (tuple[int, Ellipsis] | None) – dimensions in original space, e.g. could be 3x3 for a 2d system with system_dim = 9. Defaults to (system_dim), i.e. 1d.
random_seed (int) – random seed, defaults to 37
delta_t (float) – the timestep of the data (assumed uniform)
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
x0 (ArrayLike | None)

calc_lyapunov_exponents_final(total_time=None, rescale_time=1, convergence=0.01, x0=None)

Computes the final Lyapunov Exponents

Notes

See self.calc_lyapunov_exponents_series for full info

Parameters:

total_time (float | None) – Time to integrate over to compute LEs. Usually there’s a tradeoff between accuracy and computation time (more total_time leads to higher accuracy but more computation time). Default depends on model type and are based roughly on how long it takes for satisfactory convergence: For Lorenz63: n_steps=15000 (total_time=150 for delta_t=0.01) For Lorenz96: n_steps=50000 (total_time=500 for delta_t=0.01)
rescale_time (float) – Time for when the algorithm rescales the propagator to reduce the exponential growth in errors. Default is 1 (i.e. 100 timesteps when delta_t = 0.01).
convergence (float) – Prints warning if LE convergence is below this number. Default is 0.01.
x0 (ArrayLike | None) – initial condition to start computing LE. Needs to be on the attractor (i.e., remove transients). Default is None, which will fallback to use the x0 set during model object initialization.

Returns:

Lyapunov exponents array of size (system_dim)

Return type:

ArrayLike

calc_lyapunov_exponents_series(total_time=None, rescale_time=1, convergence=0.01, x0=None)

Computes the spectrum of Lyapunov Exponents.

Notes

Lyapunov exponents help describe the degree of “chaos” in the model. Make sure to plot the output to check that the algorithm converges. There are three ways to make the estimate more accurate:

Decrease the delta_t of the model

Increase total_time

Decrease rescale time (try increasing total_time first)

Algorithm: Eckmann 85, https://www.ihes.fr/~ruelle/PUBLICATIONS/%5B81%5D.pdf pg 651 This method computes the full time series of Lyapunov Exponents, which is useful for plotting for debugging. To get only the final Lyapunov Exponent, use self.calc_lyapunov_exponents.

Parameters:

total_time (float | None) – Time to integrate over to compute LEs. Usually there’s a tradeoff between accuracy and computation time (more total_time leads to higher accuracy but more computation time). Default depends on model type and are based roughly on how long it takes for satisfactory convergence: For Lorenz63: n_steps=15000 (total_time=150 for delta_t=0.01) For Lorenz96: n_steps=50000 (total_time=500 for delta_t=0.01)
rescale_time (float) – Time for when the algorithm rescales the propagator to reduce the exponential growth in errors. Default is 1 (i.e. 100 timesteps when delta_t = 0.01).
convergence (float) – Prints warning if LE convergence is below this number. Default is 0.01.
x0 (ArrayLike | None) – initial condition to start computing LE. Needs to be on the attractor (i.e., remove transients). Default is None, which will fallback to use the x0 set during model object initialization.

Returns:

Lyapunov exponents for all timesteps, array of size (total_time/rescale_time - 1, system_dim)

Return type:

ArrayLike

generate(n_steps=None, t_final=None, x0=None, M0=None, return_tlm=False, stride=None, **kwargs)

Generates a dataset and returns xarray state vector.

Notes

Either provide n_steps or t_final in order to indicate the length of the forecast.

Parameters:

n_steps (int | None) – Number of timesteps. One of n_steps OR t_final must be specified.
t_final (float | None) – Final time of trajectory. One of n_steps OR t_final must be specified.
x0 (ArrayLike | None) – initial conditions state vector of shape (system_dim)
M0 (ArrayLike | None) – the initial condition of the TLM matrix computation shape (system_dim, system_dim).
return_tlm (bool) – specifies whether to compute and return the integrated Jacobian as a Tangent Linear Model for each timestep.
stride (int | None) – specify how many steps to skip in the output data versus the model timestep (delta_t).
**kwargs – arguments to the integrate function (permits changes in convergence tolerance, etc.).

Returns:

Xarray Dataset of output vector, and if return_tlm=True Xarray DataArray of TLMs corresponding to the system trajectory.

Return type:

xarray.Dataset | tuple[xarray.Dataset | xarray.DataArray]

load_netcdf(filepath=None, include_vars=None, exclude_vars=None, years_select=None, dates_select=None)

Loads values from netCDF file, saves them in values attribute

Parameters:

filepath (str | None) – Path to netCDF file to load. If not given, defaults to loading ERA5 ECMWF SLP data over Japan from 2018 to 2021.
include_vars (list | ArrayLike | None) – Data variables to load from NetCDF. If None (default), loads all variables. Can be used to exclude bad variables.
exclude_vars (list | ArrayLike | None) – Data variabes to exclude from NetCDF loading. If None (default), loads all vars (or only those specified in include_vars). It’s recommended to only specify include_vars OR exclude_vars (unless you want to do extra typing).
years_select (list | ArrayLike | None) – Years to load (ints). If None, loads all timesteps.
dates_select (list | ArrayLike | None) – Dates to load. Elements must be datetime date or datetime objects, depending on type of time indices in NetCDF. If both years_select and dates_select are specified, time_stamps overwrites “years” argument. If None, loads all timesteps.

Return type:

xarray.Dataset

rhs_aux(x, t)

The auxiliary model used to compute the TLM.

Parameters:

x (ArrayLike) – State vector with size (system_dim)
t (ArrayLike) – Array of times with size (time_dim)

Returns:

State vector [size: (system_dim,)]

Return type:

dxaux (ndarray)

save_netcdf(ds, filename)

Saves values in values attribute to netCDF file

Parameters:

ds (xarray.Dataset) – Xarray dataset
filepath – Path to netCDF file to save
filename (str)

class dabench.data.ENSOIndices(file_dict=None, var_types=None, system_dim=None, store_as_jax=False, **kwargs)

Gets ENSO indices from CPC website

Notes

Source: https://www.cpc.ncep.noaa.gov/data/indices/

Parameters:

system_dim (int | None) – system dimension
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
file_dict (dict | None) –

Lists of files to get. Dict keys are type of data:
’wnd’: Wind ‘slp’: Sea level pressure ‘soi’: Southern Oscillation Index ‘soi3m’ Southerm Oscillation Index 3-month running mean ‘sst’: SST indices ‘desst’: Detrended Nino3.4 index ‘rsst’: Regional SST index (North and South Atlantic, and

Tropics)

’olr’: Outgoing long-wave radiance ‘cpolr’: Central Pacific OLR

Dict values are individual files from the website, see full list at
https://www.cpc.ncep.noaa.gov/data/indices/

Default is {‘wnd’: [‘zwnd200’], ‘slp’: [‘darwin’]}
var_types (dict | None) –
List of variables within file to get. Dict keys are type of data (see list in file_dict description). Dict values are type of variable:

’ori’ = Original ‘ano’ = Anomaly ‘std’ = Standardized Anomaly

sst (sea surface temp) has nino3, nino34, and nino12 options, see
for more info: https://climatedataguide.ucar.edu/climate-data/nino-sst-indices-nino-12-3-34-4-oni-and-tni

rsst (regional sea surface temp) has prefixes: ‘sa_’ is
South Atlantic, ‘na_’ is North Atlantic, ‘tr_’ is Tropics.

Default depends on data type (first string in each value list)
’wnd’=’ori’, ‘lp’=’ori’, ‘soi’=’ano’, ‘soi3m’=’ori’, ‘eqsoi’=’std’, ‘sst’=’nino12’, ‘desst’=’ori’, ‘rsst’=’na’, ‘olr’=’ori’, ‘cpolr’=’ano’

generate()

Alias for _load_gcp_era5

Return type:: xarray.Dataset

class dabench.data.GCP(variables=['2m_temperature'], date_start='2020-01-01', date_end='2020-12-31', min_lat=19.8554808619, max_lat=23.1886107447, min_lon=-84.9749110583, max_lon=-74.1780248685, store_as_jax=False, **kwargs)

Loads ERA5 data from Google Cloud Platform

Notes

Source: https://cloud.google.com/storage/docs/public-datasets/era5 Data is hourly

Parameters:

variables (list[str]) – Names of ERA5 variables to load. For description of variables, see: https://github.com/google-research/arco-era5?tab=readme-ov-file#full_37-1h-0p25deg-chunk-1zarr-v3 Default is [‘2m_temperature’] (Air temperature at 2 metres)
date_start (str) – Start of time range to download, in ‘yyyy-mm-dd’ format. Can also just specify year (‘yyyy’) or year and month (‘yyyy-mm’). Default is ‘2020-06-01’.
date_end (str) – End of time range to download, in ‘yyyy-mm-dd’ format. Can also just specify year (‘yyyy’) or year and month (‘yyyy-mm’). Default is ‘2020-9-30’.
min_lat (float) – Minimum latitude for bounding box. If None, loads global data (which can be VERY large). Bounding box default covers Cuba.
max_lat (float) – Max latitude for bounding box (see min_lat for info).
min_lon (float) – Min latitude for bounding box (see min_lat for info).
max_lon (float) – Max latitude for bounding box (see min_lat for info).
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).

generate()

Alias for _load_gcp_era5

Return type:: xarray.Dataset

load()

Alias for _load_gcp_era5

Return type:: xarray.Dataset

class dabench.data.Lorenz63(sigma=10.0, rho=28.0, beta=8.0 / 3.0, delta_t=0.01, x0=jnp.array([-10.0, -15.0, 21.3]), system_dim=3, values=None, store_as_jax=False, **kwargs)

Lorenz 63 model data generator.

Parameters:

sigma (float) – Lorenz 63 param. Default is 10., the original value used in Lorenz, 1963. https://doi.org/10.1175/1520-0469(1963)020<0130:DNF>2.0.CO;2
rho (float) – Lorenz 63 param. Default is 28., the value used in Lorenz, 1963 (see DOI above)
beta (float) – Lorenz 63 param. Default is 8./3., the value used in Lorenz, 1963 (see DOI above)
delta_t (float) – length of one time step
x0 (ArrayLike | None) – Initial state, array of floats of size (system_dim). Default is jnp.array([-10.0, -15.0, 21.3]), which is the system state after a 6000 step spinup with delta_t=0.01 and initial conditions [0., 1., 0.], a spinup which replicates the simulation described in Lorenz, 1963.
system_dim (int) – system dimension. Must be 3 for Lorenz63.
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
values (ArrayLike | None)

Jacobian(x)

Jacobian of the L63 system

Parameters:: x (ArrayLike) – state vector with shape (system_dim)
Returns:: ndarray of Jacobian matrix with shape (system_dim, system_dim)
Return type:: jax.Array

rhs(x, t=None)

vector field of Lorenz 63

Parameters:

x (ArrayLike) – state vector with shape (system_dim)
t (ArrayLike | None) – times vector. Needed as argument slot for odeint but unused

Returns:

vector field of Lorenz 63 with shape (system_dim)

Return type:

jax.Array

class dabench.data.Lorenz96(forcing_term=8.0, delta_t=0.05, x0=None, system_dim=36, values=None, store_as_jax=False, **kwargs)

Lorenz 96 model data generator.

Notes

Default values come from Lorenz, 1996: eapsweb.mit.edu/sites/default/files/Predicability_a_Problem_2006.pdf

Parameters:

forcing_term (float) – Forcing constant for Lorenz96, prevents energy from decaying to 0. Default is 8.0.
x0 (ArrayLike | None) – Initial state vector, array of floats of size (system_dim). For system_dim of 5, 6, or 36, defaults are the final state of a 14400 timestep spinup starting with an initial state (x0) of all 0s except the first element, which is set to 0.01. This setup is taken from Lorenz, 1996. Defaults are provided from spinups usin delta_t = 0.05 (Lorenz, 1996 original value) and delta_t = 0.01 (more frequently used today) For all other system_dim settings, default is all 0s except the first element, which is set to 0.01.
system_dim (int) – System dimension, must be between 4 and 40. Default is 36.
delta_t (float) – Length of one time step. Default is 0.05 from Lorenz, 1996, but on modern computers 0.01 is often used.
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
values (ArrayLike | None)

Jacobian(x)

Computes the Jacobian of the Lorenz96 system

Parameters:: x (ArrayLike) – state vector with shape (system_dim)
Returns:: ndarray of Jacobian matrix with shape (system_dim, system_dim)
Return type:: jax.Array

rhs(x, t=None)

Computes vector field of Lorenz 96

Parameters:

x (ArrayLike) – state vector with shape (system_dim)
t (ArrayLike | None) – times vector. Required as argument slot for some numerical integrators but unused

Returns:

vector field of Lorenz 96 with shape (system_dim)

Return type:

jax.Array

class dabench.data.PyQG(beta=1.5e-11, rd=15000.0, delta=0.25, H1=500, H2=None, U1=0.025, U2=0.0, x0=None, twrite=10000, nx=64, ny=None, delta_t=7200, ntd=1, time_dim=None, values=None, times=None, store_as_jax=False, **kwargs)

PyQG quasi-geotropic model data generator.

The PyQG class is simply a wrapper of a “optional” pyqg package. See https://pyqg.readthedocs.io

Notes

DEPRECATED Uses default attribute values from pyqg.QGModel: https://pyqg.readthedocs.io/en/latest/api.html#pyqg.QGModel

Parameters:

beta (float) – Gradient of coriolis parameter. Units: meters^-1 * seconds^-1
rek (float) – Linear drag in lower layer. Units: seconds^-1
rd (float) – Deformation radius. Units: meters.
delta (float) – Layer thickness ratio (H1/H2)
U1 (float) – Upper layer flow. Units: m/s
U2 (float) – Lower layer flow. Units: m/s
H1 (float) – Layer thickness (sets both H1 and H2).
nx (int) – Number of grid points in the x direction.
ny (int) – Number of grid points in the y direction (default: nx).
L (float) – Domain length in x direction. Units: meters.
W (float) – Domain width in y direction. Units: meters (default: L).
filterfac (float) – amplitude of the spectral spherical filter (originally 18.4, later changed to 23.6).
delta_t (float) – Numerical timestep. Units: seconds.
twrite (int) – Interval for cfl writeout. Units: number of timesteps.
tmax (float) – Total time of integration (overwritten by t_final). Units: seconds.
ntd (int) – Number of threads to use. Should not exceed the number of cores on your machine.
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).

__advance__()

Advances the QG model according to set attributes

Returns:: Array of absolute potential vorticity (relative potential vorticity + background vorticity).

forecast(n_steps=None, t_final=None, x0=None): Alias for self.generate(), except returns values as output

generate(n_steps=None, t_final=None, x0=None)

Generates values and times, saves them to the data object

Notes

Either provide n_steps or t_final in order to indicate the length of the forecast. These are used to set the values, times, and time_dim attributes.

Parameters:

n_steps (int) – Number of timesteps. One of n_steps OR t_final must be specified.
t_final (float) – Final time of trajectory. One of n_steps OR t_final must be specified.
x0 (ndarray, optional) – the initial conditions. Can also be provided when initializing model object. If provided by both, the generate() arg takes precedence.

class dabench.data.PyQGJax(beta=1.5e-11, rd=15000.0, delta=0.25, H1=500, H2=None, U1=0.025, U2=0.0, x0=None, nx=64, ny=None, delta_t=7200, random_seed=37, store_as_jax=False, **kwargs)

PyQGJax quasi-geotropic model data generator.

The PyQGJax class is simply a wrapper of the “optional” pyqg-jax package. See https://pyqg-jax.readthedocs.io

Notes

Uses default attribute values from pyqg_jax.QGModel: https://pyqg.readthedocs.io/en/latest/api.html#pyqg.QGModel

Parameters:

beta (float) – Gradient of coriolis parameter. Units: meters^-1 * seconds^-1
rd (float) – Deformation radius. Units: meters.
delta (float) – Layer thickness ratio (H1/H2)
H1 (float) – Layer thickness (sets both H1 and H2 if H2 not specified).
H2 (float | None) – Layer 2 thickness.
U1 (float) – Upper layer flow. Units: m/s
U2 (float) – Lower layer flow. Units: m/s
x0 (ArrayLike | None) – the initial conditions. Can also be provided when initializing model object. If provided by both, the generate() arg takes precedence.
nx (int) – Number of grid points in the x direction.
ny (int | None) – Number of grid points in the y direction (default: nx).
delta_t (float) – Numerical timestep. Units: seconds.
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
random_seed (int)

forecast(n_steps=None, t_final=None, x0=None): Alias for self.generate(), except returns values as output

generate(n_steps=None, t_final=40, x0=None)

Generates values and times, saves them to the data object

Notes

Either provide n_steps or t_final in order to indicate the length of the forecast.

Parameters:

n_steps (int | None) – Number of timesteps. One of n_steps OR t_final must be specified.
t_final (float) – Final time of trajectory. One of n_steps OR t_final must be specified.
x0 (ArrayLike | None) – the initial conditions. Can also be provided when initializing model object. If provided by both, the generate() arg takes precedence.

Return type:

xarray.Dataset

class dabench.data.QGS(model_params=None, x0=None, delta_t=0.1, system_dim=None, store_as_jax=False, random_seed=37, **kwargs)

QGS quasi-geostrophic model data generator.

The QGS class is simply a wrapper of an optional qgs package. See https://qgs.readthedocs.io/

Parameters:

model_params (qgs.params.params.QgParams | None) – qgs parameter object. See: https://qgs.readthedocs.io/en/latest/files/technical/configuration.html#qgs.params.params.QgParams If None, will use defaults specified by: De Cruz, et al. (2016). Geosci. Model Dev., 9, 2793-2808.
x0 (ArrayLike | None) – Initial state vector, array of floats. Default is:
delta_t (ArrayLike | None) – Numerical timestep. Units: seconds.
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
system_dim (int | None)
random_seed (int)

Jacobian(x, t=0)

Jacobian of the qgs system

Parameters:

x (ArrayLike) – State vector of size (system_dim)
t (float | None) – times vector. Required as argument slot for some numerical integrators but unused.

Returns:

Jacobian matrix of size (system_dim, system_dim)

Return type:

numpy.ndarray

calc_lyapunov_exponents_series(total_time=None, rescale_time=1, convergence=0.01, x0=None)

Computes the spectrum of Lyapunov Exponents.

Notes

Lyapunov exponents help describe the degree of “chaos” in the model. Make sure to plot the output to check that the algorithm converges. There are three ways to make the estimate more accurate:

Decrease the delta_t of the model

Increase total_time

Decrease rescale time (try increasing total_time first)

Algorithm: Eckmann 85, https://www.ihes.fr/~ruelle/PUBLICATIONS/%5B81%5D.pdf pg 651 This method computes the full time series of Lyapunov Exponents, which is useful for plotting for debugging. To get only the final Lyapunov Exponent, use self.calc_lyapunov_exponents.

Parameters:

total_time (float | None) – Time to integrate over to compute LEs. Usually there’s a tradeoff between accuracy and computation time (more total_time leads to higher accuracy but more computation time). Default depends on model type and are based roughly on how long it takes for satisfactory convergence: For Lorenz63: n_steps=15000 (total_time=150 for delta_t=0.01) For Lorenz96: n_steps=50000 (total_time=500 for delta_t=0.01)
rescale_time (float) – Time for when the algorithm rescales the propagator to reduce the exponential growth in errors. Default is 1 (i.e. 100 timesteps when delta_t = 0.01).
convergence (float) – Prints warning if LE convergence is below this number. Default is 0.01.
x0 (ArrayLike | None) – initial condition to start computing LE. Needs to be on the attractor (i.e., remove transients). Default is None, which will fallback to use the x0 set during model object initialization.

Returns:

Lyapunov exponents for all timesteps, array of size (total_time/rescale_time - 1, system_dim)

Return type:

ArrayLike

generate(n_steps=None, t_final=None, x0=None, M0=None, return_tlm=False, stride=None, **kwargs)

Generates a dataset and assigns values and times to the data object.

Notes

Either provide n_steps or t_final in order to indicate the length of the forecast.

Parameters:

n_steps (int) – Number of timesteps. One of n_steps OR t_final must be specified.
t_final (float) – Final time of trajectory. One of n_steps OR t_final must be specified.
M0 (ndarray) – the initial condition of the TLM matrix computation shape (system_dim, system_dim).
return_tlm (bool) – specifies whether to compute and return the integrated Jacobian as a Tangent Linear Model for each timestep.
x0 (ndarray) – initial conditions state vector of shape (system_dim)
stride (int) – specify how many steps to skip in the output data versus the model timestep (delta_t).
**kwargs – arguments to the integrate function (permits changes in convergence tolerance, etc.).

Returns:

Xarray Dataset of output vector, and if return_tlm=True Xarray DataArray of TLMs corresponding to the system trajectory.

Return type:

xarray.Dataset | tuple[xarray.Dataset | xarray.DataArray]

rhs(x, t=0)

Vector field (tendencies) of qgs system

Parameters:

x (ArrayLike) – State vector of size (system_dim)
t (float | None) – times vector. Required as argument slot for some numerical integrators but unused.

Returns:

Vector field of qgs

Return type:

numpy.ndarray

rhs_aux(x, t)

The auxiliary model used to compute the TLM.

Parameters:

x (ArrayLike) – State vector with size (system_dim)
t (ArrayLike) – Array of times with size (time_dim)

Returns:

State vector of size (system_dim,)

Return type:

jax.Array

class dabench.data.SQGTurb(pv=None, f=0.0001, nsq=0.0001, L=20000000.0, H=10000.0, U=30.0, r=0.0, tdiab=10.0 * 86400, diff_order=8, diff_efold=86400.0 / 3, symmetric=True, dealias=True, precision='single', tstart=0, delta_t=900, store_as_jax=False, **kwargs)

SQGTurb model data generator.

Parameters:

pv (ArrayLike | None) – Potential vorticity array. If None (default), loads data from 57600 step spinup with initial conditions taken from Jeff Whitaker’s original implementation: https://github.com/jswhit/sqgturb. 57600 steps matches the “nature run” spin up in that repository.
system_dim – The dimension of the system state
delta_t (float) – model time step (seconds)
x0 – Initial state, array of floats of size (system_dim).
f (float) – coriolis
nqr – Brunt-Vaisalla (buoyancy) freq squared
L (float) – size of square domain
H (float) – height of upper boundary
U (float) – basic state velocity at z = H
r (float) – Ekman damping (at z=0)
tdiab (float) – thermal relaxation damping
diff_order (int) – hyperdiffusion order
diff_efold (float) – hyperdiff time scale
symmetric (bool) – symmetric jet, or jet with U=0 at sf
dealias (bool) – if True, dealiasing applied using 2/3 rule
precision (str) – ‘single’ or ‘double’. Default is ‘single’
tstart (float) – initialize time counter
delta_t – the timestep of the data (assumed uniform)
store_as_jax (bool) – Store values as jax array instead of numpy array. Default is False (store as numpy).
nsq (float)

fft2(pv)

Alias method for FFT of PV

Parameters:: pv (jax.Array)
Return type:: jax.Array

fft2_2dto1d(pv)

Runs FFT then maps from 2D to 1D

Parameters:: pv (jax.Array)
Return type:: jax.Array

ifft2(pvspec)

Alias method for inverse FFT of PV Spectral

Parameters:: pvspec (jax.Array)
Return type:: jax.Array

ifft2_2dto1d(pvspec)

Runs inverse FFT then maps from 2D to 1D

Parameters:: pvspec (jax.Array)
Return type:: jax.Array

integrate(f, x0, t_final, delta_t=None, include_x0=True, t=None, **kwargs)

Advances pv forward number of timesteps given by self.n_steps.

Note

If pv not specified, use pvspec instance variable.

Parameters:

f (function) – right hand side (rhs) of the ODE. Not used, but needed to function with generate() from _data.Data().
x0 (ndarray) – potential vorticity (pvspec) initial condition in spectral space
t_final (float)
delta_t (float | None)
include_x0 (bool)
t (float | None)

Return type:

tuple[jax.Array, jax.Array]

map1dto2d(x)

Maps 1D state vector to 2D PV

Parameters:: x (jax.Array)
Return type:: jax.Array

map1dto2d_fft2(x)

Maps for 1D to 2D then runs FFT

Parameters:: x (jax.Array)
Return type:: jax.Array

map1dto2d_ifft2(x)

Maps for 1D to 2D then runs inverse FFT

Parameters:: x (jax.Array)
Return type:: jax.Array

map2dto1d(pv)

Maps 2D PV to 1D system state

Parameters:: pv (jax.Array)
Return type:: jax.Array

rhs(pvspec)

Computes pv deriv on z=0, inverts pv to get streamfunction.

Parameters:: pvspec (ArrayLike)
Return type:: jax.Array