diffkalman 0.1.0__tar.gz

diffkalman-0.1.0/.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

diffkalman-0.1.0/.python-version

@@ -0,0 +1 @@
+3.11

diffkalman-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 HadesX
+
+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.

diffkalman-0.1.0/PKG-INFO

@@ -0,0 +1,183 @@
+Metadata-Version: 2.4
+Name: diffkalman
+Version: 0.1.0
+Summary: A diffrentiable kalman filter library for auto-tuning kalman filters.
+Author-email: hades <nischalbhattaraipi@gmail.com>
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: numpy>=2.2.1
+Requires-Dist: tqdm>=4.67.1
+Provides-Extra: cpu
+Requires-Dist: torch>=2.5.1; extra == 'cpu'
+Provides-Extra: cu124
+Requires-Dist: torch>=2.5.1; extra == 'cu124'
+Description-Content-Type: text/markdown
+
+# Differentiable Kalman Filter
+
+A PyTorch-based implementation of a differentiable Kalman Filter designed for both linear and non-linear dynamical systems with Gaussian noise. This module seamlessly integrates with neural networks, enabling learnable dynamics, observation, and noise models optimized through Stochastic Variational Inference (SVI).
+
+## Features
+
+- **Fully Differentiable**: End-to-end differentiable implementation compatible with PyTorch's autograd
+- **Flexible Models**: Support for both linear and non-linear state transition and observation models
+- **Neural Network Integration**: Models can be parameterized using neural networks
+- **Automatic Jacobian Computation**: Utilizes PyTorch's autograd for derivative calculations
+- **Monte Carlo Sampling**: Supports evaluation of expected joint log-likelihood to perform Expectation-Maximization (EM) learning
+- **Rauch-Tung-Striebel Smoothing**: Implements forward-backward smoothing for improved state estimation using RTS algorithm
+
+## Installation
+
+```bash
+pip install torch  # Required dependency
+# Add your package installation command here
+```
+
+## Quick Start
+
+Here's a simple example of using the Differentiable Kalman Filter:
+
+```python
+import torch
+from diffkalman import DifferentiableKalmanFilter
+from diffkalman.utils import SymmetricPositiveDefiniteMatrix
+from diffkalman.em_loop import em_updates
+
+# Define custom state transition and observation functions
+class StateTransition(torch.nn.Module):
+    def forward(self, x, *args):
+        # Your state transition logic here
+        return x
+
+class ObservationModel(torch.nn.Module):
+    def forward(self, x, *args):
+        # Your observation logic here
+        return x
+
+# Initialize the filter
+f = StateTransition()
+h = ObservationModel()
+Q = SymmetricPositiveDefiniteMatrix(dim=4, trainable=True)
+R = SymmetricPositiveDefiniteMatrix(dim=2, trainable=True)
+kalman_filter = DifferentiableKalmanFilter(
+    dim_x=4,  # State dimension
+    dim_z=2,  # Observation dimension
+    f=f,      # State transition function
+    h=h       # Observation function
+)
+
+# Run the filter
+results = kalman_filter.sequence_filter(
+    z_seq=observations,           # Shape: (T, dim_z)
+    x0=initial_state,            # Shape: (dim_x,)
+    P0=initial_covariance,       # Shape: (dim_x, dim_x)
+    Q=Q().repeat(len(observations), 1, 1),  # Shape: (T, dim_x, dim_x)
+    R=R().repeat(len(observations), 1, 1)   # Shape: (T, dim_z, dim_z)
+)
+```
+
+## Detailed Usage
+
+### State Estimation
+
+The module provides three main estimation methods:
+
+1. **Filtering**: Forward pass only
+```python
+filtered_results = kalman_filter.sequence_filter(
+    z_seq=observations,
+    x0=initial_state,
+    P0=initial_covariance,
+    Q=process_noise,
+    R=observation_noise
+)
+```
+
+2. **Smoothing**: Forward-backward pass
+```python
+smoothed_results = kalman_filter.sequence_smooth(
+    z_seq=observations,
+    x0=initial_state,
+    P0=initial_covariance,
+    Q=process_noise,
+    R=observation_noise
+)
+```
+
+3. **Single-step Prediction**: For real-time applications
+```python
+step_result = kalman_filter.predict_update(
+    z=current_observation,
+    x=current_state,
+    P=current_covariance,
+    Q=process_noise,
+    R=observation_noise
+)
+```
+
+### Parameter Learning
+
+The module supports learning model parameters through using backpropagation using the negative expected joint log-likelihood of the 
+data as the loss function.
+
+```python
+# Define optimizer
+optimizer = torch.optim.Adam(params=[
+    {'params': kalman_filter.f.parameters()},
+    {'params': kalman_filter.h.parameters()},
+    {'params': Q.parameters()},
+    {'params': R.parameters()}
+]
+
+NUM_EPOCHS = 10
+NUM_CYCLES = 10
+
+# Run the EM loop
+marginal_likelihoods = em_updates(
+    kalman_filter=kalman_filter,
+    z_seq=observations,
+    x0=initial_state,
+    P0=initial_covariance,
+    Q=Q,
+    R=R,
+    optimizer=optimizer,
+    num_cycles=NUM_CYCLES,
+    num_epochs=NUM_EPOCHS
+) 
+    
+```
+
+## API Reference
+
+### DifferentiableKalmanFilter
+
+Main class implementing the Kalman Filter algorithm.
+
+#### Constructor Parameters:
+- `dim_x` (int): State space dimension
+- `dim_z` (int): Observation space dimension
+- `f` (nn.Module): State transition function
+- `h` (nn.Module): Observation function
+- `mc_samples` (int, optional): Number of Monte Carlo samples for log-likelihood estimation
+
+#### Key Methods:
+- `predict`: State prediction step
+- `update`: Measurement update step
+- `predict_update`: Combined prediction and update
+- `sequence_filter`: Full sequence filtering
+- `sequence_smooth`: Full sequence smoothing
+- `marginal_log_likelihood`: Compute marginal log-likelihood
+- `monte_carlo_expected_joint_log_likekihood`: Estimate expected joint log-likelihood
+
+## Requirements
+
+- PyTorch >= 1.9.0
+- Python >= 3.7
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

