Metadata-Version: 2.4
Name: koopman
Version: 0.1.0
Summary: Exact custom-gradient ops for the linear Koopman recurrence (PyTorch + JAX)
Author-email: Matthew Tamayo-Rios <matthew@geekbeast.com>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/geekbeast/koopman
Project-URL: Repository, https://github.com/geekbeast/koopman
Keywords: machine-learning,koopman,custom-gradient,autodiff,jax,pytorch,sequence,time-series
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Natural Language :: English
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: Microsoft :: Windows
Classifier: Operating System :: Unix
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Provides-Extra: torch
Requires-Dist: torch>=2.2.0; extra == "torch"
Provides-Extra: jax-cpu
Requires-Dist: jax[cpu]>=0.4.34; extra == "jax-cpu"
Provides-Extra: jax-gpu
Requires-Dist: jax[cuda12]>=0.4.34; extra == "jax-gpu"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Provides-Extra: all
Requires-Dist: torch>=2.2.0; extra == "all"
Requires-Dist: jax>=0.4.34; extra == "all"
Dynamic: license-file

# koopman

[![CI](https://github.com/geekbeast/koopman/actions/workflows/ci.yml/badge.svg)](https://github.com/geekbeast/koopman/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/koopman.svg)](https://pypi.org/project/koopman/)

Exact, closed-form custom-gradient ops for the linear **Koopman recurrence**
`z_{t+1} = K @ z_t`. Computes the full state sequence

```
Z = [z_0, K z_0, K^2 z_0, ..., K^T z_0]
```

with a single forward scan and an **exact adjoint backward** — `O(T n^2)` time
and `O(T n)` memory — instead of taping every per-step matmul through autodiff.
Available in a **PyTorch** flavor (`torch.autograd.Function`) and a **JAX**
flavor (`jax.custom_vjp`).

## Installation

Install only the backend you need:

```bash
pip install koopman[torch]      # PyTorch (CPU or GPU)
pip install koopman[jax-cpu]    # JAX on CPU
pip install koopman[jax-gpu]    # JAX on GPU (CUDA 12)
```

For a **CPU-only PyTorch** wheel (e.g. on free CI runners), install torch from
the CPU index first, then the package:

```bash
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install koopman
```

Requires Python 3.10+.

## Usage

### PyTorch

```python
import torch
from koopman.torch import linear_powers_sequence

n, T = 16, 1024
K = torch.randn(n, n, dtype=torch.float64, requires_grad=True)
z0 = torch.randn(n, dtype=torch.float64, requires_grad=True)

Z = linear_powers_sequence(K, z0, T)   # [T+1, n]
Z.sum().backward()                     # exact grad_K, grad_z0 via adjoint recurrence
```

There is also a `torch.nn.Module` wrapper, `LinearPowersSequence(T)`, and the
raw `LinearPowersSequenceFn` autograd function.

### JAX

```python
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from koopman.jax import linear_powers_sequence

n, T = 16, 1024
K = jax.random.normal(jax.random.PRNGKey(0), (n, n))
z0 = jax.random.normal(jax.random.PRNGKey(1), (n,))

Z = linear_powers_sequence(K, z0, T)   # [T+1, n]
loss = lambda K_, z0_: jnp.sum(linear_powers_sequence(K_, z0_, T))
gK, gz0 = jax.grad(loss, argnums=(0, 1))(K, z0)   # jit/vmap-friendly
```

## Why a custom gradient?

Naive autodiff through the Python/`lax.scan` loop tapes all `T` matmuls and
their activations, costing `O(T n^2)` memory. The backward here is the
closed-form adjoint recurrence:

```
a_T      = grad_Z[T]
a_t      = grad_Z[t] + Kᵀ @ a_{t+1}        (t = T-1, ..., 0)
grad_K  += outer(a_{t+1}, z_t)
grad_z0  = a_0
```

It reuses the forward states (`O(T n)` memory) and avoids materializing the
autodiff tape. See `benchmarks/` for the comparison against naive autograd and
`expm`-based baselines.

## Development

```bash
pip install -e ".[dev]"
pytest tests/test_linear_powers_torch.py -v   # PyTorch
JAX_PLATFORMS=cpu pytest tests/test_linear_powers_jax.py -v   # JAX
```

## License

Apache-2.0.
