Skip to main content

2D Schrödinger Bridge Solver via Entropic Optimal Transport with Log-Domain Sinkhorn

Project description

anyakrakusuma: 2D Schrödinger Bridge Solver via Entropic Optimal Transport

DOI Python PyPI License: MIT Code style: black NumPy SciPy Matplotlib Pandas netCDF4 Numba Pillow tqdm

The nomenclature anyakrakusuma is directly derived from the regnal titulature of Sultan Agung of Mataram (r. 1613–1645), the third and most formidable ruler of the Mataram Sultanate in Java. Born as Raden Mas Jatmika (also known as Raden Mas Rangsang), he initially adopted the title Susuhunan Anyakrakusuma (or Prabu Pandita Anyakrakusuma) upon his accession in 1613, a period marked by his aggressive consolidation of power across Central and East Java. His reign is historically distinguished by his staunch military campaigns against the Dutch East India Company (VOC) in Batavia (1628–1629) and his cultural synthesis of Javanese and Islamic traditions, most notably the creation of the Javanese calendar. In 1641, following an investiture by the envoy of Sharif Zaid ibnu Muhsin Al Hasyimi of Mecca, he consecrated his authority with the supreme title Sultan Abdullah Muhammad Maulana Matarani al-Jawi, yet he remains immortalized in Javanese historiography by his syncretic honorific: Sultan Agung Adi Prabu Anyakrakusuma Senapati ing Ngalaga Abdurrahman Sayyidin Panatagama Khalifatullah Tanah Jawi, identifying him as the "Great Sultan," the supreme ruler, and the commander of the faithful in the land of Java.

Overview

anyakrakusuma is a high-performance Python solver for the Schrödinger Bridge Problem (SBP) via Entropic Optimal Transport (EOT). The Schrödinger Bridge computes the most likely stochastic evolution between two probability distributions, combining the geometric structure of optimal transport with the diffusive dynamics of Brownian motion.

This solver implements the log-domain Sinkhorn-Knopp algorithm for numerical stability, achieving robust convergence even with small regularization parameters. The implementation leverages Numba JIT compilation with parallel processing for significant performance gains (10-100× speedup) on large-scale particle systems.

Key Applications:

  • Generative modeling and distribution morphing
  • Stochastic process interpolation
  • Computational biology (cell trajectory inference)
  • Machine learning (diffusion models, score matching)
  • Statistical physics (entropy-regularized transport)

Note: This package provides 2D visualization of particle transport dynamics.

Physics

The Schrödinger Bridge Problem

The Schrödinger Bridge Problem seeks the most likely evolution of a diffusion process $X_t$ constrained to have marginals $\rho_0$ at $t=0$ and $\rho_1$ at $t=T$.

The solution satisfies the Schrödinger system (coupled via Feynman-Kac):

$$\partial_t \varphi(x,t) = \frac{\varepsilon}{2} \Delta \varphi(x,t) \quad \text{(forward heat equation)}$$

$$\partial_t \psi(x,t) = -\frac{\varepsilon}{2} \Delta \psi(x,t) \quad \text{(backward heat equation)}$$

with boundary conditions:

  • $\varphi(x,0)\psi(x,0) = \rho_0(x)$
  • $\varphi(x,T)\psi(x,T) = \rho_1(x)$

The marginal at time $t$ is: $\rho_t(x) = \varphi(x,t)\psi(x,t)$

Entropic Optimal Transport Formulation

For discrete measures $\mu = \sum_i a_i \delta_{x_i}$ and $\nu = \sum_j b_j \delta_{y_j}$, the entropic optimal transport problem is:

$$\min_{\pi \in \Pi(\mu,\nu)} \langle C, \pi \rangle + \varepsilon H(\pi | \mu \otimes \nu)$$

where:

  • $C_{ij} = |x_i - y_j|^2$ is the squared Euclidean cost
  • $H(\pi | \mu \otimes \nu)$ is the Kullback-Leibler divergence
  • $\varepsilon > 0$ is the entropic regularization (diffusivity)
  • $\Pi(\mu,\nu)$ is the set of couplings with marginals $\mu$ and $\nu$

The optimal coupling has the form $\pi^*{ij} = u_i K{ij} v_j$ where:

  • $K = \exp(-C/\varepsilon)$ is the Gibbs kernel (heat kernel for squared distance)
  • $u, v$ are the Sinkhorn potentials found by iterative scaling

Schrödinger Bridge Displacement Interpolation

The Schrödinger bridge displacement interpolation at time $t \in [0,1]$:

