Metadata-Version: 2.4
Name: koopman-graph-evolver
Version: 0.1.1
Summary: Geometry-Preserving Latent Dynamics & SE(3)-Equivariant Koopman Operators for Long-Horizon Graph Simulation
Project-URL: Homepage, https://github.com/Abhilash437/Koopman-Graph-Evolver
Project-URL: Repository, https://github.com/Abhilash437/Koopman-Graph-Evolver
Project-URL: Paper, https://github.com/Abhilash437/Koopman-Graph-Evolver/blob/main/paper/main.tex
Project-URL: Bug Tracker, https://github.com/Abhilash437/Koopman-Graph-Evolver/issues
Author-email: Abhilash Hathwar <abhilashhathwar@gmail.com>
License: MIT
License-File: LICENSE
Keywords: equivariant-neural-networks,geometric-deep-learning,graph-neural-networks,koopman-operator,molecular-dynamics,n-body-simulation,physics-informed-machine-learning
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
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 :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.9
Requires-Dist: networkx>=3.0
Requires-Dist: numpy>=1.21.0
Requires-Dist: scikit-learn>=1.0.0
Requires-Dist: scipy>=1.7.0
Requires-Dist: torch-geometric>=2.3.0
Requires-Dist: torch>=2.0.0
Provides-Extra: all
Requires-Dist: build>=1.0.0; extra == 'all'
Requires-Dist: kagglehub; extra == 'all'
Requires-Dist: matplotlib>=3.5.0; extra == 'all'
Requires-Dist: pytest-cov>=4.0.0; extra == 'all'
Requires-Dist: pytest>=7.0.0; extra == 'all'
Requires-Dist: ruff>=0.1.0; extra == 'all'
Requires-Dist: streamlit>=1.20.0; extra == 'all'
Requires-Dist: twine>=4.0.0; extra == 'all'
Provides-Extra: data
Requires-Dist: kagglehub; extra == 'data'
Provides-Extra: dev
Requires-Dist: build>=1.0.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Requires-Dist: twine>=4.0.0; extra == 'dev'
Provides-Extra: gui
Requires-Dist: matplotlib>=3.5.0; extra == 'gui'
Requires-Dist: streamlit>=1.20.0; extra == 'gui'
Description-Content-Type: text/markdown

# Beyond MSE: Geometry-Preserving Latent Dynamics for Long-Horizon Graph Simulation

Official implementation of the paper **"Beyond MSE: Geometry-Preserving Latent Dynamics for Long-Horizon Graph Simulation"**.

---

## Executive Summary & Abstract

Graph neural networks (GNNs) achieve high short-horizon accuracy for physical simulation, yet accumulate severe latent drift over long-horizon autoregressive rollouts. Unconstrained temporal transitions (e.g., GRUs) allow latent representation norms to progressively expand or contract, producing trajectories with competitive pointwise rollout MSE despite catastrophic physical structural breakdown (bond stretching, angle distortion, and centroid collapse).

We introduce the **Koopman Graph Evolver (KGE)** and its SE(3)-equivariant extension **E-GKN**, which replace unconstrained recurrent transitions with an orthogonal Koopman operator acting in latent space. Parameterizing the transition via a matrix exponential $K = \exp(A_{\text{skew}} \Delta t)$ guarantees $K \in \text{SO}(n)$, preserving latent norm and volume by construction ($R_{\text{norm}} = 1.0000$).

Evaluated across **14 physical systems** (8 MD17 molecules, 4 MD22 macromolecules, and 2 N-body particle systems), geometry-preserving latent transitions yield statistically significant reductions in structural drift ($p \le 3.1 \times 10^{-4}$, Wilcoxon signed-rank test), maintaining physical coordinate edge length ratios ($R_{\text{edge}} \approx 1.0$) across extended rollouts.

---

## Key Contributions & Mathematical Framework

