kkf 0.3.0__py3-none-any.whl

KKF/__init__.py

@@ -0,0 +1,48 @@
+"""Kernel Koopman Kalman Filter (KKF) - A library for nonlinear state estimation.
+
+This package provides a complete implementation of the Koopman-Kalman Filter,
+combining Koopman operator theory with Kalman filtering for accurate state
+estimation in nonlinear dynamical systems.
+
+Main Components
+---------------
+- systems : Dynamical system definitions
+- koopman : Koopman operator approximation via kEDMD
+- covariances : Covariance matrix computations
+- solution : Filter solution data structure
+- filter : Main filtering algorithm
+
+Quick Start
+-----------
+>>> from KKF import DynamicalSystem, KoopmanOperator, apply_koopman_kalman_filter
+>>> # or
+>>> from KKF.filter import apply_koopman_kalman_filter
+"""
+
+# Core classes and utilities
+from .covariances import (
+    compute_dynamics_covariance,
+    compute_initial_covariance,
+    compute_observation_covariance,
+)
+
+# Main filtering function
+from .filter import apply_koopman_kalman_filter
+from .koopman import KoopmanOperator
+from .solution import KoopmanKalmanFilterSolution
+from .systems import DynamicalSystem, create_additive_system
+
+__all__ = [
+    # Classes
+    "DynamicalSystem",
+    "KoopmanOperator",
+    "KoopmanKalmanFilterSolution",
+    # Functions
+    "create_additive_system",
+    "apply_koopman_kalman_filter",
+    "compute_initial_covariance",
+    "compute_dynamics_covariance",
+    "compute_observation_covariance",
+]
+
+__version__ = "0.3.0"

KKF/covariances.py

@@ -0,0 +1,159 @@
+from typing import Any, Callable, Union
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.stats import rv_continuous
+
+
+def compute_initial_covariance(
+    x: NDArray[np.float64],
+    n_features: int,
+    initial_distribution: rv_continuous,
+    koopman_operator: Any,
+    n_samples: int,
+) -> NDArray[np.float64]:
+    """
+    Compute the covariance matrix for the initial distribution in feature space.
+
+    This function samples from an initial distribution and computes the covariance
+    matrix of the transformed samples using a Koopman operator's feature map.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Current state vector or reference point.
+    n_features : int
+        Number of features in the transformed space.
+    initial_distribution : scipy.stats.rv_continuous
+        Initial probability distribution to sample from.
+    koopman_operator : object
+        Object containing the feature map phi method for state transformation.
+        Must have a method phi(x) that maps states to feature space.
+    n_samples : int
+        Number of samples to use for covariance estimation.
+
+    Returns
+    -------
+    np.ndarray
+        Covariance matrix of size (n_features, n_features) in the transformed space.
+
+    Notes
+    -----
+    The function performs the following steps:
+    1. Samples from the initial distribution
+    2. Applies the Koopman operator's feature map to each sample
+    3. Computes the covariance matrix of the transformed samples
+    """
+    # Sample from initial distribution
+    samples = initial_distribution.rvs(size=n_samples).reshape((n_samples, len(x)))
+
+    # Initialize array for transformed samples
+    transformed_samples = np.zeros((n_samples, n_features))
+
+    # Transform each sample using the feature map
+    for i in range(n_samples):
+        transformed_samples[i, :] = koopman_operator.phi(samples[i, :])
+
+    # Compute and return covariance matrix
+    return np.cov(transformed_samples, rowvar=False)
+
+
+def compute_dynamics_covariance(
+    x: NDArray[np.float64], n_features: int, dynamics: Any, koopman_operator: Any, n_samples: int
+) -> NDArray[np.float64]:
+    """
+    Compute the covariance matrix for the system dynamics in feature space.
+
+    This function samples from the dynamics distribution, applies the system
+    dynamics, and computes the covariance matrix of the transformed results.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Current state vector or reference point.
+    n_features : int
+        Number of features in the transformed space.
+    dynamics : object
+        Dynamical system object containing:
+        - dist_dyn: scipy.stats.rv_continuous distribution for dynamics noise
+        - dynamics(x, w): method implementing the system dynamics
+    koopman_operator : object
+        Object containing the feature map phi method for state transformation.
+        Must have a method phi(x) that maps states to feature space.
+    n_samples : int
+        Number of samples to use for covariance estimation.
+
+    Returns
+    -------
+    np.ndarray
+        Covariance matrix of size (n_features, n_features) in the transformed space.
+
+    Notes
+    -----
+    The function performs the following steps:
+    1. Samples from the dynamics noise distribution
+    2. Applies the system dynamics to each sample
+    3. Transforms the results using the Koopman operator's feature map
+    4. Computes the covariance matrix of the transformed samples
+    """
+    # Sample from dynamics distribution
+    noise_samples = dynamics.dist_dyn.rvs(size=n_samples).reshape((n_samples, len(x)))
+
+    # Initialize array for transformed samples
+    transformed_samples = np.zeros((n_samples, n_features))
+
+    # Apply dynamics and transform each sample
+    for i in range(n_samples):
+        state_evolution = dynamics.dynamics(x, noise_samples[i, :])
+        transformed_samples[i, :] = koopman_operator.phi(state_evolution)
+
+    # Compute and return covariance matrix
+    return np.cov(transformed_samples, rowvar=False)
+
+
+def compute_observation_covariance(
+    x: NDArray[np.float64], n_outputs: int, dynamics: Any, n_samples: int
+) -> NDArray[np.float64]:
+    """
+    Compute the covariance matrix for the observation/measurement process.
+
+    This function samples from the measurement noise distribution and computes
+    the covariance matrix of the measurement process.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Current state vector or reference point.
+    n_outputs : int
+        Number of measurement outputs.
+    dynamics : object
+        Dynamical system object containing:
+        - dist_obs: scipy.stats.rv_continuous distribution for measurement noise
+        - measurements(x, w): method implementing the measurement process
+    n_samples : int
+        Number of samples to use for covariance estimation.
+
+    Returns
+    -------
+    np.ndarray
+        Covariance matrix of size (n_outputs, n_outputs) for the measurement process.
+
+    Notes
+    -----
+    The function performs the following steps:
+    1. Samples from the measurement noise distribution
+    2. Applies the measurement function to each sample
+    3. Computes the covariance matrix of the measurements
+    """
+    # Sample from measurement distribution
+    noise_samples = dynamics.dist_obs.rvs(size=n_samples).reshape((n_samples, n_outputs))
+
+    # Initialize array for measurements
+    measurement_samples = np.zeros((n_samples, n_outputs))
+
+    # Apply measurement function to each sample
+    for i in range(n_samples):
+        measurement_samples[i, :] = dynamics.measurements(x, noise_samples[i, :])
+
+    # Compute and return covariance matrix (ensure 2D even for single output)
+    return np.atleast_2d(np.cov(measurement_samples, rowvar=False))