$$X_t = (1-t)X_0 + t X_1 + \sqrt{\varepsilon \cdot t(1-t)} \cdot Z, \quad Z \sim \mathcal{N}(0, I)$$

where $(X_0, X_1)$ are coupled via the optimal entropic plan $\pi^*$.

Physical Interpretation:

  • Deterministic drift: $(1-t)X_0 + tX_1$ — linear interpolation (McCann geodesic)
  • Stochastic diffusion: $\sqrt{\varepsilon \cdot t(1-t)} \cdot Z$ — Brownian bridge correction
  • The diffusion term vanishes at $t=0$ and $t=1$, ensuring exact boundary conditions

Key Theoretical Properties

  1. Entropy Production: The entropy of $\rho_t$ follows a parabolic profile
  2. Cost Decomposition: Transport cost = Wasserstein cost + Entropic regularization
  3. Marginal Constraints: $\sum_j \pi_{ij} = a_i$ and $\sum_i \pi_{ij} = b_j$
  4. Plan Positivity: $\pi_{ij} > 0$ for all $i,j$ (full support)

Numerical Methods

Log-Domain Sinkhorn Algorithm

The Sinkhorn iterations in log-domain for numerical stability:

$$f^{n+1} = \varepsilon \log(a) - \varepsilon \cdot \text{LSE}_j\left(\frac{g^n - C}{\varepsilon}\right)$$

$$g^{n+1} = \varepsilon \log(b) - \varepsilon \cdot \text{LSE}_i\left(\frac{f^{n+1} - C^\top}{\varepsilon}\right)$$

where $\text{LSE}$ is the log-sum-exp operator with max-trick for stability:

$$\text{LSE}(x) = \max(x) + \log\left(\sum_i \exp(x_i - \max(x))\right)$$

Convergence Criterion

Marginal constraint violation:

$$\text{error} = \left|\pi \mathbf{1} - a\right|_1$$

Novel Diagnostic Metrics

Metric Formula Physical Meaning
Wasserstein Cost $W = \langle C, \pi^* \rangle$ Total transport work
Plan Entropy $H = -\sum_{ij} \pi_{ij} \log(\pi_{ij})$ Coupling uncertainty
Effective Sparsity $S_{\text{eff}} = \exp(H) / nm$ Normalized plan support
Bridge Diffusivity $D_{\text{bridge}} = \varepsilon \cdot \mathbb{E}[t(1-t)]$ Average stochastic spread
Marginal Fidelity $1 - |\pi\mathbf{1} - a|_1$ Constraint satisfaction
Iteration Efficiency $\eta = -\log(\text{tol})/n_{\text{iter}}$ Convergence rate per iteration

Features

  • Log-domain stability: Robust convergence for small $\varepsilon$
  • Adaptive convergence: Automatic iteration based on marginal error
  • High-performance computing: Numba JIT compilation (10-100× speedup)
  • Parallel processing: Multi-core CPU utilization via prange
  • Professional 2D visualization: Animated particle transport
  • Comprehensive diagnostics: Static PNG with convergence curves
  • Scientific data formats: NetCDF4 (CF-1.8) and CSV output
  • Comparable test cases: Four scenarios with increasing complexity

Directory Structure

anyakrakusuma/
├── configs/                          # Simulation configuration files
│   ├── case1_circle_to_circle.txt   # Baseline: concentric circles
│   ├── case2_spiral_to_gaussian.txt # Topology change
│   ├── case3_moons_to_moons.txt     # Symmetric rotation
│   └── case4_lissajous_to_trefoil.txt # Complex curves
│
├── src/anyakrakusuma_eot/           # Main package source
│   ├── __init__.py                  # Package initialization
│   ├── cli.py                       # Command-line interface
│   │
│   ├── core/                        # Core numerical algorithms
│   │   ├── __init__.py
│   │   ├── solver.py                # Sinkhorn solver (log-domain)
│   │   └── distributions.py         # Point cloud generators
│   │
│   ├── io/                          # Input/output handlers
│   │   ├── __init__.py
│   │   ├── config_manager.py        # Configuration file parser
│   │   └── data_handler.py          # NetCDF/CSV writer
│   │
│   ├── visualization/               # Animation and plotting
│   │   ├── __init__.py
│   │   └── animator.py              # 2D GIF generator + diagnostics
│   │
│   └── utils/                       # Utilities
│       ├── __init__.py
│       ├── logger.py                # Simulation logging
│       └── timer.py                 # Performance profiling
│
├── tests/                           # Unit tests
│   ├── __init__.py
│   ├── test_solver.py               # Solver correctness tests
│   └── test_distributions.py        # Distribution generator tests
│
├── outputs/                         # Generated results (created at runtime)
│   ├── *.nc                         # NetCDF data files
│   ├── *.csv                        # CSV metric summaries
│   ├── *_diagnostics.png            # Static convergence plots
│   └── *.gif                        # Animation files
│
├── logs/                            # Simulation logs (created at runtime)
│   └── *.log
│
├── pyproject.toml                   # Poetry build configuration
├── README.md                        # This file
├── LICENSE                          # MIT license
└── .gitignore                       # Git ignore patterns

