pygeoinf.data_assimilation package

Subpackages

Submodules

pygeoinf.data_assimilation.core module

core.py

A dimension-agnostic engine for numerical integration, statistical analysis, Bayesian inference, and generic visualisation.

class pygeoinf.data_assimilation.core.AssimilationEngine[source]

Bases: ABC

Abstract Base Class defining the interface for all assimilation methods (Grid, KF, EnKF). Enforces a consistent ‘setup -> observe -> run’ workflow.

add_observation(time: float, covariance: ndarray, value: ndarray | None = None, operator: ndarray | Callable | None = None)[source]

Registers an observation to be assimilated during the run.

abstract reanalyse_initial_condition(history: List[Dict[str, Any]]) → Any[source]

Performs Reanalysis (Smoothing) to estimate the state at t=0 given all observations up to t_final.

abstract run(initial_state: Any, t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the assimilation cycle. Must return a history list of dicts containing ‘time’, ‘forecast’, ‘analysis’.

class pygeoinf.data_assimilation.core.BayesianAssimilationProblem(eom_func: Callable[[float, ndarray, Any], ndarray], eom_args: Tuple = ())[source]

Bases: AssimilationEngine

Manages the definition and execution of a Grid-based Bayesian assimilation cycle. Inherits observation management from AssimilationEngine.

generate_synthetic_data(true_initial_condition: ndarray, dt_render: float = 0.05, seed: int | None = None) → Dict[str, Any][source]

Runs the physics from t=0, samples noisy observations at registered times, updates the internal Likelihood models with these values, and returns ground truth.

reanalyse_initial_condition(history: List[Dict[str, Any]]) → ProbabilityGrid[source]

Pulls the final posterior back to t=0 using inverse advection. Returns the smoothed ProbabilityGrid at t=0.

run(initial_state: ProbabilityGrid, t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the assimilation cycle starting from t=0.

Parameters:

  • initial_state – The ProbabilityGrid at t=0 (Prior).

  • t_final – Optional end time.

Returns:

A list of dictionaries containing {‘time’, ‘forecast’, ‘analysis’, ‘evidence’}.

class pygeoinf.data_assimilation.core.EnsembleKalmanFilter(eom_func: Callable[[float, ndarray, Any], ndarray], eom_args: Tuple = (), n_ensemble: int = 50)[source]

Bases: AssimilationEngine

Stochastic Ensemble Kalman Filter (EnKF). Suitable for non-linear dynamics. Uses the “Perturbed Observation” method.

System Model:

x_k = f(x_{k-1}) (Deterministic or Stochastic) y_k = H(x_k) + v_k

reanalyse_initial_condition(history: List[Dict[str, Any]]) → Tuple[ndarray, ndarray][source]

Reanalysis for deterministic physics. Takes the MEAN of the final analysis ensemble and integrates it BACKWARDS.

run(initial_state: Tuple[ndarray, ndarray] | ndarray, t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the EnKF.

Parameters:

initial_state – Either a tuple (mean, cov) to sample from, OR an explicit ensemble array (N_ens, N_dim).

class pygeoinf.data_assimilation.core.GaussianLikelihood(observation_value: ndarray | None, observation_covariance: ndarray, obs_operator_func: Callable[[ndarray], ndarray] | None = None)[source]

Bases: object

Handles Gaussian likelihoods for generic non-linear observation operators. Likelihood L(x) = P(y_obs | x) ~ N(y_obs; H(x), R)

evaluate(prob_grid: ProbabilityGrid) → ProbabilityGrid[source]

Evaluates the likelihood on the given ProbabilityGrid. Returns a new ProbabilityGrid instance containing the Likelihood values.

sample(true_state: ndarray) → ndarray[source]

Generates a noisy observation y given a true state x. y = H(x) + N(0, R)

Parameters:

true_state – The true state vector x.

Returns:

The noisy observation vector y.

class pygeoinf.data_assimilation.core.LinearGaussianLikelihood(observation_value: ndarray | None, observation_covariance: ndarray, observation_matrix: ndarray)[source]

Bases: GaussianLikelihood

Optimised subclass for Linear Observation Operators: y = Hx Avoids generic function calls in favour of matrix multiplication.

evaluate(prob_grid: ProbabilityGrid) → ProbabilityGrid[source]

Optimised evaluation for linear H.

class pygeoinf.data_assimilation.core.LinearKalmanFilter(transition_matrix_func: Callable[[float], ndarray])[source]

Bases: AssimilationEngine

Deterministic Linear Kalman Filter. Assumes perfect physics (Process Noise Q = 0).

System Model:

x_k = F_k * x_{k-1} (Deterministic) y_k = H * x_k + v_k, v_k ~ N(0, R)

Because the physics are deterministic, we can exactly reconstruct the smoothed initial condition by inverting the final state.

reanalyse_initial_condition(history: List[Dict[str, Any]]) → Tuple[ndarray, ndarray][source]

Back-propagates the final estimate to t=0 by inverting the physics. Exact for deterministic systems.

run(initial_state: Tuple[ndarray, ndarray], t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the Deterministic KF. :param initial_state: Tuple (mean_0, cov_0)

class pygeoinf.data_assimilation.core.ProbabilityGrid(axes: List[ndarray], values: ndarray)[source]

Bases: object

Encapsulates an N-dimensional probability density function discretised on a rectilinear grid. Handles normalisation, marginalisation, and Bayesian updates.

bayes_update(likelihood: ProbabilityGrid | ndarray | float) → Tuple[ProbabilityGrid, float][source]

Calculates Posterior = (Prior * Likelihood) / Evidence.

Parameters:

likelihood – Can be a ProbabilityGrid, a numpy array (same shape), or a scalar representing P(y|x).

Returns:

Normalised posterior density. evidence (float): The normalisation constant (integral of Prior * Likelihood).

Return type:

posterior (ProbabilityGrid)

classmethod from_bounds(bounds: List[Tuple[float, float]], resolution: int | List[int], pdf_func: Callable[[...], ndarray] | None = None) → ProbabilityGrid[source]

Factory: Creates a grid from bounds [(min, max), …] and resolution. Optionally evaluates a pdf_func on that grid immediately.

Parameters:

  • bounds – List of (min, max) tuples for each dimension.

  • resolution – Number of points per dimension (int or list of ints).

  • pdf_func – Optional callable f(x1, x2…) -> density to evaluate on initialisation.

Returns:

A new ProbabilityGrid instance.

marginalise(keep_indices: Tuple[int, ...]) → ProbabilityGrid[source]

Integrates out all axes NOT in keep_indices.

Parameters:

keep_indices – Tuple of dimension indices to retain (e.g., (0, 1)).

Returns:

A lower-dimensional ProbabilityGrid.

property mean: ndarray

Computes the expected value vector E[x] of the distribution.

normalise() → ProbabilityGrid[source]

Returns a NEW ProbabilityGrid that sums to 1.0. If mass is zero, returns the original grid to avoid division errors.

push_forward(eom_func: Callable, t_final: float, eom_args: Tuple = ()) → ProbabilityGrid[source]

Performs Liouville Advection (Method of Characteristics). Propagates the probability density forward in time by ‘t_final’.

Parameters:

  • eom_func – The differential equation f(t, y, *args).

  • t_final – The time duration to propagate forward.

  • eom_args – Physics arguments for the eom_func.

Returns:

A NEW ProbabilityGrid at t=t_final (normalised).

sample(n_samples: int = 1000) → ndarray[source]

Efficiently samples from the grid using PMF approximation + jitter.

Parameters:

n_samples – Number of samples to draw.

Returns:

Array of shape (n_samples, n_dim) containing the samples.

to_interpolator(fill_value: float = 0.0) → Callable[[...], ndarray][source]

Returns a callable function f(x1, x2, …) backed by this grid. Uses linear interpolation.

property total_mass: float

Computes the total integral (volume) of the grid using the trapezoidal rule.

pygeoinf.data_assimilation.core.display_animation_html(anim: Any) → HTML[source]

Standard helper to render animations in notebooks.

Parameters:

anim – The Matplotlib FuncAnimation object.

Returns:

IPython HTML object for display.

pygeoinf.data_assimilation.core.get_gaussian_pdf(mean: ndarray, cov: ndarray) → Callable[[...], ndarray][source]

Returns a callable PDF function for an N-dimensional multivariate Gaussian. The returned function is compatible with ProbabilityGrid factories.

Parameters:

  • mean – Mean vector of shape (N,).

  • cov – Covariance matrix of shape (N, N).

Returns:

A function pdf(*coordinates) -> density_array.

pygeoinf.data_assimilation.core.get_independent_gaussian_pdf(means: List[float] | ndarray, stds: List[float] | ndarray) → Callable[[...], ndarray][source]

Returns a callable PDF for N independent Gaussians (diagonal covariance). Faster and simpler than the full multivariate version.

Parameters:

  • means – List or array of means for each dimension.

  • stds – List or array of standard deviations for each dimension.

Returns:

A function pdf(*coordinates) -> density_array.

pygeoinf.data_assimilation.core.plot_1d_slice(prob_grid: ProbabilityGrid, dim: int = 0, ax: Axes | None = None, **kwargs) → Axes[source]

Marginalises down to a single dimension and plots a line graph (PDF).

Parameters:

  • prob_grid – core.ProbabilityGrid instance.

  • dim – Index of the dimension to keep.

  • ax – Matplotlib Axes object.

  • **kwargs – Passed to ax.plot (e.g., color, linewidth).

Returns:

The matplotlib Axis.

pygeoinf.data_assimilation.core.plot_ensemble_scatter(trajectories: ndarray, dim_indices: Tuple[int, int] = (0, 1), time_idx: int = -1, ax: Axes | None = None, **kwargs) → Axes[source]

Generic scatter plot of ensemble particles at a specific time snapshot.

Parameters:

  • trajectories – Array of shape (n_samples, n_dim, n_times).

  • dim_indices – Tuple of (x_dim, y_dim) indices to plot.

  • time_idx – Integer index of the time step to plot.

  • ax – Matplotlib Axes object.

  • **kwargs – Passed to ax.scatter (e.g., c, s, alpha, label).

Returns:

The matplotlib Axis.

pygeoinf.data_assimilation.core.plot_gaussian_ellipsoid(mean: ndarray, cov: ndarray, ax: Axes, n_std: float = 2.0, dims: Tuple[int, int] = (0, 1), **kwargs)[source]

Plots a 2D covariance ellipse representing the Gaussian distribution.

Parameters:

  • mean – Mean vector (N,).

  • cov – Covariance matrix (N, N).

  • ax – Matplotlib axes.

  • n_std – Number of standard deviations for the ellipse radius.

  • dims – Tuple of (x_dim, y_dim) indices.

  • **kwargs – Passed to ax.plot (color, linestyle, etc).

pygeoinf.data_assimilation.core.plot_grid_marginal(prob_grid: ProbabilityGrid, dims: Tuple[int, int] = (0, 1), ax: Axes | None = None, filled: bool = True, **kwargs) → Tuple[Axes, Any][source]

Marginalises an N-dimensional ProbabilityGrid down to 2 dimensions and plots the result as a contour.

Parameters:

  • prob_grid – core.ProbabilityGrid instance.

  • dims – Tuple of (x_dim_index, y_dim_index) to keep.

  • ax – Matplotlib Axes object. If None, creates a new figure.

  • filled – Boolean, whether to use contourf (filled) or contour (lines).

  • **kwargs – Passed directly to ax.contourf/ax.contour (e.g., levels, cmap).

Returns:

The matplotlib Axis. contour: The contour plot object.

Return type:

ax

pygeoinf.data_assimilation.core.plot_kf_step(forecast_stats: Tuple[ndarray, ndarray], analysis_stats: Tuple[ndarray, ndarray], dims: Tuple[int, int] = (0, 1), ax: Axes | None = None, title: str | None = None)[source]

Visualises a single Kalman Filter step (Forecast vs Analysis). Plots both Gaussian ellipsoids (2-sigma).

Parameters:

  • forecast_stats – (mean, cov) tuple for the prior.

  • analysis_stats – (mean, cov) tuple for the posterior.

  • dims – Dimensions to plot.

pygeoinf.data_assimilation.core.plot_tracker_1d(history: List[Dict[str, Any]], dim: int = 0, ground_truth: Dict[str, Any] | None = None, ax: Axes | None = None)[source]

Plots the time-evolution of a single state variable with uncertainty bounds.

Parameters:

  • history – The output list from kf.run().

  • dim – Index of the state dimension to plot.

  • ground_truth – Optional dict from generate_synthetic_data to compare against.

pygeoinf.data_assimilation.core.solve_ensemble(eom_func: Callable[[float, ndarray, Any], ndarray], initial_conditions: ndarray, t_points: ndarray, args: Tuple = (), **solver_kwargs) → ndarray[source]

Propagates an ensemble of particles forward in time.

Parameters:

  • eom_func – The Equation of Motion function.

  • initial_conditions – Array of shape (n_samples, n_dim).

  • t_points – Array of time points.

  • args – Physics arguments passed to eom_func.

  • **solver_kwargs – Extra kwargs for solve_trajectory (rtol, atol, etc).

Returns:

Trajectories array of shape (n_samples, n_dim, n_times).

pygeoinf.data_assimilation.core.solve_trajectory(eom_func: Callable[[float, ndarray, Any], ndarray], y0: ndarray, t_points: ndarray, args: Tuple = (), rtol: float = 1e-09, atol: float = 1e-12, method: str = 'RK45') → ndarray[source]

Integrates a single ODE trajectory over time.

Parameters:

  • eom_func – The Equation of Motion function f(t, y, *args) -> dy/dt.

  • y0 – Initial state vector of shape (n_dim,).

  • t_points – Array of time points to evaluate at.

  • args – Tuple of extra arguments to pass to eom_func.

  • rtol – Relative tolerance for solver.

  • atol – Absolute tolerance for solver.

  • method – Integration method (e.g., ‘RK45’, ‘DOP853’).

Returns:

Solution array of shape (n_dim, n_times).

pygeoinf.data_assimilation.core.wrap_angle(theta: float | ndarray) → float | ndarray[source]

Wraps an angle or array of angles to the interval [-pi, pi].

Parameters:

theta – Input angle(s) in radians.

Returns:

The wrapped angle(s) in [-pi, pi].

Module contents

QCES Data Assimilation Package.

A dimension-agnostic engine for numerical integration, statistical analysis, and Bayesian inference, with specific examples for dynamical systems.

class pygeoinf.data_assimilation.AssimilationEngine[source]

Bases: ABC

Abstract Base Class defining the interface for all assimilation methods (Grid, KF, EnKF). Enforces a consistent ‘setup -> observe -> run’ workflow.

add_observation(time: float, covariance: ndarray, value: ndarray | None = None, operator: ndarray | Callable | None = None)[source]

Registers an observation to be assimilated during the run.

abstract reanalyse_initial_condition(history: List[Dict[str, Any]]) → Any[source]

Performs Reanalysis (Smoothing) to estimate the state at t=0 given all observations up to t_final.

abstract run(initial_state: Any, t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the assimilation cycle. Must return a history list of dicts containing ‘time’, ‘forecast’, ‘analysis’.

class pygeoinf.data_assimilation.BayesianAssimilationProblem(eom_func: Callable[[float, ndarray, Any], ndarray], eom_args: Tuple = ())[source]

Bases: AssimilationEngine

Manages the definition and execution of a Grid-based Bayesian assimilation cycle. Inherits observation management from AssimilationEngine.

generate_synthetic_data(true_initial_condition: ndarray, dt_render: float = 0.05, seed: int | None = None) → Dict[str, Any][source]

Runs the physics from t=0, samples noisy observations at registered times, updates the internal Likelihood models with these values, and returns ground truth.

reanalyse_initial_condition(history: List[Dict[str, Any]]) → ProbabilityGrid[source]

Pulls the final posterior back to t=0 using inverse advection. Returns the smoothed ProbabilityGrid at t=0.

run(initial_state: ProbabilityGrid, t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the assimilation cycle starting from t=0.

Parameters:

  • initial_state – The ProbabilityGrid at t=0 (Prior).

  • t_final – Optional end time.

Returns:

A list of dictionaries containing {‘time’, ‘forecast’, ‘analysis’, ‘evidence’}.

class pygeoinf.data_assimilation.EnsembleKalmanFilter(eom_func: Callable[[float, ndarray, Any], ndarray], eom_args: Tuple = (), n_ensemble: int = 50)[source]

Bases: AssimilationEngine

Stochastic Ensemble Kalman Filter (EnKF). Suitable for non-linear dynamics. Uses the “Perturbed Observation” method.

System Model:

x_k = f(x_{k-1}) (Deterministic or Stochastic) y_k = H(x_k) + v_k

reanalyse_initial_condition(history: List[Dict[str, Any]]) → Tuple[ndarray, ndarray][source]

Reanalysis for deterministic physics. Takes the MEAN of the final analysis ensemble and integrates it BACKWARDS.

run(initial_state: Tuple[ndarray, ndarray] | ndarray, t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the EnKF.

Parameters:

initial_state – Either a tuple (mean, cov) to sample from, OR an explicit ensemble array (N_ens, N_dim).

class pygeoinf.data_assimilation.GaussianLikelihood(observation_value: ndarray | None, observation_covariance: ndarray, obs_operator_func: Callable[[ndarray], ndarray] | None = None)[source]

Bases: object

Handles Gaussian likelihoods for generic non-linear observation operators. Likelihood L(x) = P(y_obs | x) ~ N(y_obs; H(x), R)

evaluate(prob_grid: ProbabilityGrid) → ProbabilityGrid[source]

Evaluates the likelihood on the given ProbabilityGrid. Returns a new ProbabilityGrid instance containing the Likelihood values.

sample(true_state: ndarray) → ndarray[source]

Generates a noisy observation y given a true state x. y = H(x) + N(0, R)

Parameters:

true_state – The true state vector x.

Returns:

The noisy observation vector y.

class pygeoinf.data_assimilation.LinearGaussianLikelihood(observation_value: ndarray | None, observation_covariance: ndarray, observation_matrix: ndarray)[source]

Bases: GaussianLikelihood

Optimised subclass for Linear Observation Operators: y = Hx Avoids generic function calls in favour of matrix multiplication.

evaluate(prob_grid: ProbabilityGrid) → ProbabilityGrid[source]

Optimised evaluation for linear H.

class pygeoinf.data_assimilation.LinearKalmanFilter(transition_matrix_func: Callable[[float], ndarray])[source]

Bases: AssimilationEngine

Deterministic Linear Kalman Filter. Assumes perfect physics (Process Noise Q = 0).

System Model:

x_k = F_k * x_{k-1} (Deterministic) y_k = H * x_k + v_k, v_k ~ N(0, R)

Because the physics are deterministic, we can exactly reconstruct the smoothed initial condition by inverting the final state.

reanalyse_initial_condition(history: List[Dict[str, Any]]) → Tuple[ndarray, ndarray][source]

Back-propagates the final estimate to t=0 by inverting the physics. Exact for deterministic systems.

run(initial_state: Tuple[ndarray, ndarray], t_final: float | None = None) → List[Dict[str, Any]][source]

Executes the Deterministic KF. :param initial_state: Tuple (mean_0, cov_0)

class pygeoinf.data_assimilation.ProbabilityGrid(axes: List[ndarray], values: ndarray)[source]

Bases: object

Encapsulates an N-dimensional probability density function discretised on a rectilinear grid. Handles normalisation, marginalisation, and Bayesian updates.

bayes_update(likelihood: ProbabilityGrid | ndarray | float) → Tuple[ProbabilityGrid, float][source]

Calculates Posterior = (Prior * Likelihood) / Evidence.

Parameters:

likelihood – Can be a ProbabilityGrid, a numpy array (same shape), or a scalar representing P(y|x).

Returns:

Normalised posterior density. evidence (float): The normalisation constant (integral of Prior * Likelihood).

Return type:

posterior (ProbabilityGrid)

classmethod from_bounds(bounds: List[Tuple[float, float]], resolution: int | List[int], pdf_func: Callable[[...], ndarray] | None = None) → ProbabilityGrid[source]

Factory: Creates a grid from bounds [(min, max), …] and resolution. Optionally evaluates a pdf_func on that grid immediately.

Parameters:

  • bounds – List of (min, max) tuples for each dimension.

  • resolution – Number of points per dimension (int or list of ints).

  • pdf_func – Optional callable f(x1, x2…) -> density to evaluate on initialisation.

Returns:

A new ProbabilityGrid instance.

marginalise(keep_indices: Tuple[int, ...]) → ProbabilityGrid[source]

Integrates out all axes NOT in keep_indices.

Parameters:

keep_indices – Tuple of dimension indices to retain (e.g., (0, 1)).

Returns:

A lower-dimensional ProbabilityGrid.

property mean: ndarray

Computes the expected value vector E[x] of the distribution.

normalise() → ProbabilityGrid[source]

Returns a NEW ProbabilityGrid that sums to 1.0. If mass is zero, returns the original grid to avoid division errors.

push_forward(eom_func: Callable, t_final: float, eom_args: Tuple = ()) → ProbabilityGrid[source]

Performs Liouville Advection (Method of Characteristics). Propagates the probability density forward in time by ‘t_final’.

Parameters:

  • eom_func – The differential equation f(t, y, *args).

  • t_final – The time duration to propagate forward.

  • eom_args – Physics arguments for the eom_func.

Returns:

A NEW ProbabilityGrid at t=t_final (normalised).

sample(n_samples: int = 1000) → ndarray[source]

Efficiently samples from the grid using PMF approximation + jitter.

Parameters:

n_samples – Number of samples to draw.

Returns:

Array of shape (n_samples, n_dim) containing the samples.

to_interpolator(fill_value: float = 0.0) → Callable[[...], ndarray][source]

Returns a callable function f(x1, x2, …) backed by this grid. Uses linear interpolation.

property total_mass: float

Computes the total integral (volume) of the grid using the trapezoidal rule.

pygeoinf.data_assimilation.display_animation_html(anim: Any) → HTML[source]

Standard helper to render animations in notebooks.

Parameters:

anim – The Matplotlib FuncAnimation object.

Returns:

IPython HTML object for display.

pygeoinf.data_assimilation.get_gaussian_pdf(mean: ndarray, cov: ndarray) → Callable[[...], ndarray][source]

Returns a callable PDF function for an N-dimensional multivariate Gaussian. The returned function is compatible with ProbabilityGrid factories.

Parameters:

  • mean – Mean vector of shape (N,).

  • cov – Covariance matrix of shape (N, N).

Returns:

A function pdf(*coordinates) -> density_array.

pygeoinf.data_assimilation.get_independent_gaussian_pdf(means: List[float] | ndarray, stds: List[float] | ndarray) → Callable[[...], ndarray][source]

Returns a callable PDF for N independent Gaussians (diagonal covariance). Faster and simpler than the full multivariate version.

Parameters:

  • means – List or array of means for each dimension.

  • stds – List or array of standard deviations for each dimension.

Returns:

A function pdf(*coordinates) -> density_array.

pygeoinf.data_assimilation.plot_1d_slice(prob_grid: ProbabilityGrid, dim: int = 0, ax: Axes | None = None, **kwargs) → Axes[source]

Marginalises down to a single dimension and plots a line graph (PDF).

Parameters:

  • prob_grid – core.ProbabilityGrid instance.

  • dim – Index of the dimension to keep.

  • ax – Matplotlib Axes object.

  • **kwargs – Passed to ax.plot (e.g., color, linewidth).

Returns:

The matplotlib Axis.

pygeoinf.data_assimilation.plot_ensemble_scatter(trajectories: ndarray, dim_indices: Tuple[int, int] = (0, 1), time_idx: int = -1, ax: Axes | None = None, **kwargs) → Axes[source]

Generic scatter plot of ensemble particles at a specific time snapshot.

Parameters:

  • trajectories – Array of shape (n_samples, n_dim, n_times).

  • dim_indices – Tuple of (x_dim, y_dim) indices to plot.

  • time_idx – Integer index of the time step to plot.

  • ax – Matplotlib Axes object.

  • **kwargs – Passed to ax.scatter (e.g., c, s, alpha, label).

Returns:

The matplotlib Axis.

pygeoinf.data_assimilation.plot_gaussian_ellipsoid(mean: ndarray, cov: ndarray, ax: Axes, n_std: float = 2.0, dims: Tuple[int, int] = (0, 1), **kwargs)[source]

Plots a 2D covariance ellipse representing the Gaussian distribution.

Parameters:

  • mean – Mean vector (N,).

  • cov – Covariance matrix (N, N).

  • ax – Matplotlib axes.

  • n_std – Number of standard deviations for the ellipse radius.

  • dims – Tuple of (x_dim, y_dim) indices.

  • **kwargs – Passed to ax.plot (color, linestyle, etc).

pygeoinf.data_assimilation.plot_grid_marginal(prob_grid: ProbabilityGrid, dims: Tuple[int, int] = (0, 1), ax: Axes | None = None, filled: bool = True, **kwargs) → Tuple[Axes, Any][source]

Marginalises an N-dimensional ProbabilityGrid down to 2 dimensions and plots the result as a contour.

Parameters:

  • prob_grid – core.ProbabilityGrid instance.

  • dims – Tuple of (x_dim_index, y_dim_index) to keep.

  • ax – Matplotlib Axes object. If None, creates a new figure.

  • filled – Boolean, whether to use contourf (filled) or contour (lines).

  • **kwargs – Passed directly to ax.contourf/ax.contour (e.g., levels, cmap).

Returns:

The matplotlib Axis. contour: The contour plot object.

Return type:

ax

pygeoinf.data_assimilation.plot_kf_step(forecast_stats: Tuple[ndarray, ndarray], analysis_stats: Tuple[ndarray, ndarray], dims: Tuple[int, int] = (0, 1), ax: Axes | None = None, title: str | None = None)[source]

Visualises a single Kalman Filter step (Forecast vs Analysis). Plots both Gaussian ellipsoids (2-sigma).

Parameters:

  • forecast_stats – (mean, cov) tuple for the prior.

  • analysis_stats – (mean, cov) tuple for the posterior.

  • dims – Dimensions to plot.

pygeoinf.data_assimilation.plot_tracker_1d(history: List[Dict[str, Any]], dim: int = 0, ground_truth: Dict[str, Any] | None = None, ax: Axes | None = None)[source]

Plots the time-evolution of a single state variable with uncertainty bounds.

Parameters:

  • history – The output list from kf.run().

  • dim – Index of the state dimension to plot.

  • ground_truth – Optional dict from generate_synthetic_data to compare against.

pygeoinf.data_assimilation.solve_ensemble(eom_func: Callable[[float, ndarray, Any], ndarray], initial_conditions: ndarray, t_points: ndarray, args: Tuple = (), **solver_kwargs) → ndarray[source]

Propagates an ensemble of particles forward in time.

Parameters:

  • eom_func – The Equation of Motion function.

  • initial_conditions – Array of shape (n_samples, n_dim).

  • t_points – Array of time points.

  • args – Physics arguments passed to eom_func.

  • **solver_kwargs – Extra kwargs for solve_trajectory (rtol, atol, etc).

Returns:

Trajectories array of shape (n_samples, n_dim, n_times).

pygeoinf.data_assimilation.solve_trajectory(eom_func: Callable[[float, ndarray, Any], ndarray], y0: ndarray, t_points: ndarray, args: Tuple = (), rtol: float = 1e-09, atol: float = 1e-12, method: str = 'RK45') → ndarray[source]

Integrates a single ODE trajectory over time.

Parameters:

  • eom_func – The Equation of Motion function f(t, y, *args) -> dy/dt.

  • y0 – Initial state vector of shape (n_dim,).

  • t_points – Array of time points to evaluate at.

  • args – Tuple of extra arguments to pass to eom_func.

  • rtol – Relative tolerance for solver.

  • atol – Absolute tolerance for solver.

  • method – Integration method (e.g., ‘RK45’, ‘DOP853’).

Returns:

Solution array of shape (n_dim, n_times).

pygeoinf.data_assimilation.wrap_angle(theta: float | ndarray) → float | ndarray[source]

Wraps an angle or array of angles to the interval [-pi, pi].

Parameters:

theta – Input angle(s) in radians.

Returns:

The wrapped angle(s) in [-pi, pi].