Implement Diffusion Schrödinger Bridge for optimal transport

[Copilot]

Implements a minimal DSB algorithm for learning stochastic optimal transport between probability distributions using score-based diffusion models.

Core Components

• SDE solvers (core/sde.py): Forward/backward SDEs with Euler-Maruyama discretization. Forward diffuses from source, backward samples using learned scores.

• Score model (core/score_model.py): Time-conditioned MLP that estimates ∇log p_t(x) with sinusoidal embeddings.

• DSBridge (core/bridge.py): Training loop using Brownian bridge score matching. Iteratively refines drift term to match source and target marginals.

Usage

from dsb import DSBridge
from dsb.utils import gaussian_sampler

source = gaussian_sampler([0.0], 1.0)  # N(0, 1)
target = gaussian_sampler([3.0], 0.25) # N(3, 0.5²)

bridge = DSBridge(dim=1, T=1.0, num_steps=50)
bridge.train(source, target, n_iterations=500, batch_size=256)

# Transport samples
x_transported = bridge.sample(source(1000))

Examples

• simple_1d.py: 1D Gaussian-to-Gaussian with visualization
• 2d_moons.py: 2D Gaussian-to-TwoMoons with visualization

Both generate plots comparing source, transported, and target distributions.

Implementation Notes

• Uses simplified score matching with closed-form Brownian bridge targets for Gaussian approximation
• No iterative proportional fitting (IPF) - single-pass training for simplicity
• Constant diffusion coefficient; zero drift in forward SDE

Original prompt

a simple diffusion schrodinger bridge

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.