diffkalman-0.1.0/README.md

@@ -0,0 +1,168 @@
+# Differentiable Kalman Filter
+
+A PyTorch-based implementation of a differentiable Kalman Filter designed for both linear and non-linear dynamical systems with Gaussian noise. This module seamlessly integrates with neural networks, enabling learnable dynamics, observation, and noise models optimized through Stochastic Variational Inference (SVI).
+
+## Features
+
+- **Fully Differentiable**: End-to-end differentiable implementation compatible with PyTorch's autograd
+- **Flexible Models**: Support for both linear and non-linear state transition and observation models
+- **Neural Network Integration**: Models can be parameterized using neural networks
+- **Automatic Jacobian Computation**: Utilizes PyTorch's autograd for derivative calculations
+- **Monte Carlo Sampling**: Supports evaluation of expected joint log-likelihood to perform Expectation-Maximization (EM) learning
+- **Rauch-Tung-Striebel Smoothing**: Implements forward-backward smoothing for improved state estimation using RTS algorithm
+
+## Installation
+
+```bash
+pip install torch  # Required dependency
+# Add your package installation command here
+```
+
+## Quick Start
+
+Here's a simple example of using the Differentiable Kalman Filter:
+
+```python
+import torch
+from diffkalman import DifferentiableKalmanFilter
+from diffkalman.utils import SymmetricPositiveDefiniteMatrix
+from diffkalman.em_loop import em_updates
+
+# Define custom state transition and observation functions
+class StateTransition(torch.nn.Module):
+    def forward(self, x, *args):
+        # Your state transition logic here
+        return x
+
+class ObservationModel(torch.nn.Module):
+    def forward(self, x, *args):
+        # Your observation logic here
+        return x
+
+# Initialize the filter
+f = StateTransition()
+h = ObservationModel()
+Q = SymmetricPositiveDefiniteMatrix(dim=4, trainable=True)
+R = SymmetricPositiveDefiniteMatrix(dim=2, trainable=True)
+kalman_filter = DifferentiableKalmanFilter(
+    dim_x=4,  # State dimension
+    dim_z=2,  # Observation dimension
+    f=f,      # State transition function
+    h=h       # Observation function
+)
+
+# Run the filter
+results = kalman_filter.sequence_filter(
+    z_seq=observations,           # Shape: (T, dim_z)
+    x0=initial_state,            # Shape: (dim_x,)
+    P0=initial_covariance,       # Shape: (dim_x, dim_x)
+    Q=Q().repeat(len(observations), 1, 1),  # Shape: (T, dim_x, dim_x)
+    R=R().repeat(len(observations), 1, 1)   # Shape: (T, dim_z, dim_z)
+)
+```
+
+## Detailed Usage
+
+### State Estimation
+
+The module provides three main estimation methods:
+
+1. **Filtering**: Forward pass only
+```python
+filtered_results = kalman_filter.sequence_filter(
+    z_seq=observations,
+    x0=initial_state,
+    P0=initial_covariance,
+    Q=process_noise,
+    R=observation_noise
+)
+```
+
+2. **Smoothing**: Forward-backward pass
+```python
+smoothed_results = kalman_filter.sequence_smooth(
+    z_seq=observations,
+    x0=initial_state,
+    P0=initial_covariance,
+    Q=process_noise,
+    R=observation_noise
+)
+```
+
+3. **Single-step Prediction**: For real-time applications
+```python
+step_result = kalman_filter.predict_update(
+    z=current_observation,
+    x=current_state,
+    P=current_covariance,
+    Q=process_noise,
+    R=observation_noise
+)
+```
+
+### Parameter Learning
+
+The module supports learning model parameters through using backpropagation using the negative expected joint log-likelihood of the 
+data as the loss function.
+
+```python
+# Define optimizer
+optimizer = torch.optim.Adam(params=[
+    {'params': kalman_filter.f.parameters()},
+    {'params': kalman_filter.h.parameters()},
+    {'params': Q.parameters()},
+    {'params': R.parameters()}
+]
+
+NUM_EPOCHS = 10
+NUM_CYCLES = 10
+
+# Run the EM loop
+marginal_likelihoods = em_updates(
+    kalman_filter=kalman_filter,
+    z_seq=observations,
+    x0=initial_state,
+    P0=initial_covariance,
+    Q=Q,
+    R=R,
+    optimizer=optimizer,
+    num_cycles=NUM_CYCLES,
+    num_epochs=NUM_EPOCHS
+) 
+    
+```
+
+## API Reference
+
+### DifferentiableKalmanFilter
+
+Main class implementing the Kalman Filter algorithm.
+
+#### Constructor Parameters:
+- `dim_x` (int): State space dimension
+- `dim_z` (int): Observation space dimension
+- `f` (nn.Module): State transition function
+- `h` (nn.Module): Observation function
+- `mc_samples` (int, optional): Number of Monte Carlo samples for log-likelihood estimation
+
+#### Key Methods:
+- `predict`: State prediction step
+- `update`: Measurement update step
+- `predict_update`: Combined prediction and update
+- `sequence_filter`: Full sequence filtering
+- `sequence_smooth`: Full sequence smoothing
+- `marginal_log_likelihood`: Compute marginal log-likelihood
+- `monte_carlo_expected_joint_log_likekihood`: Estimate expected joint log-likelihood
+
+## Requirements
+
+- PyTorch >= 1.9.0
+- Python >= 3.7
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