KKF/filter.py

@@ -0,0 +1,250 @@
+"""Koopman Kalman Filter algorithm implementation.
+
+Combines Koopman operator theory with Kalman filtering for nonlinear
+state estimation.
+"""
+
+from typing import Any, Callable, Tuple
+
+import numpy as np
+from numpy.linalg import pinv
+from numpy.typing import NDArray
+
+from .covariances import (
+    compute_dynamics_covariance,
+    compute_initial_covariance,
+    compute_observation_covariance,
+)
+from .solution import KoopmanKalmanFilterSolution
+
+
+def apply_koopman_kalman_filter(
+    koopman_operator: Any,
+    observations: NDArray[np.float64],
+    initial_distribution: Any,
+    n_features: int,
+    optimize: bool = True,
+    n_restarts_optimizer: int = 10,
+    noise_samples: int = 100,
+) -> KoopmanKalmanFilterSolution:
+    """
+    Implementation of the Koopman-Kalman Filter algorithm.
+
+    This function combines Koopman operator theory with Kalman filtering to perform
+    state estimation for nonlinear dynamical systems. It operates by lifting the
+    state space to a higher-dimensional feature space where the dynamics are
+    approximately linear.
+
+    Parameters
+    ----------
+    koopman_operator : object
+        Object containing Koopman operator methods and attributes:
+        - compute_edmd(n_features): Extended Dynamic Mode Decomposition method
+        - dynamical_system: Dynamical system object
+        - U, B, C: Matrices for Koopman approximation
+        - phi: Feature map function
+    observations : np.ndarray
+        Array of shape (n_timesteps, n_outputs) containing system measurements.
+    initial_distribution : object
+        Distribution object for initial state with methods:
+        - mean: Returns mean of initial state
+        - rvs: Random sampling method
+    n_features : int
+        Number of features in the lifted space.
+    optimize : bool, optional
+        Whether to optimize the kernel parameters. Default is True.
+    n_restarts_optimizer : int, optional
+        Number of restarts for kernel optimization. Default is 10.
+    noise_samples : int, optional
+        Number of samples for noise covariance estimation. Default is 100.
+
+    Returns
+    -------
+    KoopmanKalmanFilterSolution
+        Object containing all filter estimates and covariances.
+
+    Notes
+    -----
+    The algorithm consists of three main phases:
+    1. Initialization:
+       - Sets up initial states and covariances
+       - Computes Koopman approximation
+
+    2. Prediction Step:
+       - Projects state forward using system dynamics
+       - Updates covariances using Koopman operator
+
+    3. Update Step:
+       - Incorporates new measurements
+       - Updates state and covariance estimates
+
+    The filter operates in both the original state space (x) and the lifted
+    feature space (z), maintaining estimates and covariances in both spaces.
+    """
+    # Compute Koopman approximation
+    koopman_operator.compute_edmd(n_features, optimize, n_restarts_optimizer)
+    dynamical_system = koopman_operator.dynamical_system
+
+    # Extract system and Koopman components
+    state_dynamics = dynamical_system.f
+    measurement_model = dynamical_system.g
+    U, B, C = koopman_operator.U, koopman_operator.B, koopman_operator.C
+    phi = koopman_operator.phi
+
+    ### Initialization ###
+    mean_val = initial_distribution.mean  # Initial state estimate
+    x0 = np.atleast_1d(np.asarray(mean_val() if callable(mean_val) else mean_val))
+    z0 = phi(x0)  # Initial state in feature space
+
+    # Get problem dimensions
+    n_timesteps, n_states = observations.shape[0], len(x0)
+    n_outputs = observations.shape[1]
+
+    # Initialize state arrays
+    x_minus, x_plus = _initialize_state_arrays(n_timesteps, n_states)
+    z_minus, z_plus = _initialize_state_arrays(n_timesteps, n_features)
+
+    # Initialize covariance arrays
+    Px_minus, Px_plus = _initialize_covariance_arrays(n_timesteps, n_states)
+    Pz_minus, Pz_plus = _initialize_covariance_arrays(n_timesteps, n_features)
+
+    # Initialize filter matrices
+    S = np.zeros((n_timesteps, n_outputs, n_outputs))  # Innovation covariance
+    K = np.zeros((n_timesteps, n_features, n_outputs))  # Kalman gain
+
+    # Set initial conditions
+    x_minus[0, :], x_plus[0, :] = x0, x0
+    z_minus[0, :], z_plus[0, :] = z0, z0
+
+    # Initialize covariances
+    initial_covariance = compute_initial_covariance(
+        x_minus[0, :], n_features, initial_distribution, koopman_operator, noise_samples
+    )
+    Pz_minus[0, :, :], Pz_plus[0, :, :] = initial_covariance, initial_covariance
+
+    ### Main Filter Loop ###
+    for t in range(1, n_timesteps):
+        # Prediction Step
+        x_minus[t], z_minus[t], Pz_minus[t] = _prediction_step(
+            x_plus[t - 1],
+            z_plus[t - 1],
+            Pz_plus[t - 1],
+            state_dynamics,
+            phi,
+            U,
+            dynamical_system,
+            koopman_operator,
+            n_features,
+            noise_samples,
+        )
+
+        # Update Step
+        x_plus[t], z_plus[t], Pz_plus[t], S[t], K[t] = _update_step(
+            x_minus[t],
+            z_minus[t],
+            Pz_minus[t],
+            observations[t],
+            measurement_model,
+            C,
+            B,
+            dynamical_system,
+            noise_samples,
+        )
+
+        # Update state-space covariances
+        Px_minus[t], Px_plus[t] = _update_state_covariances(Pz_minus[t], Pz_plus[t], B)
+
+    return KoopmanKalmanFilterSolution(x_plus, x_minus, Pz_plus, Pz_minus, Px_plus, Px_minus, S, K)
+
+
+def _initialize_state_arrays(
+    n_timesteps: int, n_dim: int
+) -> Tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Initialize arrays for state estimates."""
+    return (np.zeros((n_timesteps, n_dim)), np.zeros((n_timesteps, n_dim)))
+
+
+def _initialize_covariance_arrays(
+    n_timesteps: int, n_dim: int
+) -> Tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Initialize arrays for covariance matrices."""
+    return (np.zeros((n_timesteps, n_dim, n_dim)), np.zeros((n_timesteps, n_dim, n_dim)))
+
+
+def _prediction_step(
+    x_prev: NDArray[np.float64],
+    z_prev: NDArray[np.float64],
+    Pz_prev: NDArray[np.float64],
+    state_dynamics: Callable,
+    phi: Callable,
+    U: NDArray[np.float64],
+    dynamical_system: Any,
+    koopman_operator: Any,
+    n_features: int,
+    noise_samples: int,
+) -> Tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
+    """Perform the prediction step of the filter."""
+    # Compute dynamics covariance
+    Qz = compute_dynamics_covariance(
+        x_prev, n_features, dynamical_system, koopman_operator, noise_samples
+    )
+
+    # Predict states
+    x_pred = state_dynamics(
+        x_prev,
+    )
+    z_pred = phi(x_pred)
+
+    # Predict covariance
+    Pz_pred = U @ Pz_prev @ U.T + Qz
+
+    return x_pred, z_pred, Pz_pred
+
+
+def _update_step(
+    x_minus: NDArray[np.float64],
+    z_minus: NDArray[np.float64],
+    Pz_minus: NDArray[np.float64],
+    observation: NDArray[np.float64],
+    measurement_model: Callable,
+    C: NDArray[np.float64],
+    B: NDArray[np.float64],
+    dynamical_system: Any,
+    noise_samples: int,
+) -> Tuple[
+    NDArray[np.float64],
+    NDArray[np.float64],
+    NDArray[np.float64],
+    NDArray[np.float64],
+    NDArray[np.float64],
+]:
+    """Perform the update step of the filter."""
+    # Compute measurement covariance
+    Rz = compute_observation_covariance(x_minus, len(observation), dynamical_system, noise_samples)
+
+    # Compute innovation
+    innovation = observation - measurement_model(x_minus)
+
+    # Compute innovation covariance
+    S = C @ Pz_minus @ C.T + Rz
+
+    # Compute Kalman gain
+    K = Pz_minus @ C.T @ pinv(S)
+
+    # Update states
+    z_plus = z_minus + K @ innovation
+    x_plus = B @ z_plus
+
+    # Update covariance
+    Pz_plus = (np.eye(len(z_minus)) - K @ C) @ Pz_minus
+
+    return x_plus, z_plus, Pz_plus, S, K
+
+
+def _update_state_covariances(
+    Pz_minus: NDArray[np.float64], Pz_plus: NDArray[np.float64], B: NDArray[np.float64]
+) -> Tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Update covariances in the state space."""
+    Px_minus = B @ Pz_minus @ B.T
+    Px_plus = B @ Pz_plus @ B.T
+    return Px_minus, Px_plus

KKF/koopman.py

@@ -0,0 +1,155 @@
+"""Koopman operator approximation module.
+
+Implements kernel-based Extended Dynamic Mode Decomposition (kEDMD)
+for computing Koopman operator approximations.
+"""
+
+from typing import Callable, Optional
+
+import numpy as np
+from numpy.typing import NDArray
+from sklearn.gaussian_process import GaussianProcessRegressor
+from sklearn.gaussian_process.kernels import Kernel
+
+from .systems import DynamicalSystem, create_additive_system
+
+
+class KoopmanOperator:
+    """
+    Implementation of Koopman operator approximation using kernel-based Extended
+    Dynamic Mode Decomposition (kEDMD).
+
+    This class provides methods to compute finite-dimensional approximations of the
+    Koopman operator for nonlinear dynamical systems.
+
+    Attributes
+    ----------
+    kernel_function : Callable
+        Kernel function for computing feature space mappings.
+    dynamical_system : DynamicalSystem
+        The underlying dynamical system.
+    X : Optional[np.ndarray]
+        Dictionary of states used for kernel computations.
+    phi : Optional[Callable]
+        Feature map function.
+    U : Optional[np.ndarray]
+        Koopman operator matrix.
+    G : Optional[np.ndarray]
+        Gram matrix.
+    C : Optional[np.ndarray]
+        Output matrix.
+    B : Optional[np.ndarray]
+        State-to-feature space transformation matrix.
+
+    Notes
+    -----
+    The Koopman operator framework lifts nonlinear dynamics to a linear setting
+    in a higher-dimensional feature space. This implementation uses kernel methods
+    to compute the necessary feature spaces and operators.
+    """
+
+    def __init__(
+        self,
+        kernel_function: Callable[[NDArray[np.float64], NDArray[np.float64]], NDArray[np.float64]],
+        dynamical_system: DynamicalSystem,
+    ) -> None:
+        self.kernel_function = kernel_function
+        self.dynamical_system = dynamical_system
+        self.X: Optional[NDArray[np.float64]] = None
+        self.phi: Optional[Callable] = None
+        self.U: Optional[NDArray[np.float64]] = None
+        self.G: Optional[NDArray[np.float64]] = None
+        self.C: Optional[NDArray[np.float64]] = None
+        self.B: Optional[NDArray[np.float64]] = None
+
+    def compute_edmd(
+        self, n_features: int, optimize: bool = True, n_restarts_optimizer: int = 10
+    ) -> None:
+        """
+        Compute the kernel-based Extended Dynamic Mode Decomposition (kEDMD).
+
+        This method constructs finite-dimensional approximations of the Koopman
+        operator and associated matrices using kernel methods.
+
+        Parameters
+        ----------
+        n_features : int
+            Number of features to use in the approximation.
+        optimize : bool
+            Whether to optimize the kernel function. If True, the method will
+            optimize the kernel function using Gaussian Process Regression. If False,
+            the provided kernel function will be used without optimization. Default is True.
+        n_restarts_optimizer : int
+            Number of restarts for the optimizer. If optimize is False, will be ignored. Default is 10.
+
+        Notes
+        -----
+        The method performs the following steps:
+        1. Generates dictionary points using the state distribution
+        2. Constructs the feature map using the kernel function
+        3. Computes the Gram matrix and its inverse
+        4. Constructs the Koopman operator approximation
+        5. Computes output and state transformation matrices
+        """
+        # Extract system components
+        f, g = self.dynamical_system.f, self.dynamical_system.g
+
+        # Generate dictionary points
+        self.X = self.dynamical_system.sample_state(n_features)
+
+        # Optimize kernel function
+        if optimize:
+            self.opt_kernel(X=self.X, n_restarts_optimizer=n_restarts_optimizer)
+
+        # Define feature map
+        self.phi = lambda x: self.kernel_function(x, self.X)[0]
+
+        # Compute Gram matrix
+        self.G = self.kernel_function(self.X, self.X)
+        G_inv = np.linalg.inv(self.G)
+
+        # Compute Koopman operator approximation
+        next_states = f(self.X.T).T
+        self.U = self.kernel_function(self.X, next_states) @ G_inv
+
+        # Compute output and state transformation matrices
+        self.C = g(self.X.T) @ G_inv
+        self.B = self.X.T @ G_inv
+
+    def get_feature_dimension(self) -> Optional[int]:
+        """
+        Get the dimension of the feature space.
+
+        Returns
+        -------
+        Optional[int]
+            Dimension of the feature space, or None if EDMD hasn't been computed.
+        """
+        return self.X.shape[0] if self.X is not None else None
+
+    def opt_kernel(self, X: NDArray[np.float64], n_restarts_optimizer: int = 10) -> None:
+        """
+        Optimize the kernel function for the Koopman operator.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            Set of points.
+        n_restarts_optimizer : int
+            Number of restarts for the optimizer. Default is 10.
+
+        Notes
+        -----
+        This method uses Gaussian Process Regression to optimize the kernel function.
+        """
+        # Compute the output of the dynamical system
+        y = self.dynamical_system.f(X.T).T
+
+        # Fit the Gaussian process
+        gp = GaussianProcessRegressor(
+            kernel=self.kernel_function, n_restarts_optimizer=n_restarts_optimizer
+        )
+        gp.fit(X.T, y.T)
+
+        # Update with optimal kernel
+        self.kernel_function = gp.kernel_

KKF/solution.py

@@ -0,0 +1,145 @@
+"""Solution data structure for Koopman Kalman Filter results.
+
+Provides the KoopmanKalmanFilterSolution dataclass for storing and accessing
+filter outputs.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+@dataclass
+class KoopmanKalmanFilterSolution:
+    """
+    A class to store the solution of a Koopman-Kalman Filter iteration.
+
+    This class maintains the state estimates, covariance matrices, and filter gains
+    for both the prior (minus) and posterior (plus) estimates in both state and
+    feature spaces.
+
+    Attributes
+    ----------
+    x_plus : np.ndarray
+        Posterior state estimate after measurement update.
+    x_minus : np.ndarray
+        Prior state estimate from prediction step.
+    Pz_plus : np.ndarray
+        Posterior covariance matrix in feature space after measurement update.
+    Pz_minus : np.ndarray
+        Prior covariance matrix in feature space from prediction step.
+    Px_plus : np.ndarray
+        Posterior covariance matrix in state space after measurement update.
+    Px_minus : np.ndarray
+        Prior covariance matrix in state space from prediction step.
+    S : np.ndarray
+        Innovation (residual) covariance matrix.
+    K : np.ndarray
+        Kalman gain matrix.
+
+    Notes
+    -----
+    The class uses the common Kalman filter notation where:
+    - (-) denotes prior estimates before measurement update
+    - (+) denotes posterior estimates after measurement update
+    - Pz refers to covariances in the feature/transformed space
+    - Px refers to covariances in the original state space
+
+    Examples
+    --------
+    >>> solution = KoopmanKalmanFilterSolution(
+    ...     x_plus=np.array([1.0, 2.0]),
+    ...     x_minus=np.array([0.9, 1.9]),
+    ...     Pz_plus=np.eye(2),
+    ...     Pz_minus=np.eye(2) * 1.1,
+    ...     Px_plus=np.eye(2) * 0.9,
+    ...     Px_minus=np.eye(2),
+    ...     S=np.eye(2) * 0.5,
+    ...     K=np.array([[0.1, 0], [0, 0.1]])
+    ... )
+    """
+
+    x_plus: NDArray[np.float64]
+    x_minus: NDArray[np.float64]
+    Pz_plus: NDArray[np.float64]
+    Pz_minus: NDArray[np.float64]
+    Px_plus: NDArray[np.float64]
+    Px_minus: NDArray[np.float64]
+    S: NDArray[np.float64]
+    K: NDArray[np.float64]
+
+    def __post_init__(self) -> None:
+        """Validate the dimensions of the input arrays."""
+        # Ensure all inputs are numpy arrays
+        for attr in ["x_plus", "x_minus", "Pz_plus", "Pz_minus", "Px_plus", "Px_minus", "S", "K"]:
+            value = getattr(self, attr)
+            if not isinstance(value, np.ndarray):
+                setattr(self, attr, np.array(value))
+
+    def get_state_dimension(self) -> int:
+        """
+        Get the dimension of the state vector.
+
+        Returns
+        -------
+        int
+            The dimension of the state vector.
+        """
+        return self.x_plus.shape[-1]
+
+    def get_feature_dimension(self) -> int:
+        """
+        Get the dimension of the feature space.
+
+        Returns
+        -------
+        int
+            The dimension of the feature space.
+        """
+        return self.Pz_plus.shape[-1]
+
+    def get_estimation_error(self) -> NDArray[np.float64]:
+        """
+        Calculate the difference between prior and posterior estimates.
+
+        Returns
+        -------
+        np.ndarray
+            The difference between posterior and prior state estimates.
+        """
+        return self.x_plus - self.x_minus
+
+    def get_trace_reduction(self) -> float:
+        """
+        Calculate the reduction in uncertainty as measured by trace of covariance.
+
+        Returns
+        -------
+        float
+            The relative reduction in trace of the state covariance matrix.
+        """
+        trace_minus = np.trace(self.Px_minus)
+        trace_plus = np.trace(self.Px_plus)
+        return (trace_minus - trace_plus) / trace_minus if trace_minus != 0 else 0.0
+
+    def to_dict(self) -> dict:
+        """
+        Convert the solution to a dictionary format.
+
+        Returns
+        -------
+        dict
+            Dictionary containing all solution components.
+        """
+        return {
+            "x_plus": self.x_plus,
+            "x_minus": self.x_minus,
+            "Pz_plus": self.Pz_plus,
+            "Pz_minus": self.Pz_minus,
+            "Px_plus": self.Px_plus,
+            "Px_minus": self.Px_minus,
+            "S": self.S,
+            "K": self.K,
+        }