1. **Characterization of Structural Drift & The MSE Paradox:** We show that standard rollout MSE fails to reflect physical degradation, rewarding models that expand uniformly or collapse toward spatial centroids. Physical topology metrics (bond, angle, torsion drift, and coordinate edge ratios) are required for faithful physical evaluation.
2. **Volume-Preserving Koopman Transitions:** Enforcing $K = \exp(A_{\text{skew}} \Delta t) \in \text{SO}(n)$ guarantees:
   - $K^T K = I$ (Orthogonality)
   - $\det(K) = 1$ (Orientation and volume preservation)
   - $\|Kz\|_2 = \|z\|_2$ (Norm preservation in latent feature space)
3. **Physical Regularization in 3D Space:** While the non-linear spatial decoder does not mathematically mandate 3D coordinate edge ratios $R_{\text{edge}} = 1.0$ identically, latent orthogonality strongly regularizes spatial decoding, keeping $R_{\text{edge}} \approx 1.0$ ($0.9416$–$1.0112$) across extended rollouts.
4. **SE(3)-Equivariant Extension (E-GKN):** Augmenting equivariant message passing with shared node-local Koopman transitions prevents numerical divergence ($10^{27}$ / NaNs) present in standard EGNNs on large flexible macromolecules.

---

## Empirical Benchmark Results (14 Physical Systems)

### 1. Multi-Seed Robustness (Averaged Over Seeds {42, 1337, 2026})

| System | Model | Rollout MSE (29-step) | Bond Drift (Å) | Angle Drift (°) | Torsion Drift (°) | Physical Coord Edge Ratio ($R_{\text{edge}}$) |
|:---|:---|:---:|:---:|:---:|:---:|:---:|
| **aspirin** | Flat Koopman | 0.0715 ± 0.008 | 0.0816 ± 0.004 | 4.55 ± 0.45 | 5.38 ± 0.32 | 0.9707 |
| | **Graph Koopman** | 0.2411 ± 0.003 | **0.0045 ± 0.004** | **0.09 ± 0.02** | **0.15 ± 0.06** | **0.9974** |
| | Graph GRU | 0.1388 ± 0.031 | 0.0689 ± 0.012 | 5.49 ± 1.13 | 6.42 ± 0.62 | 0.9584 |
| **malonaldehyde** | Flat Koopman | 0.4002 ± 0.003 | 0.1699 ± 0.022 | 10.51 ± 1.29 | 16.73 ± 1.28 | 0.9367 |
| | **Graph Koopman** | 0.9151 ± 0.048 | **0.0905 ± 0.030** | **0.46 ± 0.34** | **0.83 ± 0.40** | **0.9416** |
| | Graph GRU | 0.3532 ± 0.005 | 0.0981 ± 0.005 | 3.72 ± 0.66 | 4.19 ± 1.36 | 0.9262 |
| **at-at** | Flat Koopman | 3.6341 ± 1.041 | 0.4499 ± 0.146 | 37.14 ± 9.74 | 48.21 ± 11.20 | 0.9972 |
| | **Graph Koopman** | 6.2917 ± 0.514 | **0.0240 ± 0.010** | **0.63 ± 0.26** | **1.15 ± 0.48** | **0.9868** |
| | Graph GRU | 2.6390 ± 0.158 | 0.2514 ± 0.038 | 17.98 ± 1.96 | 26.37 ± 3.13 | 0.8675 |
| **springs** | Flat Koopman | 0.1756 ± 0.001 | 0.1075 ± 0.026 | 14.55 ± 2.87 | 29.38 ± 5.56 | 0.9827 |
| | **Graph Koopman** | 0.1764 ± 0.003 | **0.0248 ± 0.009** | **2.59 ± 1.01** | **6.05 ± 1.79** | **1.0112** |
| | Graph GRU | 0.0531 ± 0.002 | 0.6167 ± 0.013 | 46.51 ± 0.87 | 80.21 ± 0.95 | 1.6290 |

### 2. Statistical Significance Across All 14 Systems

One-sided Wilcoxon signed-rank test results comparing Graph Koopman (KGE) vs. Graph GRU (G-GRU) across 14 physical systems:

| Metric | KGE Win Rate | Wilcoxon Statistic | p-value |
|:---|:---:|:---:|:---:|
| **Bond Drift (Å)** | 13/14 | 1.5 | $3.05 \times 10^{-4}$ |
| **Angle Drift (°)** | 14/14 | 0.0 | $6.10 \times 10^{-5}$ |
| **Torsion Drift (°)** | 14/14 | 0.0 | $6.10 \times 10^{-5}$ |
| **Latent Norm Ratio \|R_norm - 1\|** | 14/14 | 0.0 | $6.10 \times 10^{-5}$ |

