mdot-tnt 0.1.0__tar.gz → 1.0.0__tar.gz

{mdot_tnt-0.1.0 → mdot_tnt-1.0.0}/LICENSE

@@ -23,4 +23,7 @@ For commercial use, a separate license must be obtained from the Licensor. To in
 This license automatically terminates if the Licensee breaches any of its terms. Upon termination, all rights granted under this license are revoked, and the Licensee must cease using and distributing the Software.
 
 ## 4. Governing Law and Enforcement
-This license shall be governed by and construed in accordance with the laws of Ontario, Canada. However, violations of this license may also be pursued un
+This license shall be governed by and construed in accordance with the laws of Ontario, Canada. However, violations of this license may also be pursued under applicable copyright laws in the jurisdiction where infringement occurs.
+
+## 5. Contact
+For licensing inquiries, please contact: **kemertas@cs.toronto.edu**.

mdot_tnt-1.0.0/PKG-INFO

@@ -0,0 +1,216 @@
+Metadata-Version: 2.4
+Name: mdot-tnt
+Version: 1.0.0
+Summary: A fast, GPU-parallel, PyTorch-compatible optimal transport solver.
+Author-email: Mete Kemertas <kemertas@cs.toronto.edu>
+License: # Non-Commercial Research License (NCRL-1.0)
+        
+        Copyright (C) 2025 Mete Kemertas
+        
+        ## 1. License Grant
+        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to use, copy, modify, merge, publish, and distribute copies of the Software **solely for non-commercial research, educational, and personal purposes**, subject to the following conditions:
+        
+        ## 2. Restrictions
+        ### 2.1 **Non-Commercial Use Only**
+        - The Software **may NOT** be used for any commercial purpose without explicit written permission from the Licensor.
+        - "Commercial purpose" includes, but is not limited to:
+          - Selling or licensing the Software.
+          - Using the Software in proprietary products or services.
+          - Offering the Software as part of a paid or revenue-generating service.
+        
+        ### 2.2 **No Warranty & Liability**
+        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 ARISING FROM THE USE OF THE SOFTWARE.
+        
+        ### 2.3 **Commercial Licensing**
+        For commercial use, a separate license must be obtained from the Licensor. To inquire about licensing, please contact: **kemertas@cs.toronto.edu**.
+        
+        ## 3. Termination
+        This license automatically terminates if the Licensee breaches any of its terms. Upon termination, all rights granted under this license are revoked, and the Licensee must cease using and distributing the Software.
+        
+        ## 4. Governing Law and Enforcement
+        This license shall be governed by and construed in accordance with the laws of Ontario, Canada. However, violations of this license may also be pursued under applicable copyright laws in the jurisdiction where infringement occurs.
+        
+        ## 5. Contact
+        For licensing inquiries, please contact: **kemertas@cs.toronto.edu**.
+        
+Project-URL: Homepage, https://github.com/metekemertas/mdot_tnt
+Project-URL: Documentation, https://mdot-tnt.readthedocs.io
+Project-URL: Repository, https://github.com/metekemertas/mdot_tnt
+Project-URL: Issues, https://github.com/metekemertas/mdot_tnt/issues
+Keywords: optimal-transport,sinkhorn,entropy-regularization,machine-learning,pytorch,gpu
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+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: Operating System :: OS Independent
+Classifier: Typing :: Typed
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: pre-commit>=3.0; extra == "dev"
+Requires-Dist: numpy>=1.20; extra == "dev"
+Dynamic: license-file
+
+# MDOT-TNT
+
+<img src="assets/logo.png" alt="MDOT-TNT Logo" width="180" align="right"/>
+
+**A Truncated Newton Method for Optimal Transport**
+
+[![PyPI version](https://badge.fury.io/py/mdot-tnt.svg)](https://badge.fury.io/py/mdot-tnt)
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
+[![License](https://img.shields.io/badge/license-Non--Commercial-green.svg)](LICENSE)
+
+A fast, GPU-accelerated solver for entropic-regularized optimal transport (OT) problems. MDOT-TNT combines mirror descent with a truncated Newton projection method to achieve high numerical precision while remaining stable under weak regularization.
+
+<br clear="right"/>
+
+## Features
+
+- **High Precision**: Stable under extremely weak regularization  (γ up to 2¹⁸), enabling highly precise approximations of unregularized OT
+- **GPU Accelerated**: Fully compatible with CUDA for fast computation on large problems
+- **Batched Solving**: Solve multiple OT problems simultaneously in batched mode
+- **Memory Efficient**: Log-domain computations and efficient rounding avoid storing full transport plans
+- **PyTorch Native**: Seamless integration with PyTorch, supporting autograd-compatible inputs
+
+## Installation
+
+**Prerequisites**: Install [PyTorch](https://pytorch.org/get-started/locally/) for your system configuration first.
+
+```bash
+pip install mdot-tnt
+```
+
+For development:
+
+```bash
+git clone https://github.com/metekemertas/mdot_tnt.git
+cd mdot_tnt
+pip install -e ".[dev]"
+```
+
+## Quick Start
+
+### Single Problem
+
+```python
+import torch
+import mdot_tnt
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Create marginals (probability distributions)
+n, m = 512, 512
+r = torch.rand(n, device=device, dtype=torch.float64)
+r = r / r.sum()
+c = torch.rand(m, device=device, dtype=torch.float64)
+c = c / c.sum()
+
+# Cost matrix (e.g., pairwise distances)
+C = torch.rand(n, m, device=device, dtype=torch.float64)
+
+# Solve for optimal transport cost
+cost = mdot_tnt.solve_OT(r, c, C, gamma_f=1024)
+
+# Or get the full transport plan
+plan = mdot_tnt.solve_OT(r, c, C, gamma_f=1024, return_plan=True)
+```
+
+### Batched Solving
+
+When solving multiple OT problems, use the batched solver for significant speedup compared to sequential solution:
+
+```python
+import torch
+import mdot_tnt
+
+device = "cuda"
+batch_size, n, m = 32, 512, 512
+
+# Multiple marginal pairs
+r = torch.rand(batch_size, n, device=device, dtype=torch.float64)
+r = r / r.sum(-1, keepdim=True)
+c = torch.rand(batch_size, m, device=device, dtype=torch.float64)
+c = c / c.sum(-1, keepdim=True)
+
+# Shared cost matrix (or per-problem: shape [batch_size, n, m])
+C = torch.rand(n, m, device=device, dtype=torch.float64)
+
+# Solve all problems at once
+costs = mdot_tnt.solve_OT_batched(r, c, C, gamma_f=1024)  # Returns (batch_size,) tensor
+```
+
+The batched solver achieves speedup by amortizing GPU synchronization overhead across all problems in the batch.
+
+## API Reference
+
+### `solve_OT`
+
+```python
+mdot_tnt.solve_OT(r, c, C, gamma_f=1024., return_plan=False, round=True, log=False)
+```
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `r` | `Tensor` | Row marginal of shape `(n,)`, must sum to 1 |
+| `c` | `Tensor` | Column marginal of shape `(m,)`, must sum to 1 |
+| `C` | `Tensor` | Cost matrix of shape `(n, m)`, recommended to normalize to [0, 1] |
+| `gamma_f` | `float` | Temperature parameter (inverse regularization). Higher = more accurate. Default: 1024 |
+| `return_plan` | `bool` | If True, return transport plan instead of cost |
+| `round` | `bool` | If True, round solution onto feasible set |
+| `log` | `bool` | If True, also return optimization logs |
+
+**Returns**: Transport cost (scalar) or plan `(n, m)`, optionally with logs dict.
+
+### `solve_OT_batched`
+
+```python
+mdot_tnt.solve_OT_batched(r, c, C, gamma_f=1024., return_plan=False, round=True, log=False)
+```
+
+Same parameters as `solve_OT`, but with batched inputs:
+- `r`: Shape `(batch, n)`
+- `c`: Shape `(batch, m)`  
+- `C`: Shape `(n, m)` for shared cost, or `(batch, n, m)` for per-problem costs
+
+**Returns**: Costs `(batch,)` or plans `(batch, n, m)`.
+
+## Performance Tips
+
+1. **Use float64** for `gamma_f > 1024` (automatic conversion with warning)
+2. **Normalize cost matrices** to [0, 1] for numerical stability
+3. **Use batched solver** when solving multiple problems with shared structure
+4. **Increase `gamma_f`** for higher precision (error scales as O(log n / γ) in the worst case, but can be much better)
+
+## Citation
+
+If you use MDOT-TNT in your research, please cite:
+
+```bibtex
+@inproceedings{kemertas2025truncated,
+  title={A Truncated Newton Method for Optimal Transport},
+  author={Kemertas, Mete and Farahmand, Amir-massoud and Jepson, Allan Douglas},
+  booktitle={The Thirteenth International Conference on Learning Representations},
+  year={2025},
+  url={https://openreview.net/forum?id=gWrWUaCbMa}
+}
+```
+
+## License
+
+This code is released under a [non-commercial use license](LICENSE). For commercial licensing inquiries, please contact the authors.
+
+## Contact
+
+For questions or issues, please [open an issue](https://github.com/metekemertas/mdot_tnt/issues) or email: kemertas [at] cs [dot] toronto [dot] edu

mdot_tnt-1.0.0/README.md

@@ -0,0 +1,152 @@
+# MDOT-TNT
+
+<img src="assets/logo.png" alt="MDOT-TNT Logo" width="180" align="right"/>
+
+**A Truncated Newton Method for Optimal Transport**
+
+[![PyPI version](https://badge.fury.io/py/mdot-tnt.svg)](https://badge.fury.io/py/mdot-tnt)
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
+[![License](https://img.shields.io/badge/license-Non--Commercial-green.svg)](LICENSE)
+
+A fast, GPU-accelerated solver for entropic-regularized optimal transport (OT) problems. MDOT-TNT combines mirror descent with a truncated Newton projection method to achieve high numerical precision while remaining stable under weak regularization.
+
+<br clear="right"/>
+
+## Features
+
+- **High Precision**: Stable under extremely weak regularization  (γ up to 2¹⁸), enabling highly precise approximations of unregularized OT
+- **GPU Accelerated**: Fully compatible with CUDA for fast computation on large problems
+- **Batched Solving**: Solve multiple OT problems simultaneously in batched mode
+- **Memory Efficient**: Log-domain computations and efficient rounding avoid storing full transport plans
+- **PyTorch Native**: Seamless integration with PyTorch, supporting autograd-compatible inputs
+
+## Installation
+
+**Prerequisites**: Install [PyTorch](https://pytorch.org/get-started/locally/) for your system configuration first.
+
+```bash
+pip install mdot-tnt
+```
+
+For development:
+
+```bash
+git clone https://github.com/metekemertas/mdot_tnt.git
+cd mdot_tnt
+pip install -e ".[dev]"
+```
+
+## Quick Start
+
+### Single Problem
+
+```python
+import torch
+import mdot_tnt
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Create marginals (probability distributions)
+n, m = 512, 512
+r = torch.rand(n, device=device, dtype=torch.float64)
+r = r / r.sum()
+c = torch.rand(m, device=device, dtype=torch.float64)
+c = c / c.sum()
+
+# Cost matrix (e.g., pairwise distances)
+C = torch.rand(n, m, device=device, dtype=torch.float64)
+
+# Solve for optimal transport cost
+cost = mdot_tnt.solve_OT(r, c, C, gamma_f=1024)
+
+# Or get the full transport plan
+plan = mdot_tnt.solve_OT(r, c, C, gamma_f=1024, return_plan=True)
+```
+
+### Batched Solving
+
+When solving multiple OT problems, use the batched solver for significant speedup compared to sequential solution:
+
+```python
+import torch
+import mdot_tnt
+
+device = "cuda"
+batch_size, n, m = 32, 512, 512
+
+# Multiple marginal pairs
+r = torch.rand(batch_size, n, device=device, dtype=torch.float64)
+r = r / r.sum(-1, keepdim=True)
+c = torch.rand(batch_size, m, device=device, dtype=torch.float64)
+c = c / c.sum(-1, keepdim=True)
+
+# Shared cost matrix (or per-problem: shape [batch_size, n, m])
+C = torch.rand(n, m, device=device, dtype=torch.float64)
+
+# Solve all problems at once
+costs = mdot_tnt.solve_OT_batched(r, c, C, gamma_f=1024)  # Returns (batch_size,) tensor
+```
+
+The batched solver achieves speedup by amortizing GPU synchronization overhead across all problems in the batch.
+
+## API Reference
+
+### `solve_OT`
+
+```python
+mdot_tnt.solve_OT(r, c, C, gamma_f=1024., return_plan=False, round=True, log=False)
+```
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `r` | `Tensor` | Row marginal of shape `(n,)`, must sum to 1 |
+| `c` | `Tensor` | Column marginal of shape `(m,)`, must sum to 1 |
+| `C` | `Tensor` | Cost matrix of shape `(n, m)`, recommended to normalize to [0, 1] |
+| `gamma_f` | `float` | Temperature parameter (inverse regularization). Higher = more accurate. Default: 1024 |
+| `return_plan` | `bool` | If True, return transport plan instead of cost |
+| `round` | `bool` | If True, round solution onto feasible set |
+| `log` | `bool` | If True, also return optimization logs |
+
+**Returns**: Transport cost (scalar) or plan `(n, m)`, optionally with logs dict.
+
+### `solve_OT_batched`
+
+```python
+mdot_tnt.solve_OT_batched(r, c, C, gamma_f=1024., return_plan=False, round=True, log=False)
+```
+
+Same parameters as `solve_OT`, but with batched inputs:
+- `r`: Shape `(batch, n)`
+- `c`: Shape `(batch, m)`  
+- `C`: Shape `(n, m)` for shared cost, or `(batch, n, m)` for per-problem costs
+
+**Returns**: Costs `(batch,)` or plans `(batch, n, m)`.
+
+## Performance Tips
+
+1. **Use float64** for `gamma_f > 1024` (automatic conversion with warning)
+2. **Normalize cost matrices** to [0, 1] for numerical stability
+3. **Use batched solver** when solving multiple problems with shared structure
+4. **Increase `gamma_f`** for higher precision (error scales as O(log n / γ) in the worst case, but can be much better)
+
+## Citation
+
+If you use MDOT-TNT in your research, please cite:
+
+```bibtex
+@inproceedings{kemertas2025truncated,
+  title={A Truncated Newton Method for Optimal Transport},
+  author={Kemertas, Mete and Farahmand, Amir-massoud and Jepson, Allan Douglas},
+  booktitle={The Thirteenth International Conference on Learning Representations},
+  year={2025},
+  url={https://openreview.net/forum?id=gWrWUaCbMa}
+}
+```
+
+## License
+
+This code is released under a [non-commercial use license](LICENSE). For commercial licensing inquiries, please contact the authors.
+
+## Contact
+
+For questions or issues, please [open an issue](https://github.com/metekemertas/mdot_tnt/issues) or email: kemertas [at] cs [dot] toronto [dot] edu

{mdot_tnt-0.1.0 → mdot_tnt-1.0.0}/mdot_tnt/__init__.py

@@ -1,11 +1,48 @@
+"""
+MDOT-TNT: A Truncated Newton Method for Optimal Transport
+
+This package provides efficient solvers for the entropic-regularized optimal transport
+problem, as introduced in the paper "A Truncated Newton Method for Optimal Transport"
+by Mete Kemertas, Amir-massoud Farahmand, Allan D. Jepson (ICLR, 2025).
+URL: https://openreview.net/forum?id=gWrWUaCbMa
+
+Main functions:
+    solve_OT: Solve a single OT problem.
+    solve_OT_batched: Solve multiple OT problems simultaneously (5-10x faster).
+
+Example:
+    >>> import torch
+    >>> from mdot_tnt import solve_OT, solve_OT_batched
+    >>>
+    >>> # Single problem
+    >>> r = torch.rand(512, device='cuda', dtype=torch.float64)
+    >>> r = r / r.sum()
+    >>> c = torch.rand(512, device='cuda', dtype=torch.float64)
+    >>> c = c / c.sum()
+    >>> C = torch.rand(512, 512, device='cuda', dtype=torch.float64)
+    >>> cost = solve_OT(r, c, C, gamma_f=1024.)
+    >>>
+    >>> # Batched (32 problems at once)
+    >>> r_batch = torch.rand(32, 512, device='cuda', dtype=torch.float64)
+    >>> r_batch = r_batch / r_batch.sum(-1, keepdim=True)
+    >>> c_batch = torch.rand(32, 512, device='cuda', dtype=torch.float64)
+    >>> c_batch = c_batch / c_batch.sum(-1, keepdim=True)
+    >>> costs = solve_OT_batched(r_batch, c_batch, C, gamma_f=1024.)
+"""
+
+import math
+import warnings
 
 import torch as th
 
-from mdot_tnt.mdot import mdot
+from mdot_tnt.batched import solve_OT_batched
+from mdot_tnt.mdot import mdot, preprocess_marginals
 from mdot_tnt.rounding import round_altschuler, rounded_cost_altschuler
 
+__all__ = ["solve_OT", "solve_OT_batched"]
+
 
-def solve_OT(r, c, C, gamma_f=4096., drop_tiny=False, return_plan=False, round=True, log=False):
+def solve_OT(r, c, C, gamma_f=1024.0, drop_tiny=False, return_plan=False, round=True, log=False):
     """
     Solve the entropic-regularized optimal transport problem. Inputs r, c, C are required to be torch tensors.
     :param r: n-dimensional row marginal.
@@ -23,21 +60,28 @@ def solve_OT(r, c, C, gamma_f=4096., drop_tiny=False, return_plan=False, round=T
     """
     assert all(isinstance(x, th.Tensor) for x in [r, c, C]), "r, c, and C must be torch tensors"
     dtype = r.dtype
-    if gamma_f > 2 ** 10:
+    # Require high precision for gamma_f > 2^10
+    if gamma_f > 2**10 and dtype != th.float64:
+        warnings.warn(
+            "Switching to double precision for gamma_f > 2^10 during execution. "
+            f"Output will be input dtype: {dtype}."
+        )
         r, c, C = r.double(), c.double(), C.double()
+
     if drop_tiny:
-        drop_lessthan = math.log(min(r.size(-1), c.size(-1))) / (gamma_f ** 2)
+        drop_lessthan = math.log(min(r.size(-1), c.size(-1))) / (gamma_f**2)
         (r_, r_keep), (c_, c_keep), C_ = preprocess_marginals(r, c, C, drop_lessthan)
 
         u_, v_, gamma_f_, k_total, opt_logs = mdot(r_, c_, C_, gamma_f)
 
-        u = -th.ones_like(r) * float('inf')
-        u[:, r_keep] = u_
-        v = -th.ones_like(c) * float('inf')
-        v[:, c_keep] = v_
+        u = -th.ones_like(r) * float("inf")
+        u[r_keep] = u_
+        v = -th.ones_like(c) * float("inf")
+        v[c_keep] = v_
     else:
         u, v, gamma_f_, k_total, opt_logs = mdot(r, c, C, gamma_f)
 
+    # Switch back to original dtype
     u, v = u.to(dtype), v.to(dtype)
 
     if return_plan: