kkf 0.3.0__tar.gz

kkf-0.3.0/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-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-0.3.0/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-0.3.0/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_