Skip to content

types

cuthbert.enkf.types

Provides types for EnKF callback functions.

InitSample

Bases: Protocol

Protocol for sampling from the initial distribution.

__call__(key, model_inputs)

Sample from the initial distribution.

Parameters:

Name Type Description Default
key KeyArray

JAX PRNG key.

 required
model_inputs ArrayTreeLike

Model inputs.

 required

Returns:

Type Description
Array

Sample from the initial distribution, shape (x_dim,).
Source code in cuthbert/enkf/types.py 
12
13
14
15
16
17
18
19
20
21
22
def __call__(self, key: KeyArray, model_inputs: ArrayTreeLike) -> Array:
    """Sample from the initial distribution.

    Args:
        key: JAX PRNG key.
        model_inputs: Model inputs.

    Returns:
        Sample from the initial distribution, shape (x_dim,).
    """
    ...

GetEnKFDynamics

Bases: Protocol

Protocol for getting dynamics function that describes the general simulator p(x_{t+1} | x_t).

__call__(model_inputs)

Get dynamics function that describes the general simulator p(x_{t+1} | x_t) from model inputs.

Parameters:

Name Type Description Default
model_inputs ArrayTreeLike

Model inputs.

 required

Returns:

Type Description
DynamicsFn

Dynamics function that describes the general simulator (x_t, key) -> x_{t+1} ~ p(x_{t+1} | x_t).
Source code in cuthbert/enkf/types.py 
28
29
30
31
32
33
34
35
36
37
def __call__(self, model_inputs: ArrayTreeLike) -> DynamicsFn:
    """Get dynamics function that describes the general simulator p(x_{t+1} | x_t) from model inputs.

    Args:
        model_inputs: Model inputs.

    Returns:
        Dynamics function that describes the general simulator (x_t, key) -> x_{t+1} ~ p(x_{t+1} | x_t).
    """
    ...

GetEnKFObservations

Bases: Protocol

Protocol for getting observation function, observation noise Cholesky factor chol_R, and observation vector y for an EnKF. i.e., for state space x_t, y_t = h(x_t, model_inputs) + R_t, where R_t ~ N(0, chol_R), return (h, chol_R, y).

__call__(model_inputs)

Get observation function, observation noise Cholesky factor chol_R, and observation vector y from model inputs.

Parameters:

Name Type Description Default
model_inputs ArrayTreeLike

Model inputs.

 required

Returns:

Type Description
ObservationFn

Tuple with observation function, Cholesky factor of the observation noise covariance, and observation vector.
Array

observation noise covariance and y is the observation vector.
Source code in cuthbert/enkf/types.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
def __call__(
    self, model_inputs: ArrayTreeLike
) -> tuple[ObservationFn, Array, Array]:
    """Get observation function, observation noise Cholesky factor chol_R, and observation vector y from model inputs.

    Args:
        model_inputs: Model inputs.

    Returns:
        Tuple with observation function, Cholesky factor of the observation noise covariance, and observation vector.
        observation noise covariance and y is the observation vector.
    """
    ...