mdot-tnt 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

mdot_tnt/__init__.py

@@ -1,11 +1,19 @@
+"""
+Code for solving the entropic-regularized optimal transport problem via the MDOT-TruncatedNewton (MDOT-TNT)
+method 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
+"""
 
+import math
 import torch as th
+import warnings
 
-from mdot_tnt.mdot import mdot
+from mdot_tnt.mdot import mdot, preprocess_marginals
 from mdot_tnt.rounding import round_altschuler, rounded_cost_altschuler
 
 
-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., 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,8 +31,12 @@ 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. "
+                      "Output will be input dtype: {}.".format(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)
         (r_, r_keep), (c_, c_keep), C_ = preprocess_marginals(r, c, C, drop_lessthan)
@@ -32,12 +44,13 @@ def solve_OT(r, c, C, gamma_f=4096., drop_tiny=False, return_plan=False, round=T
         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_
+        u[r_keep] = u_
         v = -th.ones_like(c) * float('inf')
-        v[:, c_keep] = v_
+        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:

mdot_tnt/mdot.py

@@ -1,12 +1,4 @@
-"""
-Code for solving the entropic-regularized optimal transport problem via the MDOT-TruncatedNewton (MDOT-TNT)
-method 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
-"""
 
-
-import math
 import warnings
 
 from mdot_tnt.rounding import *
@@ -16,18 +8,18 @@ from mdot_tnt.truncated_newton import TruncatedNewtonProjector
 def preprocess_marginals(r, c, C, eps):
     """
     This function drops the smallest entries whose cumulative sum equals 
-    :param r:
-    :param c:
-    :param C:
-    :param eps:
-    :return:
+    :param r: the row marginal.
+    :param c: the column marginal.
+    :param C: the cost matrix.
+    :param eps: the threshold for the cumulative sum of the marginal entries to be dropped
+    :return: the new marginals and the new cost matrix with the corresponding rows and columns dropped.
     """
 
     def preprocess_marginal(m, eps):
         m_sorted, m_idx = th.sort(m, dim=-1, descending=False)
         m_cumsum = th.cumsum(m_sorted, dim=-1)
         m_keep = m_idx[m_cumsum > eps]
-        m_new = m[:, m_keep]
+        m_new = m[m_keep]
         mass_removed = 1 - m_new.sum(-1)
         m_new = m_new + mass_removed / m_new.size(-1)
 
@@ -35,6 +27,8 @@ def preprocess_marginals(r, c, C, eps):
 
     r_new, r_keep = preprocess_marginal(r, eps)
     c_new, c_keep = preprocess_marginal(c, eps)
+    print("Dropped {} entries from r and {} entries from c.".format(
+        r.size(-1) - r_new.size(-1), c.size(-1) - c_new.size(-1)))
 
     C = C[r_keep][:, c_keep]
 
@@ -42,6 +36,15 @@ def preprocess_marginals(r, c, C, eps):
 
 
 def smooth_marginals(r, c, eps, w_r=0.5, w_c=0.5):
+    """
+    Smooth the marginals by adding a small amount of uniform mass to each entry.
+    :param r: the row marginal.
+    :param c: the column marginal.
+    :param eps: the amount of mass to add to each entry.
+    :param w_r: the weight for the row marginal.
+    :param w_c: the weight for the column marginal.
+    :return: the smoothed marginals with a total TV distance at most eps from the original marginals.
+    """
     assert w_r + w_c == 1, "w_r and w_c must sum to 1"
     eps = eps.clamp(max=1.).unsqueeze(-1)
     r_hat = (1 - w_r * eps) * r + w_r * eps * th.ones_like(r) / r.size(-1)
@@ -51,6 +54,12 @@ def smooth_marginals(r, c, eps, w_r=0.5, w_c=0.5):
 
 
 def adjust_schedule(q, deltas=None):
+    """
+    Adjust the temperature annealing schedule based on the success of the Truncated Newton method.
+    :param q: the current temperature annealing schedule adjustment factor.
+    :param deltas: the list of deltas from the Truncated Newton method; see Sec. 3.3 of Kemertas et al. (2025).
+    :return: the new temperature annealing schedule adjustment factor.
+    """
     if deltas is None:
         return q
 
@@ -65,7 +74,7 @@ def adjust_schedule(q, deltas=None):
     return q
 
 
-def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0**(1/3)):
+def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0):
     """
     Solve the entropic-regularized optimal transport problem using the MDOT method introduced in the paper:
     "Efficient and Accurate Optimal Transport with Mirror Descent and Conjugate Gradients" by Mete Kemertas,
@@ -78,7 +87,7 @@ def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0**(1/3)):
     :param gamma_i: The initial temperature.
     :param p: The exponent for the epsilon function, used to determine the stopping criterion for the dual gradient.
     :param q: The temperature annealing (or mirror descent step size) schedule adjustment factor.
-    :return:
+    :return: The dual variables u, v, the final temperature, the total number of O(n^2) primitive ops, and logs.
     """
     projector = TruncatedNewtonProjector(device=C.device, dtype=C.dtype)
 