---

## Quickstart & Installation

[![PyPI Version](https://img.shields.io/pypi/v/koopman-graph-evolver.svg?v=1)](https://pypi.org/project/koopman-graph-evolver/)
[![Python Versions](https://img.shields.io/pypi/pyversions/koopman-graph-evolver.svg?v=1)](https://pypi.org/project/koopman-graph-evolver/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

### 1. Install via PyPI

```bash
# Core package (PyTorch & PyG models + PhysicsEval suite)
pip install koopman-graph-evolver

# Optional extras: GUI dashboard or dataset downloaders
pip install "koopman-graph-evolver[gui,data]"
```

---

### 2. Python API Quickstart

**Use Graph-Aware Koopman Net or E-GKN in your PyTorch code:**

```python
import torch
from koopman_evolver import GraphAwareKoopmanNet, EGKN, PhysicsEval

# 1. Define graph connectivity (e.g. molecular bond edges)
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)

# 2. Instantiate Graph-Aware Koopman Net or Equivariant Koopman Net (E-GKN)
model = GraphAwareKoopmanNet(
    edge_index=edge_index,
    node_dim=6,
    hidden_dim=64,
    n_atoms=3
)

# 3. Perform volume-preserving long-horizon rollout
h0 = torch.randn(1, 5, 3, 64)  # Initial trajectory (B, T, N, D)
rollout = model.forward_rollout(h0, steps=30)

# 4. Evaluate physical structural drift metrics
evaluator = PhysicsEval(koop_model=model, gru_model=model, test_split=None, n_atoms=3, molecule_name="custom")
bonds, angles, torsions = evaluator.extract_topology(edge_index)
```

---

### 3. Execution via Executable CLI (`kge`)

**Train models via `kge` CLI:**
```bash
# Train Graph Koopman on Aspirin (MD17)
kge train --md17 aspirin --model koopman --seed 42 --epochs 100

# Train E-GKN on DHA (MD22)
kge train --md22 dha --model e-gkn --seed 42 --epochs 100

# Train Graph Koopman on N-Body Charged
kge train --nbody charged --model koopman --seed 42 --epochs 100
```

**Evaluate Checkpoints via `kge` CLI:**
```bash
kge eval --md17 aspirin \
  --koopman-ckpt checkpoints/graph_aware_koopman_aspirin_seed42.pt \
  --gru-ckpt checkpoints/graph_aware_gru_aspirin_seed42.pt \
  --flat-ckpt checkpoints/flat_koopman_aspirin_seed42.pt \
  --rollout-steps 29
```

---

### 4. Quick Start via Docker (Recommended for GUI)

**Launch Interactive Web GUI Dashboard:**
```bash
docker compose up koopman-gui
```
*Access the dashboard at `http://localhost:8501` in your browser.*

**Train a model via Docker CLI:**
```bash
docker compose run --build --rm koopman train --md22 stachyose --model koopman --epochs 100
```

---

## Repository Structure

```text
.
├── koopman_evolver/           # Modular Python package
│   ├── data/                  # MD17, MD22, N-Body adapters and Kaggle downloaders
│   ├── models/                # GraphAwareKoopmanNet, GraphAwareGRUNet, FlatKoopmanNet, E-GKN
│   ├── training/              # PyTorch training loops with matrix exponential transitions
│   ├── evaluation/            # PhysicsEval suite & multi-system ablation metrics
│   └── cli.py                 # Command-line interface entrypoint
├── paper/                     # Manuscript source files, LaTeX tables, & figures
│   └── main.tex               # Conference manuscript LaTeX source
├── eval_logs/                 # Raw experimental log files & diagnostic evaluation outputs
├── app.py                     # Interactive Streamlit Web GUI Dashboard
├── requirements.txt           # Python package dependencies
├── Dockerfile                 # Container setup
└── docker-compose.yml         # Service definitions & volume mappings
```

---

## License

MIT License - feel free to use, fork, and modify for your research.
