dabench.dacycler

Data Assimilation cyclers

Classes

DACycler

Base for all DACyclers

ETKF

Ensemble transform Kalman filter DA Cycler

Var3D

3D-Var DA Cycler

Var4D

4D-Var DA Cycler

Var4DBackprop

Backpropagation 4D-Var DA Cycler

Package Contents

class dabench.dacycler.DACycler(system_dim, delta_t, model_obj, B=None, R=None, H=None, h=None)

Base for all DACyclers

Parameters:

  • system_dim (int) – System dimension

  • delta_t (float) – The timestep of the model (assumed uniform)

  • model_obj (dabench.model.Model) – Forecast model object.

  • B (ArrayLike | None) – Initial / static background error covariance. Shape: (system_dim, system_dim). If not provided, will be calculated automatically.

  • R (ArrayLike | None) – Observation error covariance matrix. Shape (obs_dim, obs_dim). If not provided, will be calculated automatically.

  • H (ArrayLike | None) – Observation operator with shape: (obs_dim, system_dim). If not provided will be calculated automatically.

  • h (Callable | None) – Optional observation operator as function. More flexible (allows for more complex observation operator). Default is None.

cycle(input_state, start_time, obs_vector, n_cycles, obs_error_sd=None, analysis_window=0.2, analysis_time_in_window=None, return_forecast=False)

Perform DA cycle repeatedly, including analysis and forecast

Parameters:

  • input_state (XarrayDatasetLike) – Input state as a Xarray Dataset

  • start_time (float | numpy.datetime64) – Starting time.

  • obs_vector (XarrayDatasetLike) – Observations vector.

  • n_cycles (int) – Number of analysis cycles to run, each of length analysis_window.

  • analysis_window (float) – Time window from which to gather observations for DA Cycle.

  • analysis_time_in_window (float | None) – Where within analysis_window to perform analysis. For example, 0.0 is the start of the window. Default is None, which selects the middle of the window.

  • return_forecast (bool) – If True, returns forecast at each model timestep. If False, returns only analyses, one per analysis cycle.

  • obs_error_sd (float | ArrayLike | None)

Return type:

XarrayDatasetLike

class dabench.dacycler.ETKF(system_dim, delta_t, model_obj, B=None, R=None, H=None, h=None, ensemble_dim=4, multiplicative_inflation=1.0)

Ensemble transform Kalman filter DA Cycler

Parameters:

  • system_dim (int) – System dimension.

  • delta_t (float) – The timestep of the model (assumed uniform)

  • model_obj (dabench.model.Model) – Forecast model object.

  • B (ArrayLike | None) – Initial / static background error covariance. Shape: (system_dim, system_dim). If not provided, will be calculated automatically.

  • R (ArrayLike | None) – Observation error covariance matrix. Shape (obs_dim, obs_dim). If not provided, will be calculated automatically.

  • H (ArrayLike | None) – Observation operator with shape: (obs_dim, system_dim). If not provided will be calculated automatically.

  • h (Callable | None) – Optional observation operator as function. More flexible (allows for more complex observation operator). Default is None.

  • ensemble_dim (int) – Number of ensemble instances for ETKF. Default is 4. Higher ensemble_dim increases accuracy but has performance cost.

  • multiplicative_inflation (float) – Scaling factor by which to multiply ensemble deviation. Default is 1.0 (no inflation).

class dabench.dacycler.Var3D(system_dim, delta_t, model_obj, B=None, R=None, H=None, h=None)

3D-Var DA Cycler

Parameters:

  • system_dim (int) – System dimension.

  • delta_t (float) – The timestep of the model (assumed uniform)

  • model_obj (dabench.model.Model) – Forecast model object.

  • B (ArrayLike | None) – Initial / static background error covariance. Shape: (system_dim, system_dim). If not provided, will be calculated automatically.

  • R (ArrayLike | None) – Observation error covariance matrix. Shape (obs_dim, obs_dim). If not provided, will be calculated automatically.

  • H (ArrayLike | None) – Observation operator with shape: (obs_dim, system_dim). If not provided will be calculated automatically.

  • h (Callable | None) – Optional observation operator as function. More flexible (allows for more complex observation operator). Default is None.