diffkalman-0.1.0/pyproject.toml

@@ -0,0 +1,57 @@
+[project]
+name = "diffkalman"
+version = "0.1.0"
+description = "A diffrentiable kalman filter library for auto-tuning kalman filters."
+readme = "README.md"
+authors = [
+    { name = "hades", email = "nischalbhattaraipi@gmail.com" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    "numpy>=2.2.1",
+    "tqdm>=4.67.1",
+]
+
+[project.optional-dependencies]
+cpu = [
+  "torch>=2.5.1",
+]
+cu124 = [
+  "torch>=2.5.1",
+]
+
+[tool.uv]
+conflicts = [
+  [
+    { extra = "cpu" },
+    { extra = "cu124" },
+  ],
+]
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cpu", extra = "cpu" },
+  { index = "pytorch-cu124", extra = "cu124" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu124"
+url = "https://download.pytorch.org/whl/cu124"
+explicit = true
+
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "ipykernel>=6.29.5",
+    "matplotlib>=3.10.0",
+    "pandas>=2.2.3",
+]

diffkalman-0.1.0/src/diffkalman/__init__.py

@@ -0,0 +1,6 @@
+"""Top level module for diffkalman package."""
+from .filter import DiffrentiableKalmanFilter
+from .em_loop import em_updates
+
+
+__all__ = ['DiffrentiableKalmanFilter', 'em_updates']

diffkalman-0.1.0/src/diffkalman/em_loop.py

@@ -0,0 +1,115 @@
+"""The EM loop module that implements the EM algorithm for the Differentiable Kalman Filter."""
+
+from .filter import DiffrentiableKalmanFilter
+from .utils import SymmetricPositiveDefiniteMatrix
+import torch
+import torch.nn as nn
+
+
+__all__ = [
+    "em_updates",
+]
+
+
+def em_updates(
+    dkf: DiffrentiableKalmanFilter,
+    z_seq: torch.Tensor,
+    x0: torch.Tensor,
+    P0: torch.Tensor,
+    Q: SymmetricPositiveDefiniteMatrix,
+    R: SymmetricPositiveDefiniteMatrix,
+    optimizer: torch.optim.Optimizer,
+    lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
+    num_cycles: int = 20,
+    num_epochs: int = 100,
+    h_args: tuple = (),
+    f_args: tuple = (),
+) -> dict:
+    """A sample implementation of the EM algorithm for the Differentiable Kalman Filter.
+
+    Args:
+        z (torch.Tensor): The noisy measurements sequence. Dimension: (seq_len, obs_dim)
+        x0 (torch.Tensor): The initial state vector. Dimension: (state_dim,)
+        P0 (torch.Tensor): The initial covariance matrix of the state vector. Dimension: (state_dim, state_dim)
+        Q (SymmetricPositiveDefiniteMatrix): The process noise covariance matrix module.
+        R (SymmetricPositiveDefiniteMatrix): The measurement noise covariance matrix module.
+        optimizer (torch.optim.Optimizer): The optimizer
+        lr_scheduler (torch.optim.lr_scheduler._LRScheduler): The learning rate scheduler
+        num_cycles (int): The number of cycles. Default: 20
+        num_epochs (int): The number of epochs. Default: 100
+        h_args (tuple): Additional arguments for the measurement model. Default: ()
+        f_args (tuple): Additional arguments for the transition function. Default: ()
+
+    Returns:
+        dict: log likelihoods of the model with respect to the number of epochs and cycles
+
+    Note:
+        - Use the SymmetricPositiveDefiniteMatrix module to ensure that the process noise covariance matrix Q and the measurement noise covariance matrix R are symmetric and positive definite.
+        - The optimizer and the learning rate scheduler should be initialized before calling this function.
+        - The measurement model and the transition function should be differentiable torch modules.
+    """
+    likelihoods = torch.zeros(num_epochs, num_cycles)
+
+    # Perform EM updates for num_epochs
+    for e in range(num_epochs):
+
+        ## The E-step
+        ## Without gradients tracking, get the posterior state distribution wrt current values of the parameters
+        with torch.no_grad():
+            posterior = dkf.sequence_smooth(
+                z_seq=z_seq,
+                x0=x0,
+                P0=P0,
+                Q=Q().repeat(len(z_seq), 1, 1),
+                R=R().repeat(len(z_seq), 1, 1),
+                f_args=f_args,
+                h_args=h_args,
+            )
+
+        ## The M-step (Update the parameters) with respect to the current posterior state distribution for num_cycles
+        for c in range(num_cycles):
+            # Zero the gradients
+            optimizer.zero_grad()
+
+            # Compute the marginal likelihood for logging
+            with torch.no_grad():
+                marginal_likelihood = dkf.marginal_log_likelihood(
+                    z_seq=z_seq,
+                    x0=x0,
+                    P0=P0,
+                    Q=Q().repeat(len(z_seq), 1, 1),
+                    R=R().repeat(len(z_seq), 1, 1),
+                    f_args=f_args,
+                    h_args=h_args,
+                )
+                likelihoods[e, c] = marginal_likelihood
+
+            # Perform the forward pass i.e compute the expected complete joint log-likelihood with respect previous posterior state distribution and current parameters
+            complete_log_likelihood = dkf.monte_carlo_expected_joint_log_likekihood(
+                z_seq=z_seq,
+                x0=x0,
+                P0=P0,
+                # below represents the posterior state distribution
+                x0_smoothed=posterior["x0_smoothed"],
+                P0_smoothed=posterior["P0_smoothed"],
+                x_smoothed=posterior["x_smoothed"],
+                P_smoothed=posterior["P_smoothed"],
+                Q_seq=Q().repeat(len(z_seq), 1, 1),
+                R_seq=R().repeat(len(z_seq), 1, 1),
+                f_args=f_args,
+                h_args=h_args,
+            )
+
+            # Update the parameters
+            (-complete_log_likelihood).backward()
+            optimizer.step()
+            lr_scheduler.step()
+
+            # Print the log likelihood
+            print(
+                f"Epoch {e + 1}/{num_epochs} Cycle {c + 1}/{num_cycles} Log Likelihood: {marginal_likelihood.item()}"
+            )
+
+    return {
+        "likelihoods": likelihoods,
+    }