Skip to content

types

cuthbert.ensemble_kalman.types

Provides types for EnKF callback functions.

InitSample

Bases: Protocol

Protocol for sampling from the initial distribution.

__call__(key)

Sample from the initial distribution.

Parameters:

Name Type Description Default
key KeyArray

JAX PRNG key.

 required

Returns:

Type Description
Array

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

    Args:
        key: JAX PRNG key.

    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/ensemble_kalman/types.py 
30
31
32
33
34
35
36
37
38
39
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 generalised 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 generalised 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, generalised Cholesky factor of the observation noise covariance, and observation vector.
Array

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

    Args:
        model_inputs: Model inputs.

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

ModifyCrossCovariance

Bases: Protocol

Protocol for modifying an EnKF cross-covariance.

__call__(cross_covariance, model_inputs)

Modify an empirical state-observation cross-covariance.

Parameters:

Name Type Description Default
cross_covariance Array

Empirical state-observation cross-covariance.

 required
model_inputs ArrayTreeLike

Model inputs.

 required

Returns:

Type Description
Array

Modified state-observation cross-covariance.
Source code in cuthbert/ensemble_kalman/types.py 
63
64
65
66
67
68
69
70
71
72
73
def __call__(self, cross_covariance: Array, model_inputs: ArrayTreeLike) -> Array:
    """Modify an empirical state-observation cross-covariance.

    Args:
        cross_covariance: Empirical state-observation cross-covariance.
        model_inputs: Model inputs.

    Returns:
        Modified state-observation cross-covariance.
    """
    ...

ConstructCholInnovationCovariance

Bases: Protocol

Protocol for constructing an EnKF innovation covariance factor.

__call__(normalized_observation_deviations, chol_observation_covariance, model_inputs)

Construct a generalized Cholesky factor of an innovation covariance.

Parameters:

Name Type Description Default
normalized_observation_deviations Array

Observation deviations transposed and divided by the square root of one less than the ensemble size, shape (y_dim, n_particles).

 required
chol_observation_covariance Array

Cholesky factor of the observation noise covariance, shape (y_dim, y_dim).

 required
model_inputs ArrayTreeLike

Model inputs.

 required

Returns:

Type Description
Array

Generalized Cholesky factor of the complete innovation covariance,
Array

shape (y_dim, y_dim).
Source code in cuthbert/ensemble_kalman/types.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __call__(
    self,
    normalized_observation_deviations: Array,
    chol_observation_covariance: Array,
    model_inputs: ArrayTreeLike,
) -> Array:
    """Construct a generalized Cholesky factor of an innovation covariance.

    Args:
        normalized_observation_deviations: Observation deviations transposed
            and divided by the square root of one less than the ensemble size,
            shape (y_dim, n_particles).
        chol_observation_covariance: Cholesky factor of the observation noise
            covariance, shape (y_dim, y_dim).
        model_inputs: Model inputs.

    Returns:
        Generalized Cholesky factor of the complete innovation covariance,
        shape (y_dim, y_dim).
    """
    ...