@@ -109,8 +118,9 @@ def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0**(1/3)):
             u_cur, v_cur = u_init.clone(), v_init.clone()
 
         u_prev, v_prev = u_cur.clone(), v_cur.clone()
+        gamma_C = gamma * C
         u_cur, v_cur, proj_log, success = projector.project(
-            gamma * C, r_hat.log(), c_hat.log(), eps_d / 2, u_init, v_init)
+            gamma_C, r_hat.log(), c_hat.log(), eps_d / 2, u_init, v_init)
 
         logs["proj_logs"].append(proj_log)
         if not success:

mdot_tnt/truncated_newton.py

@@ -27,13 +27,13 @@ class TruncatedNewtonProjector:
         # a good enough solution to keep MDOT going.
         success_fn = lambda err_: err_ < 10 * eps_d
 
+        r = log_r.exp()
+        c = log_c.exp()
+
         # Each LSE operation costs 4 * n^2 operations.
         self.LSE_r = lambda v_: th.logsumexp(v_.unsqueeze(-2) - gamma_C, dim=-1)
         self.LSE_c = lambda u_: th.logsumexp(u_.unsqueeze(-1) - gamma_C, dim=-2)
 
-        r = log_r.exp()
-        c = log_c.exp()
-
         log_c_P = v + self.LSE_c(u)
         v += log_c - log_c_P  # Ensure c=c(P)
         log_r_P = u + self.LSE_r(v)
@@ -162,8 +162,7 @@ class TruncatedNewtonProjector:
         rho = self.rho
         tol = err * eta_k
 
-        def matmul_PPc(x_):
-            return (P @ ((x_.unsqueeze(-2) @ P).transpose(-2, -1) / c.unsqueeze(-1))).squeeze(-1)
+        matmul_PPc = lambda x_: (P @ ((x_.unsqueeze(-2) @ P).transpose(-2, -1) / c.unsqueeze(-1))).squeeze(-1)
 
         # mml = th.compile(matmul_PPc)
         mml = matmul_PPc