class dabench.dacycler.Var4D(system_dim, delta_t, model_obj, B=None, R=None, H=None, h=None, solver='bicgstab', n_outer_loops=1, steps_per_window=1, obs_window_indices=None, **kwargs)

4D-Var DA Cycler

Parameters:

  • system_dim (int) – System dimension.

  • delta_t (float) – The timestep of the model (assumed uniform)

  • model_obj (dabench.model.Model) – Forecast model object.

  • B (ArrayLike | None) – Initial / static background error covariance. Shape: (system_dim, system_dim). If not provided, will be calculated automatically.

  • R (ArrayLike | None) – Observation error covariance matrix. Shape (obs_dim, obs_dim). If not provided, will be calculated automatically.

  • H (ArrayLike | None) – Observation operator with shape: (obs_dim, system_dim). If not provided will be calculated automatically.

  • h (Callable | None) – Optional observation operator as function. More flexible (allows for more complex observation operator). Default is None.

  • solver (str) – Name of solver to use. Default is ‘bicgstab’.

  • n_outer_loops (int) – Number of times to run through outer loop over 4DVar. Increasing this may result in higher accuracy but slower performance. Default is 1.

  • steps_per_window (int) – Number of timesteps per analysis window. If None (default), will calculate automatically based on delta_t and .cycle() analysis_window length.

  • obs_window_indices (ArrayLike | None) – Timestep indices where observations fall within each analysis window. For example, if analysis window is 0 - 0.05 with delta_t = 0.01 and observations fall at 0, 0.01, 0.02, 0.03, 0.04, and 0.05, obs_window_indices = [0, 1, 2, 3, 4, 5]. If None (default), will calculate automatically.

class dabench.dacycler.Var4DBackprop(system_dim, delta_t, model_obj, B=None, R=None, H=None, h=None, learning_rate=0.5, lr_decay=0.5, num_iters=3, steps_per_window=None, obs_window_indices=None, loss_growth_limit=10, **kwargs)

Backpropagation 4D-Var DA Cycler

Parameters:

  • system_dim (int) – System dimension.

  • delta_t (float) – The timestep of the model (assumed uniform)

  • model_obj (dabench.model.Model) – Forecast model object.

  • B (ArrayLike | None) – Initial / static background error covariance. Shape: (system_dim, system_dim). If not provided, will be calculated automatically.

  • R (ArrayLike | None) – Observation error covariance matrix. Shape (obs_dim, obs_dim). If not provided, will be calculated automatically.

  • H (ArrayLike | None) – Observation operator with shape: (obs_dim, system_dim). If not provided will be calculated automatically.

  • h (Callable | None) – Optional observation operator as function. More flexible (allows for more complex observation operator). Default is None.

  • num_iters (int) – Number of iterations for backpropagation per analysis cycle. Default is 3.

  • steps_per_window (int | None) – Number of timesteps per analysis window. If None (default), will calculate automatically based on delta_t and .cycle() analysis_window length.

  • learning_rate (float) – LR for backpropogation. Default is 0.5, but DA results can be quite sensitive to this parameter.

  • lr_decay (float) – Exponential learning rate decay. If set to 1, no decay. Default is 0.5.

  • obs_window_indices (ArrayLike | list | None) – Timestep indices where observations fall within each analysis window. For example, if analysis window is 0 - 0.05 with delta_t = 0.01 and observations fall at 0, 0.01, 0.02, 0.03, 0.04, and 0.05, obs_window_indices = [0, 1, 2, 3, 4, 5]. If None (default), will calculate automatically.

  • loss_growth_limit (float) – If loss grows by more than this factor during one analysis cycle, JAX will cut off computation and return an error. This prevents it from hanging indefinitely when loss grows exponentionally. Default is 10.