otloss 0.1.0__tar.gz

otloss-0.1.0/.github/workflows/tests.yml

@@ -0,0 +1,34 @@
+name: Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
+          pip install -e ".[dev]"
+
+      - name: Run tests
+        run: pytest tests/ -v --tb=short
+
+      - name: Check formatting
+        run: black --check otloss/ tests/

otloss-0.1.0/.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+.venv/
+*.egg

otloss-0.1.0/LICENSE

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

otloss-0.1.0/PKG-INFO

@@ -0,0 +1,332 @@
+Metadata-Version: 2.4
+Name: otloss
+Version: 0.1.0
+Summary: Optimal Transport training objectives for PyTorch — drop-in replacement for cross-entropy and MSE
+Project-URL: Homepage, https://github.com/Maqbool61/otloss
+Project-URL: Repository, https://github.com/Maqbool61/otloss
+Project-URL: Issues, https://github.com/Maqbool61/otloss/issues
+Author: Maqbool61
+License: MIT
+License-File: LICENSE
+Keywords: calibration,deep-learning,generative-models,loss-function,machine-learning,optimal-transport,pytorch,wasserstein
+Classifier: Development Status :: 3 - Alpha
+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: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: torch>=2.0.0
+Provides-Extra: dev
+Requires-Dist: black; extra == 'dev'
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: examples
+Requires-Dist: matplotlib>=3.5; extra == 'examples'
+Requires-Dist: numpy>=1.22; extra == 'examples'
+Description-Content-Type: text/markdown
+
+# otloss
+
+**Optimal Transport training objectives for PyTorch** — a drop-in replacement for cross-entropy and MSE that eliminates mode collapse, improves calibration, and produces robust distributional representations.
+
+[![PyPI version](https://img.shields.io/pypi/v/otloss.svg)](https://pypi.org/project/otloss/)
+[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+[![Tests](https://github.com/Maqbool61/otloss/actions/workflows/tests.yml/badge.svg)](https://github.com/Maqbool61/otloss/actions)
+
+---
+
+## Why WassersteinLoss?
+
+Cross-entropy and KL divergence have a fundamental flaw: **their gradients vanish when the model and target distributions don't overlap**. This causes:
+
+- Mode collapse in GANs and generative models
+- Overconfident, poorly calibrated probability outputs
+- Brittle behaviour under distribution shift
+
+The Wasserstein-2 distance solves all three. It defines a *geometric* distance between distributions using the ground metric of the feature space:
+
+```
+W₂(μ, ν) = inf_{γ ∈ Π(μ,ν)} ∫ ‖x - y‖² dγ(x, y)
+```
+
+It always has **meaningful gradients**, naturally respects the geometry of your output space, and produces smooth, interference-free learning signals.
+
+---
+
+## Installation
+
+```bash
+pip install otloss
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/Maqbool61/otloss.git
+cd otloss
+pip install -e ".[dev]"
+```
+
+**Requirements:** Python ≥ 3.9, PyTorch ≥ 2.0
+
+---
+
+## Quick start
+
+```python
+from otloss import WassersteinLoss
+
+# Drop-in replacement for nn.MSELoss
+criterion = WassersteinLoss(p=2, blur=0.05)
+
+pred   = torch.randn(32, 100, 2, requires_grad=True)  # (batch, particles, dim)
+target = torch.randn(32, 100, 2)
+
+loss = criterion(pred, target)
+loss.backward()   # gradients flow through pred
+```
+
+---
+
+## Layered API
+
+### High-level (nn.Module)
+
+```python
+from otloss import WassersteinLoss, SlicedWassersteinLoss
+from otloss.losses import WassersteinGANLoss
+
+# Exact Wasserstein via Sinkhorn (best accuracy)
+criterion = WassersteinLoss(
+    p=2,           # Wasserstein order (1 or 2)
+    blur=0.05,     # entropic regularisation ε = blur²
+    max_iter=100,  # Sinkhorn iterations
+    debias=True,   # Sinkhorn divergence debiasing
+    reduction="mean",
+)
+
+# Fast O(n log n) approximation via random projections
+criterion = SlicedWassersteinLoss(
+    n_projections=200,
+    p=2,
+)
+
+# WGAN-GP critic/generator losses
+criterion = WassersteinGANLoss(gp_weight=10.0)
+d_loss = criterion.critic_loss(real_scores, fake_scores)
+gp     = criterion.gradient_penalty(critic, real, fake)
+g_loss = criterion.generator_loss(fake_scores)
+```
+
+### Functional API (low-level, full control)
+
+```python
+from otloss import (
+    otloss,      # full Wasserstein via Sinkhorn
+    sliced_otloss,
+    sinkhorn,              # raw Sinkhorn solver
+    cost_matrix,           # ground cost C_{ij} = ‖xᵢ - yⱼ‖ᵖ
+    dual_variables,        # Kantorovich dual potentials (f, g)
+)
+
+# Compute cost matrix
+C = cost_matrix(x, y, p=2)          # (N, M)
+
+# Run Sinkhorn and get dual potentials + transport cost
+f, g, cost = sinkhorn(a, b, C, blur=0.05, debias=True)
+
+# Recover soft transport plan P_{ij}
+from otloss.utils import transport_plan
+P = transport_plan(f, g, C, blur=0.05)  # (N, M)
+
+# Wasserstein barycenter
+from otloss.utils import wasserstein_barycenter_weights
+barycenter = wasserstein_barycenter_weights(measures, weights=[0.3, 0.7], support=X)
+```
+
+---
+
+## Real-world use cases
+
+### 1. GAN training without mode collapse
+
+```python
+from otloss.losses import WassersteinGANLoss
+
+criterion = WassersteinGANLoss(gp_weight=10.0)
+
+# Critic update
+c_loss = criterion.critic_loss(D(real), D(fake.detach()))
+gp     = criterion.gradient_penalty(D, real, fake)
+(c_loss + gp).backward()
+
+# Generator update
+g_loss = criterion.generator_loss(D(fake))
+g_loss.backward()
+```
+
+Or even simpler — no critic needed:
+
+```python
+criterion = WassersteinLoss(blur=0.05)
+
+# Directly minimise W₂ between fake and real sample clouds
+fake = G(noise)   # (B, N, D)
+real = real_data  # (B, N, D)
+loss = criterion(fake, real)
+loss.backward()
+```
+
+### 2. LLM calibration (confidence matches accuracy)
+
+```python
+from otloss import WassersteinLoss
+from otloss.distributions import label_smoothed_weights
+
+criterion = WassersteinLoss(p=2, blur=0.05)
+
+# Class positions as 1-D support
+support = torch.linspace(0, 1, n_classes).unsqueeze(-1)   # (K, 1)
+support = support.unsqueeze(0).expand(B, -1, -1)           # (B, K, 1)
+
+pred_weights = torch.softmax(logits, dim=-1)               # (B, K)
+target_weights = label_smoothed_weights(y, n_classes)      # (B, K)
+
+loss = criterion(support, support,
+                 pred_weights=pred_weights,
+                 target_weights=target_weights)
+```
+
+### 3. Drug molecule generation (diverse scaffolds)
+
+```python
+from otloss import SlicedWassersteinLoss
+
+# Measures structural diversity in 8-D property space
+criterion = SlicedWassersteinLoss(n_projections=200)
+
+generated = model.decode(z)          # (B, N, 8)  property vectors
+reference  = real_molecules          # (B, N, 8)
+loss = criterion(generated, reference)
+loss.backward()
+# → model covers the full pharmacological distribution
+```
+
+### 4. Financial time-series generation (fat tails)
+
+```python
+criterion = WassersteinLoss(p=2, blur=0.01)   # small blur → sharp tails
+
+generated_returns = model(noise)   # (B, T, 1)
+real_returns      = historical     # (B, T, 1)
+loss = criterion(generated_returns, real_returns)
+# → VaR/CVaR of generated paths matches real market data
+```
+
+### 5. RLHF reward model training
+
+```python
+criterion = WassersteinLoss(p=2, blur=0.05, debias=True)
+
+# Reward model outputs as distributions over preference scores
+pred_rewards = reward_model(responses)    # (B, K, 1)
+human_prefs  = preference_labels         # (B, K, 1)
+loss = criterion(pred_rewards, human_prefs)
+# → smoother reward landscape → better RLHF alignment
+```
+
+---
+
+## Mathematical background
+
+### Entropic regularisation (Sinkhorn)
+
+Direct computation of W₂ is O(n³). We solve the entropy-regularised problem:
+
+```
+W_ε(a, b) = min_{P ≥ 0} ⟨C, P⟩ − ε · H(P)
+             s.t.  P·1 = a,  Pᵀ·1 = b
+```
+
+Via Sinkhorn-Knopp iterations in log-domain (numerically stable):
+
+```
+fᵢ ← ε · log(aᵢ) − ε · LSE_j[(gⱼ − Cᵢⱼ) / ε]
+gⱼ ← ε · log(bⱼ) − ε · LSE_i[(fᵢ − Cᵢⱼ) / ε]
+```
+
+### Sinkhorn divergence (debiasing)
+
+Raw Sinkhorn overestimates W due to entropic bias. We correct with:
+
+```
+S_ε(a, b) = W_ε(a, b) − ½W_ε(a, a) − ½W_ε(b, b)
+```
+
+This ensures `S_ε(a, a) = 0` (positive definite) and `S_ε → W` as `ε → 0`.
+
+### Sliced Wasserstein Distance
+
+Projects to 1-D random lines and uses the closed-form 1-D solution:
+
+```
+SW_p(μ, ν) = ( ∫_{S^{D-1}} W_p(θ#μ, θ#ν)^p dσ(θ) )^{1/p}
+```
+
+Exact W in 1-D reduces to: `W_p = ‖sort(x) − sort(y)‖_p / N^{1/p}`.
+Complexity: **O(n log n)** vs O(n³) for exact OT.
+
+---
+
+## Choosing `blur`
+
+| Scenario | Recommended blur |
+|---|---|
+| Tight distributions (calibration) | 0.01 – 0.03 |
+| Moderate spread (generation) | 0.05 – 0.1 |
+| Very spread / high-dimensional | 0.1 – 0.5 |
+| Rule of thumb | `blur ≈ std(data) × 0.05` |
+
+Smaller blur = more accurate but more Sinkhorn iterations. Blur annealing (enabled by default via `scaling=0.5`) starts coarse and refines automatically.
+
+---
+
+## Running tests
+
+```bash
+pytest tests/ -v
+```
+
+---
+
+## Citation
+
+If you use WassersteinLoss in your research:
+
+```bibtex
+@software{otloss_2026,
+  author  = {Maqbool61},
+  title   = {otloss: Optimal Transport objectives for PyTorch},
+  year    = {2026},
+  url     = {https://github.com/Maqbool61/otloss},
+}
+```
+
+**Key papers:**
+- Villani (2008) — *Optimal Transport: Old and New*
+- Cuturi (2013) — *Sinkhorn Distances: Lightspeed Computation of Optimal Transport*
+- Arjovsky et al. (2017) — *Wasserstein GAN*
+- Gulrajani et al. (2017) — *Improved Training of Wasserstein GANs*
+- Feydy et al. (2019) — *Interpolating between Optimal Transport and MMD using Sinkhorn Divergences*
+
+---
+
+## License
+
+MIT © Maqbool61

otloss-0.1.0/README.md

@@ -0,0 +1,300 @@
+# otloss
+
+**Optimal Transport training objectives for PyTorch** — a drop-in replacement for cross-entropy and MSE that eliminates mode collapse, improves calibration, and produces robust distributional representations.
+
+[![PyPI version](https://img.shields.io/pypi/v/otloss.svg)](https://pypi.org/project/otloss/)
+[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+[![Tests](https://github.com/Maqbool61/otloss/actions/workflows/tests.yml/badge.svg)](https://github.com/Maqbool61/otloss/actions)
+
+---
+
+## Why WassersteinLoss?
+
+Cross-entropy and KL divergence have a fundamental flaw: **their gradients vanish when the model and target distributions don't overlap**. This causes:
+
+- Mode collapse in GANs and generative models
+- Overconfident, poorly calibrated probability outputs
+- Brittle behaviour under distribution shift
+
+The Wasserstein-2 distance solves all three. It defines a *geometric* distance between distributions using the ground metric of the feature space:
+
+```
+W₂(μ, ν) = inf_{γ ∈ Π(μ,ν)} ∫ ‖x - y‖² dγ(x, y)
+```
+
+It always has **meaningful gradients**, naturally respects the geometry of your output space, and produces smooth, interference-free learning signals.
+
+---
+
+## Installation
+
+```bash
+pip install otloss
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/Maqbool61/otloss.git
+cd otloss
+pip install -e ".[dev]"
+```
+
+**Requirements:** Python ≥ 3.9, PyTorch ≥ 2.0
+
+---
+
+## Quick start
+
+```python
+from otloss import WassersteinLoss
+
+# Drop-in replacement for nn.MSELoss
+criterion = WassersteinLoss(p=2, blur=0.05)
+
+pred   = torch.randn(32, 100, 2, requires_grad=True)  # (batch, particles, dim)
+target = torch.randn(32, 100, 2)
+
+loss = criterion(pred, target)
+loss.backward()   # gradients flow through pred
+```
+
+---
+
+## Layered API
+
+### High-level (nn.Module)
+
+```python
+from otloss import WassersteinLoss, SlicedWassersteinLoss
+from otloss.losses import WassersteinGANLoss
+
+# Exact Wasserstein via Sinkhorn (best accuracy)
+criterion = WassersteinLoss(
+    p=2,           # Wasserstein order (1 or 2)
+    blur=0.05,     # entropic regularisation ε = blur²
+    max_iter=100,  # Sinkhorn iterations
+    debias=True,   # Sinkhorn divergence debiasing
+    reduction="mean",
+)
+
+# Fast O(n log n) approximation via random projections
+criterion = SlicedWassersteinLoss(
+    n_projections=200,
+    p=2,
+)
+
+# WGAN-GP critic/generator losses
+criterion = WassersteinGANLoss(gp_weight=10.0)
+d_loss = criterion.critic_loss(real_scores, fake_scores)
+gp     = criterion.gradient_penalty(critic, real, fake)
+g_loss = criterion.generator_loss(fake_scores)
+```
+
+### Functional API (low-level, full control)
+
+```python
+from otloss import (
+    otloss,      # full Wasserstein via Sinkhorn
+    sliced_otloss,
+    sinkhorn,              # raw Sinkhorn solver
+    cost_matrix,           # ground cost C_{ij} = ‖xᵢ - yⱼ‖ᵖ
+    dual_variables,        # Kantorovich dual potentials (f, g)
+)
+
+# Compute cost matrix
+C = cost_matrix(x, y, p=2)          # (N, M)
+
+# Run Sinkhorn and get dual potentials + transport cost
+f, g, cost = sinkhorn(a, b, C, blur=0.05, debias=True)
+
+# Recover soft transport plan P_{ij}
+from otloss.utils import transport_plan
+P = transport_plan(f, g, C, blur=0.05)  # (N, M)
+
+# Wasserstein barycenter
+from otloss.utils import wasserstein_barycenter_weights
+barycenter = wasserstein_barycenter_weights(measures, weights=[0.3, 0.7], support=X)
+```
+
+---
+
+## Real-world use cases
+
+### 1. GAN training without mode collapse
+
+```python
+from otloss.losses import WassersteinGANLoss
+
+criterion = WassersteinGANLoss(gp_weight=10.0)
+
+# Critic update
+c_loss = criterion.critic_loss(D(real), D(fake.detach()))
+gp     = criterion.gradient_penalty(D, real, fake)
+(c_loss + gp).backward()
+
+# Generator update
+g_loss = criterion.generator_loss(D(fake))
+g_loss.backward()
+```
+
+Or even simpler — no critic needed:
+
+```python
+criterion = WassersteinLoss(blur=0.05)
+
+# Directly minimise W₂ between fake and real sample clouds
+fake = G(noise)   # (B, N, D)
+real = real_data  # (B, N, D)
+loss = criterion(fake, real)
+loss.backward()
+```
+
+### 2. LLM calibration (confidence matches accuracy)
+
+```python
+from otloss import WassersteinLoss
+from otloss.distributions import label_smoothed_weights
+
+criterion = WassersteinLoss(p=2, blur=0.05)
+
+# Class positions as 1-D support
+support = torch.linspace(0, 1, n_classes).unsqueeze(-1)   # (K, 1)
+support = support.unsqueeze(0).expand(B, -1, -1)           # (B, K, 1)
+
+pred_weights = torch.softmax(logits, dim=-1)               # (B, K)
+target_weights = label_smoothed_weights(y, n_classes)      # (B, K)
+
+loss = criterion(support, support,
+                 pred_weights=pred_weights,
+                 target_weights=target_weights)
+```
+
+### 3. Drug molecule generation (diverse scaffolds)
+
+```python
+from otloss import SlicedWassersteinLoss
+
+# Measures structural diversity in 8-D property space
+criterion = SlicedWassersteinLoss(n_projections=200)
+
+generated = model.decode(z)          # (B, N, 8)  property vectors
+reference  = real_molecules          # (B, N, 8)
+loss = criterion(generated, reference)
+loss.backward()
+# → model covers the full pharmacological distribution
+```
+
+### 4. Financial time-series generation (fat tails)
+
+```python
+criterion = WassersteinLoss(p=2, blur=0.01)   # small blur → sharp tails
+
+generated_returns = model(noise)   # (B, T, 1)
+real_returns      = historical     # (B, T, 1)
+loss = criterion(generated_returns, real_returns)
+# → VaR/CVaR of generated paths matches real market data
+```
+
+### 5. RLHF reward model training
+
+```python
+criterion = WassersteinLoss(p=2, blur=0.05, debias=True)
+
+# Reward model outputs as distributions over preference scores
+pred_rewards = reward_model(responses)    # (B, K, 1)
+human_prefs  = preference_labels         # (B, K, 1)
+loss = criterion(pred_rewards, human_prefs)
+# → smoother reward landscape → better RLHF alignment
+```
+
+---
+
+## Mathematical background
+
+### Entropic regularisation (Sinkhorn)
+
+Direct computation of W₂ is O(n³). We solve the entropy-regularised problem:
+
+```
+W_ε(a, b) = min_{P ≥ 0} ⟨C, P⟩ − ε · H(P)
+             s.t.  P·1 = a,  Pᵀ·1 = b
+```
+
+Via Sinkhorn-Knopp iterations in log-domain (numerically stable):
+
+```
+fᵢ ← ε · log(aᵢ) − ε · LSE_j[(gⱼ − Cᵢⱼ) / ε]
+gⱼ ← ε · log(bⱼ) − ε · LSE_i[(fᵢ − Cᵢⱼ) / ε]
+```
+
+### Sinkhorn divergence (debiasing)
+
+Raw Sinkhorn overestimates W due to entropic bias. We correct with:
+
+```
+S_ε(a, b) = W_ε(a, b) − ½W_ε(a, a) − ½W_ε(b, b)
+```
+
+This ensures `S_ε(a, a) = 0` (positive definite) and `S_ε → W` as `ε → 0`.
+
+### Sliced Wasserstein Distance
+
+Projects to 1-D random lines and uses the closed-form 1-D solution:
+
+```
+SW_p(μ, ν) = ( ∫_{S^{D-1}} W_p(θ#μ, θ#ν)^p dσ(θ) )^{1/p}
+```
+
+Exact W in 1-D reduces to: `W_p = ‖sort(x) − sort(y)‖_p / N^{1/p}`.
+Complexity: **O(n log n)** vs O(n³) for exact OT.
+
+---
+
+## Choosing `blur`
+
+| Scenario | Recommended blur |
+|---|---|
+| Tight distributions (calibration) | 0.01 – 0.03 |
+| Moderate spread (generation) | 0.05 – 0.1 |
+| Very spread / high-dimensional | 0.1 – 0.5 |
+| Rule of thumb | `blur ≈ std(data) × 0.05` |
+
+Smaller blur = more accurate but more Sinkhorn iterations. Blur annealing (enabled by default via `scaling=0.5`) starts coarse and refines automatically.
+
+---
+
+## Running tests
+
+```bash
+pytest tests/ -v
+```
+
+---
+
+## Citation
+
+If you use WassersteinLoss in your research:
+
+```bibtex
+@software{otloss_2026,
+  author  = {Maqbool61},
+  title   = {otloss: Optimal Transport objectives for PyTorch},
+  year    = {2026},
+  url     = {https://github.com/Maqbool61/otloss},
+}
+```
+
+**Key papers:**
+- Villani (2008) — *Optimal Transport: Old and New*
+- Cuturi (2013) — *Sinkhorn Distances: Lightspeed Computation of Optimal Transport*
+- Arjovsky et al. (2017) — *Wasserstein GAN*
+- Gulrajani et al. (2017) — *Improved Training of Wasserstein GANs*
+- Feydy et al. (2019) — *Interpolating between Optimal Transport and MMD using Sinkhorn Divergences*
+
+---
+
+## License
+
+MIT © Maqbool61