KKF/systems.py

@@ -0,0 +1,200 @@
+"""Dynamical systems module.
+
+Defines classes and utilities for representing general dynamical systems
+with state transitions and measurement functions.
+"""
+
+from typing import Any, Callable, Optional
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.stats import rv_continuous
+
+
+class DynamicalSystem:
+    """
+    A class representing a general dynamical system with state and measurement equations.
+
+    This class encapsulates the dynamics, measurements, and associated probability
+    distributions for both state and measurement noise. Can be considered as:
+
+    Discrete time:
+        Dynamics: x_{k+1} = f(x_{k}) + w_{k}
+        Measurements: y_{k} = g(x_{k}) + v_{k}
+
+    Continous time:
+        Dynamics: x'(t) = f(x(t)) + w(t)
+        Measurements: y(t) = g(x(t)) + v(t)
+
+    Attributes
+    ----------
+    nx : int
+        Dimension of the state space.
+    ny : int
+        Dimension of the measurement/output space.
+    f : Callable
+        State transition function (dynamics).
+    g : Callable
+        Measurement/output function.
+    dist_X : rv_continuous
+        Probability distribution for initial states.
+    dist_dyn : rv_continuous
+        Probability distribution for dynamics noise.
+    dist_obs : rv_continuous
+        Probability distribution for measurement noise.
+    discrete_time : bool
+        Indicates if the system is in discrete time, if False the system is in continous time.
+
+    Notes
+    -----
+    The functions f and g should have the following signatures:
+    - f(x: ndarray, w: ndarray) -> ndarray
+    - g(x: ndarray, v: ndarray) -> ndarray
+    where:
+    - x is the state vector
+    - w is the process noise
+    - v is the measurement noise
+    """
+
+    def __init__(
+        self,
+        nx: int,
+        ny: int,
+        f: Callable[[NDArray[np.float64]], NDArray[np.float64]],
+        g: Callable[[NDArray[np.float64]], NDArray[np.float64]],
+        dist_X: rv_continuous,
+        dist_dyn: rv_continuous,
+        dist_obs: rv_continuous,
+        discrete_time: bool,
+    ) -> None:
+        if nx <= 0:
+            raise ValueError(f"State dimension nx must be positive, got {nx}")
+        self.nx = nx
+        self.ny = ny
+        self.f = f
+        self.g = g
+        self.dist_X = dist_X
+        self.dist_dyn = dist_dyn
+        self.dist_obs = dist_obs
+        self.discrete_time = discrete_time
+
+    def dynamics(self, x: NDArray[np.float64], w: NDArray[np.float64]) -> NDArray[np.float64]:
+        """
+        Apply the state transition function.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Current state vector.
+        w : np.ndarray
+            Process noise vector.
+
+        Returns
+        -------
+        np.ndarray
+            Next state vector.
+        """
+        return self.f(x) + w
+
+    def measurements(self, x: NDArray[np.float64], v: NDArray[np.float64]) -> NDArray[np.float64]:
+        """
+        Apply the measurement function.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Current state vector.
+        v : np.ndarray
+            Measurement noise vector.
+
+        Returns
+        -------
+        np.ndarray
+            Measurement/output vector.
+        """
+        return self.g(x) + v
+
+    def sample_state(self, size: int = 1) -> NDArray[np.float64]:
+        """
+        Sample from the state distribution.
+
+        Parameters
+        ----------
+        size : int, optional
+            Number of samples to draw. Default is 1.
+
+        Returns
+        -------
+        np.ndarray
+            Sampled state(s).
+        """
+        return self.dist_X.rvs(size=size).reshape((size, self.nx))
+
+
+def create_additive_system(
+    nx: int,
+    ny: int,
+    f: Callable[[NDArray[np.float64], NDArray[np.float64]], NDArray[np.float64]],
+    g: Callable[[NDArray[np.float64], NDArray[np.float64]], NDArray[np.float64]],
+    dist_X: rv_continuous,
+    dist_dyn: rv_continuous,
+    dist_obs: rv_continuous,
+    N_samples: int,
+    discrete_time: bool = True,
+) -> DynamicalSystem:
+    """
+    Create an additive dynamical system where noise is added to the state and observation functions.
+
+    Parameters
+    ----------
+    nx : int
+        Dimension of the state space.
+    ny : int
+        Dimension of the measurement space.
+    f : Callable
+        State transition function (with noise).
+    g : Callable
+        Measurement function (with noise).
+    dist_X : rv_continuous
+        Initial state distribution.
+    dist_dyn : rv_continuous
+        Dynamics noise distribution.
+    dist_obs : rv_continuous
+        Measurement noise distribution.
+    N_samples: int
+        Number of samples to generate empirical mean of dynamics and observation.
+    discrete_time : bool, optional
+        Indicates if the system is in discrete time. Default is True.
+
+    Returns
+    -------
+    DynamicalSystem
+        New dynamical system instance with additive noise.
+
+    Notes
+    -----
+    The resulting system has the form:
+    x[k+1] = f(x[k]) + w[k]
+    y[k] = g(x[k]) + v[k]
+    where w and v are noise terms.
+    """
+    new_f = lambda x: np.mean([f(x.reshape((len(x), 1)), dist_dyn.rvs(N_samples))], axis=1)
+    new_g = lambda x: np.mean([g(x.reshape((len(x), 1)), dist_obs.rvs(N_samples))], axis=1)
+
+    class DynDist:
+        def __init__(self) -> None:
+            pass
+
+        def rvs(self, x: NDArray[np.float64], size: int = 1) -> NDArray[np.float64]:
+            return f(x, dist_dyn.rvs(size=size)) - new_f(x)
+
+    class ObsDist:
+        def __init__(self) -> None:
+            pass
+
+        def rvs(self, x: NDArray[np.float64], size: int = 1) -> NDArray[np.float64]:
+            return g(x, dist_obs.rvs(size=size)) - new_g(x)
+
+    return DynamicalSystem(
+        nx, ny, new_f, new_g, dist_X, DynDist(), ObsDist(), discrete_time=discrete_time
+    )