Installation

From PyPI:

pip install anyakrakusuma

From source:

git clone https://github.com/sandyherho/anyakrakusuma.git
cd anyakrakusuma
pip install -e .

Quick Start

Command line:

anyakrakusuma case1              # Circle to circle
anyakrakusuma case4 --cores 8    # Complex case with 8 cores
anyakrakusuma --all              # Run all cases

Python API:

from anyakrakusuma_eot import SchrodingerBridgeSolver, generate_circle, generate_spiral

# Initialize solver
solver = SchrodingerBridgeSolver(n_cores=8, verbose=True)

# Generate distributions
X_source = generate_circle(n=1000, radius=1.0)
X_target = generate_spiral(n=1000, turns=2.0)

# Solve Schrödinger Bridge
result = solver.solve(
    X_source, X_target,
    epsilon=0.05,
    max_iter=2000,
    tol=1e-9
)

# Check convergence
print(f"Transport cost: {result['cost']:.6f}")
print(f"Converged: {result['converged']}")

Test Cases (Comparable)

All test cases use the same number of particles (n=1000) and similar computational settings to enable direct comparison of transport characteristics.

Case Source → Target $\varepsilon$ Key Insight
1 Circle (r=1) → Circle (r=2) 0.02 Baseline radial scaling
2 Spiral → Gaussian mixture 0.05 Topology dissolution
3 Two moons → Two moons (rotated) 0.03 Symmetric interchange
4 Lissajous (3:2) → Trefoil knot 0.04 Complex curve morphing

Comparison Metrics (per case):

  • Wasserstein cost (transport effort)
  • Sinkhorn iterations (convergence difficulty)
  • Plan entropy (coupling complexity)
  • Effective sparsity (plan structure)

Citation

If you use this software in your study, please cite:

@software{herho2026anyakrakusuma,
  title   = {{anyakrakusuma: 2D Schrödinger Bridge Solver via Entropic Optimal Transport}},
  author  = {Herho, Sandy H. S. and Anwar, Iwan P. and Khadami, Faruq and Suwarman, Rusmawan and Irawan, Dasapta E.},
  year    = {2026},
  version = {0.0.1},
  url     = {https://github.com/sandyherho/anyakrakusuma}
}

Authors

License

MIT License - See LICENSE for details.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

anyakrakusuma-0.0.1.tar.gz (28.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

anyakrakusuma-0.0.1-py3-none-any.whl (30.2 kB view details)

Uploaded Python 3

File details

Details for the file anyakrakusuma-0.0.1.tar.gz.

File metadata

  • Download URL: anyakrakusuma-0.0.1.tar.gz
  • Upload date:
  • Size: 28.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/2.1.3 CPython/3.12.2 Linux/6.11.9-100.fc39.x86_64

File hashes

Hashes for anyakrakusuma-0.0.1.tar.gz
Algorithm Hash digest
SHA256 764a79379b9dae30b5fb9410492bc8bc8c92bc4648be1ee7353b2c82885b129b
MD5 ccec871f48e9efcd2efcd7c90c771811
BLAKE2b-256 190688c8009b50c51bc10171293761b78caf8c729d5db5c3724963979aa1db07

See more details on using hashes here.

File details

Details for the file anyakrakusuma-0.0.1-py3-none-any.whl.

File metadata

  • Download URL: anyakrakusuma-0.0.1-py3-none-any.whl
  • Upload date:
  • Size: 30.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/2.1.3 CPython/3.12.2 Linux/6.11.9-100.fc39.x86_64

File hashes

Hashes for anyakrakusuma-0.0.1-py3-none-any.whl
Algorithm Hash digest
SHA256 9c6fbf5dd1ba8ed3c074dcef0db9699812b354353fd43a00cd5e454b7695879d
MD5 4859c870558b1e6221282118e64091b2
BLAKE2b-256 0a26ff1495e2a2571869559736be8ba6f7d7d394df6a7055819512ed112577bf

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page