{mdot_tnt-0.1.0.dist-info → mdot_tnt-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mdot-tnt
-Version: 0.1.0
+Version: 0.2.0
 Summary: A fast, GPU-parallel, PyTorch-compatible optimal transport solver.
 Author-email: Mete Kemertas <kemertas@cs.toronto.edu>
 Project-URL: Homepage, https://github.com/metekemertas/mdot_tnt
@@ -19,7 +19,7 @@ The current implementation is based on PyTorch and is compatible with both CPU a
 For installation:
 First, install PyTorch following the instructions at https://pytorch.org/get-started/locally/ to select the version that matches your system's configuration.
 ```bash
-pip install mdot_tnt
+pip3 install mdot_tnt
 ```
 
 Quickstart guide:
@@ -33,18 +33,18 @@ N, M, dim = 100, 200, 128
 r = th.distributions.Dirichlet(th.ones(N)).sample()
 c = th.distributions.Dirichlet(th.ones(M)).sample()
 
-# Cost matrix from pairwise Euclidean distances of random points in R^100
+# Cost matrix from pairwise Euclidean distances squared given random points in R^100
 x = th.distributions.MultivariateNormal(th.zeros(dim), th.eye(dim)).sample((N,))
 y = th.distributions.MultivariateNormal(th.zeros(dim), th.eye(dim)).sample((M,))
-C = th.cdist(x, y, p=2)
+C = th.cdist(x, y, p=2) ** 2
 C /= C.max()  # Normalize cost matrix to meet convention.
 
 # Use double precision for numerical stability in high precision regime.
-r, c, C = r.double(), c.double(), C.double()
+r, c, C = r.double().to(device), c.double().to(device), C.double().to(device)
 
 # Solve OT problem. Increase (decrease) gamma_f for higher (lower) precision.
-# Default is gamma_f=2**12. Expect error of order logn / gamma_f at worst, and possibly lower.
-cost = mdot_tnt.solve_OT(r, c, C, gamma_f=2**12)
+# Default is gamma_f=2**10. Expect error of order logn / gamma_f at worst, and possibly lower.
+cost = mdot_tnt.solve_OT(r, c, C, gamma_f=2**10)
 
 # To return a feasible transport plan, use the following:
 transport_plan = mdot_tnt.solve_OT(r, c, C, gamma_f=2**12, return_plan=True)

mdot_tnt-0.2.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+mdot_tnt/__init__.py,sha256=9_DJl0mXe5WiKFf_12oN4d2Bnq4Q3AcuuNz9A6FmFp8,3290
+mdot_tnt/mdot.py,sha256=8SZxldnq64ySWZlUjsdGqT_prqfcs2ZrVFBGp-EJ4B0,5562
+mdot_tnt/rounding.py,sha256=Q7QBPsFzBqnMZKnlV147ruStUCme6gwt6HnjTjqBezk,2405
+mdot_tnt/truncated_newton.py,sha256=Zp4Tb65dSE_AGDElgAE1gp3V_CaXFsbDX3SJCaNEWfc,10299
+mdot_tnt-0.2.0.dist-info/LICENSE,sha256=sXw3FpVqouAddNhfwD6nXaSgABFyLnXuSn_ghjU0AhY,1837
+mdot_tnt-0.2.0.dist-info/METADATA,sha256=sfqk_nDPsGJlSrXoz982dgLLAOGqI1F1obSdOsodIN8,3022
+mdot_tnt-0.2.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+mdot_tnt-0.2.0.dist-info/top_level.txt,sha256=HmxTNtoLH7F20hgZVFdfUowIQ2fviSX64wSG1HP8Iao,9
+mdot_tnt-0.2.0.dist-info/RECORD,,

mdot_tnt-0.1.0.dist-info/RECORD

@@ -1,9 +0,0 @@
-mdot_tnt/__init__.py,sha256=zi0ujbN19nv_dt8wdxPQgsOuDi-4gIg4sq6y9fDfAkI,2657
-mdot_tnt/mdot.py,sha256=EUch7ECZxDWDoNW3S84qyCx6-qbG6Mu1_XlWow862Wc,4620
-mdot_tnt/rounding.py,sha256=Q7QBPsFzBqnMZKnlV147ruStUCme6gwt6HnjTjqBezk,2405
-mdot_tnt/truncated_newton.py,sha256=YEkWJxL5H4XsuWQzfdZ0E5Ju3Npk3kKoe2l_go8wq2A,10314
-mdot_tnt-0.1.0.dist-info/LICENSE,sha256=sXw3FpVqouAddNhfwD6nXaSgABFyLnXuSn_ghjU0AhY,1837
-mdot_tnt-0.1.0.dist-info/METADATA,sha256=2Ou4X10K6bUGHQhnMXoFjrv-4YUeIyreSKibMLM2gq0,2972
-mdot_tnt-0.1.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-mdot_tnt-0.1.0.dist-info/top_level.txt,sha256=HmxTNtoLH7F20hgZVFdfUowIQ2fviSX64wSG1HP8Iao,9
-mdot_tnt-0.1.0.dist-info/RECORD,,