kkf-0.3.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 diegoolguinw
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

kkf-0.3.0.dist-info/METADATA

@@ -0,0 +1,367 @@
+Metadata-Version: 2.2
+Name: kkf
+Version: 0.3.0
+Summary: KKF: Kernel Koopman Kalman Filter - A library for non-linear state estimation using kernel methods and Koopman operator theory
+Author-email: Diego Olguin-Wende <dolguin@dim.uchile.cl>
+License: MIT
+Project-URL: Homepage, https://github.com/diegoolguinw/kkf
+Project-URL: Documentation, https://kkf.readthedocs.io
+Project-URL: Repository, https://github.com/diegoolguinw/kkf.git
+Project-URL: Bug Tracker, https://github.com/diegoolguinw/kkf/issues
+Keywords: kalman-filter,koopman-operator,kernel-methods,edmd,nonlinear-dynamics,state-estimation,filtering
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: scipy>=1.7.0
+Requires-Dist: scikit-learn>=1.0.0
+Provides-Extra: viz
+Requires-Dist: matplotlib>=3.5.0; extra == "viz"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Requires-Dist: black>=22.0.0; extra == "dev"
+Requires-Dist: flake8>=4.0.0; extra == "dev"
+Requires-Dist: mypy>=0.990; extra == "dev"
+Requires-Dist: isort>=5.10.0; extra == "dev"
+Requires-Dist: pylint>=2.13.0; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinx>=4.5.0; extra == "docs"
+Requires-Dist: sphinx-rtd-theme>=1.0.0; extra == "docs"
+Requires-Dist: nbsphinx>=0.8.8; extra == "docs"
+Requires-Dist: ipython>=7.0.0; extra == "docs"
+Requires-Dist: jupyter>=1.0.0; extra == "docs"
+Provides-Extra: all
+Requires-Dist: kkf[dev,docs,viz]; extra == "all"
+
+# Kernel Koopman Kalman Filter (KKF)
+
+[![Tests](https://github.com/diegoolguinw/kkf/workflows/Tests%20&%20Code%20Quality/badge.svg)](https://github.com/diegoolguinw/kkf/actions)
+[![PyPI version](https://badge.fury.io/py/kkf.svg)](https://badge.fury.io/py/kkf)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.8+](https://img.shields.io/badge/Python-3.8+-blue.svg)](https://www.python.org/downloads/)
+
+A Python library for non-linear state estimation using Kernel-based Extended Dynamic Mode Decomposition (kEDMD) and Koopman operator theory.
+
+## What is KKF?
+
+KKF is a sophisticated filtering library that combines:
+
+- **Koopman Operator Theory**: Lifts nonlinear dynamics to a linear setting in feature space
+- **Kernel Methods**: Uses kernel functions for feature space approximation
+- **Kalman Filtering**: Applies optimal filtering in the lifted feature space
+- **Extended Dynamic Mode Decomposition**: Learns accurate linear approximations
+
+This enables accurate state estimation for highly nonlinear dynamical systems where traditional Kalman filters struggle.
+
+## Features
+
+✨ **Core Capabilities**
+- Kernel Extended Dynamic Mode Decomposition (kEDMD)
+- Non-linear Kalman Filter implementation
+- Support for arbitrary dynamical systems
+- Multiple kernel function options (RBF, Matérn, etc.)
+- Robust state estimation with uncertainty quantification
+- Covariance propagation in both state and feature spaces
+
+📊 **High-Quality Implementation**
+- Type hints throughout the codebase
+- Comprehensive error handling
+- Numerical stability optimizations
+- Efficient NumPy/SciPy operations
+
+🧪 **Well-Tested**
+- 95+ unit and integration tests
+- Edge case coverage
+- Multi-platform CI/CD
+- Coverage reporting
+
+📚 **Comprehensive Documentation**
+- Detailed docstrings
+- 3 complete working examples
+- Parameter optimization guide
+- API documentation
+
+## Installation
+
+### Current Release
+
+```bash
+pip install kkf
+```
+
+### Development Version
+
+```bash
+git clone https://github.com/diegoolguinw/kkf.git
+cd kkf
+pip install -e ".[dev]"
+```
+
+### Optional Dependencies
+
+```bash
+# For visualization examples
+pip install kkf[viz]
+
+# For development and testing
+pip install kkf[dev]
+
+# For building documentation
+pip install kkf[docs]
+
+# Install everything
+pip install kkf[all]
+```
+
+## Quick Start
+
+Here's a complete example of using KKF to estimate states in a SIR (Susceptible-Infected-Recovered) epidemiological model (see also [`examples/02_sir_epidemic_model.py`](examples/02_sir_epidemic_model.py)):
+
+```python
+import numpy as np
+from scipy import stats
+from sklearn.gaussian_process.kernels import Matern
+
+from KKF import systems, koopman, filter
+
+# Define SIR dynamics
+beta, gamma = 0.12, 0.04
+
+def f(x):  # State transition
+    return x + np.array([
+        -beta * x[0] * x[1],
+        beta * x[0] * x[1] - gamma * x[1],
+        gamma * x[1]
+    ])
+
+def g(x):  # Observation function (infected population)
+    return np.array([x[1]])
+
+# Setup system
+nx, ny, n_features = 3, 1, 50
+X_dist = stats.dirichlet(alpha=np.ones(nx))
+dyn_dist = stats.multivariate_normal(mean=np.zeros(nx), cov=1e-5 * np.eye(nx))
+obs_dist = stats.multivariate_normal(mean=np.zeros(ny), cov=1e-3 * np.eye(ny))
+
+system = DynamicalSystem(nx, ny, f, g, X_dist, dyn_dist, obs_dist, discrete_time=True)
+
+# Generate synthetic observations
+np.random.seed(42)
+n_timesteps = 100
+y_obs = np.random.randn(n_timesteps, ny) * 0.1 + 0.1
+
+# Apply Koopman Kalman Filter
+kernel = Matern(length_scale=1.0, nu=0.5)
+koopman_op = KoopmanOperator(kernel, system)
+
+x0_prior = np.array([0.8, 0.15, 0.05])
+initial_prior = stats.multivariate_normal(mean=x0_prior, cov=0.1 * np.eye(nx))
+
+solution = apply_koopman_kalman_filter(
+    koopman_operator=koopman_op,
+    observations=y_obs,
+    initial_distribution=initial_prior,
+    n_features=n_features,
+    optimize=False,
+    noise_samples=100
+)
+
+# Access results
+print(f"State estimates shape: {solution.x_plus.shape}")  # (n_timesteps, nx)
+print(f"State covariances shape: {solution.Px_plus.shape}")  # (n_timesteps, nx, nx)
+print(f"Mean estimation error: {solution.get_estimation_error().mean():.6f}")
+```
+
+## Examples
+
+The library includes three comprehensive examples:
+
+### 1. Basic Linear System (`examples/01_basic_linear_system.py`)
+
+Introduces core concepts with a simple linear system:
+- System definition
+- Synthetic data generation
+- Filter application
+- Confidence interval visualization
+
+**Run:** `python examples/01_basic_linear_system.py`
+
+### 2. SIR Epidemic Model (`examples/02_sir_epidemic_model.py`)
+
+Real-world application to disease modeling:
+- Nonlinear dynamics
+- Partial state observation
+- Population dynamics estimation
+- Epidemic curve prediction
+
+**Run:** `python examples/02_sir_epidemic_model.py`
+
+### 3. Parameter Exploration (`examples/03_advanced_parameter_exploration.py`)
+
+Advanced hyperparameter tuning:
+- Multiple kernel comparison
+- Feature dimension optimization
+- Performance analysis
+- Systematic search strategy
+
+**Run:** `python examples/03_advanced_parameter_exploration.py`
+
+For more details, see [examples/README.md](examples/README.md)
+
+## Documentation
+
+### API Reference
+
+```python
+# Core classes
+from KKF import DynamicalSystem, KoopmanOperator, KoopmanKalmanFilterSolution
+from KKF.filter import apply_koopman_kalman_filter
+
+# Utility functions
+from KKF import (
+    compute_initial_covariance,
+    compute_dynamics_covariance,
+    compute_observation_covariance
+)
+```
+
+### Theory
+
+The KKF algorithm works in three main steps:
+
+1. **Lifting**: Maps nonlinear states to a higher-dimensional feature space using kernel methods
+2. **Approximation**: Constructs a linear Koopman operator approximation via kEDMD
+3. **Filtering**: Applies Kalman filtering in the lifted space, then maps back to original coordinates
+
+### Key Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `n_features` | Dimension of lifted feature space | 50 |
+| `kernel` | Kernel function (RBF, Matérn, etc.) | RBF(1.0) |
+| `optimize` | Whether to optimize kernel parameters | False |
+| `noise_samples` | Monte Carlo samples for covariance | 100 |
+
+## System Requirements
+
+- Python 3.8 or higher
+- NumPy ≥ 1.20.0
+- SciPy ≥ 1.7.0
+- scikit-learn ≥ 1.0.0
+- Matplotlib ≥ 3.5.0 (optional, for visualization)
+
+## Advanced Usage
+
+### Custom Kernel Functions
+
+```python
+from sklearn.gaussian_process.kernels import (
+    RBF, Matern, ExpSineSquared, DotProduct, ConstantKernel
+)
+
+# Different kernel choices
+kernel1 = RBF(length_scale=1.0)
+kernel2 = Matern(length_scale=1.0, nu=2.5)
+kernel3 = ExpSineSquared(length_scale=1.0, periodicity=1.0)
+kernel4 = ConstantKernel(constant_value=1.0) * Matern(nu=1.5)
+```
+
+### Manual Kernel Optimization
+
+```python
+solution = apply_koopman_kalman_filter(
+    koopman_operator=koopman_op,
+    observations=y_obs,
+    initial_distribution=initial_prior,
+    n_features=50,
+    optimize=True,           # Enable kernel optimization
+    n_restarts_optimizer=20  # Number of optimization restarts
+)
+```
+
+### Extracting Uncertainty Estimates
+
+```python
+# Posterior state covariance
+Px_plus = solution.Px_plus  # (n_timesteps, nx, nx)
+
+# Compute confidence intervals
+std_devs = np.sqrt(np.diagonal(Px_plus, axis1=1, axis2=2))
+ci_lower = solution.x_plus - 1.96 * std_devs
+ci_upper = solution.x_plus + 1.96 * std_devs
+```
+
+## Testing
+
+Run the comprehensive test suite:
+
+```bash
+# All tests
+pytest tests/
+
+# Verbose output
+pytest tests/ -v
+
+# Specific test file
+pytest tests/test_core.py
+
+# Skip slow tests
+pytest tests/ -m "not slow"
+
+# With coverage report
+pytest tests/ --cov=KKF --cov-report=html
+```
+
+## Contributing
+
+We welcome contributions! See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to:
+- Report bugs
+- Suggest features
+- Submit pull requests
+- Improve documentation
+
+## Code of Conduct
+
+Please note that this project is released with a [Contributor Code of Conduct](CODE_OF_CONDUCT.md). By participating you agree to abide by its terms.
+
+## Citation
+
+If you use KKF in your research, please cite:
+
+```bibtex
+@software{kkf2024,
+  title={KKF: Kernel Koopman Kalman Filter},
+  author={Olguin-Wende, Diego},
+  year={2024},
+  url={https://github.com/diegoolguinw/kkf}
+}
+```
+
+## License
+
+This project is licensed under the MIT License - see [LICENSE](LICENSE) file for details.
+
+## Changelog
+
+See [CHANGELOG.md](CHANGELOG.md) for version history and release notes.
+
+## Support & Contact
+
+- **Issues & Bug Reports**: [GitHub Issues](https://github.com/diegoolguinw/kkf/issues)
+- **Discussions**: [GitHub Discussions](https://github.com/diegoolguinw/kkf/discussions)
+- **Email**: dolguin@dim.uchile.cl

kkf-0.3.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+KKF/__init__.py,sha256=3W3pPuvwh-uVF9eHZ6o6xzSRLRHQEwkfXQdv_4B4fI4,1412
+KKF/covariances.py,sha256=7ET1-wvbDnYBIc8v_i_aO7k06wZwHXI7Cxj_HGBbIkQ,5735
+KKF/filter.py,sha256=evtaZFj9mbOKie6wUKesC50rX1-Dpvkk-Pf4uwOd3m0,8095
+KKF/koopman.py,sha256=KobfN1WPW-kxiCl_4-9v5NoGX4TTfChw8s8ngM0Evp0,5480
+KKF/solution.py,sha256=AqglUKkr6XiwEXnfA8YCX4edTaF1cw3ceLGlDOAGEOk,4425
+KKF/systems.py,sha256=FmKc5BuRVEPxKF94y1FAi_-uvmCteNFYuF3ps7NRL00,5798
+KKF/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kkf-0.3.0.dist-info/LICENSE,sha256=6WnhMjAMnLiDcClR0tHyWuaX8lEs0MSn5R51JqXL7Ts,1069
+kkf-0.3.0.dist-info/METADATA,sha256=t_ZymB92I2o0Iz22L7WcOq0JaXAV3gy_xc1Hk7wBBVI,10888
+kkf-0.3.0.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+kkf-0.3.0.dist-info/top_level.txt,sha256=yCY9tBJ5DJeTknZPyFujr9XvomHsm-7oP4ABMz348pw,4
+kkf-0.3.0.dist-info/RECORD,,

kkf-0.3.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (76.1.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

kkf-0.3.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+KKF