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

mdot_tnt/mdot.py

@@ -1,21 +1,33 @@
+"""Core MDOT solver using truncated Newton projection."""
 
 import warnings
+from typing import Any, Dict, List, Tuple, Union
+
+import torch as th
 
-from mdot_tnt.rounding import *
 from mdot_tnt.truncated_newton import TruncatedNewtonProjector
 
 
-def preprocess_marginals(r, c, C, eps):
+def preprocess_marginals(
+    r: th.Tensor, c: th.Tensor, C: th.Tensor, eps: float
+) -> Tuple[Tuple[th.Tensor, th.Tensor], Tuple[th.Tensor, th.Tensor], th.Tensor]:
     """
-    This function drops the smallest entries whose cumulative sum equals 
-    :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.
+    Drop the smallest marginal entries whose cumulative sum is below a threshold.
+
+    Args:
+        r: The row marginal of shape (n,).
+        c: The column marginal of shape (m,).
+        C: The cost matrix of shape (n, m).
+        eps: The threshold for the cumulative sum of the marginal entries to be dropped.
+
+    Returns:
+        A tuple containing:
+            - (r_new, r_keep): The new row marginal and indices of kept entries.
+            - (c_new, c_keep): The new column marginal and indices of kept entries.
+            - C: The cost matrix with corresponding rows and columns dropped.
     """
 
-    def preprocess_marginal(m, eps):
+    def preprocess_marginal(m: th.Tensor, eps: float) -> Tuple[th.Tensor, th.Tensor]:
         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]
@@ -27,76 +39,115 @@ 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)))
+    print(
+        f"Dropped {r.size(-1) - r_new.size(-1)} entries from r and {c.size(-1) - c_new.size(-1)} entries from c."
+    )
 
     C = C[r_keep][:, c_keep]
 
     return (r_new, r_keep), (c_new, c_keep), C
 
 
-def smooth_marginals(r, c, eps, w_r=0.5, w_c=0.5):
+def smooth_marginals(
+    r: th.Tensor,
+    c: th.Tensor,
+    eps: th.Tensor,
+    w_r: float = 0.5,
+    w_c: float = 0.5,
+) -> Tuple[th.Tensor, th.Tensor]:
     """
     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.
+
+    Args:
+        r: The row marginal of shape (n,).
+        c: The column marginal of shape (m,).
+        eps: The amount of mass to add to each entry.
+        w_r: The weight for the row marginal.
+        w_c: The weight for the column marginal.
+
+    Returns:
+        A tuple (r_hat, c_hat) of smoothed marginals with 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)
+    eps = eps.clamp(max=1.0).unsqueeze(-1)
     r_hat = (1 - w_r * eps) * r + w_r * eps * th.ones_like(r) / r.size(-1)
     c_hat = (1 - w_c * eps) * c + w_c * eps * th.ones_like(c) / c.size(-1)
 
     return r_hat, c_hat
 
 
-def adjust_schedule(q, deltas=None):
+def adjust_schedule(q: float, deltas: Union[List[float], None] = None) -> float:
     """
     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.
+
+    Args:
+        q: The current temperature annealing schedule adjustment factor.
+        deltas: The list of deltas from the Truncated Newton method;
+                see Sec. 3.3 of Kemertas et al. (2025).
+
+    Returns:
+        The new temperature annealing schedule adjustment factor.
     """
     if deltas is None:
         return q
 
-    deltas = deltas + [1.]  # If deltas is empty, we assume that the first iteration was successful
+    deltas = deltas + [1.0]  # If deltas is empty, we assume that the first iteration was successful
     delta_min = min(deltas)
 
     if delta_min < 0.5:
-        q = q ** 0.5
+        q = q**0.5
     elif delta_min > 0.9:
-        q = q ** 2
+        q = q**2
 
     return q
 
 
-def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0):
+def mdot(
+    r: th.Tensor,
+    c: th.Tensor,
+    C: th.Tensor,
+    gamma_f: float,
+    gamma_i: float = 16,
+    p: float = 1.5,
+    q: float = 2.0,
+) -> Tuple[th.Tensor, th.Tensor, float, int, Dict[str, Any]]:
     """
-    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,
-    Allan D. Jepson and Amir-massoud Farahmand. URL: https://arxiv.org/abs/2307.08507
+    Solve the entropic-regularized optimal transport problem using the MDOT method.
+
+    This implements the MDOT method introduced in the paper:
+    "Efficient and Accurate Optimal Transport with Mirror Descent and Conjugate Gradients"
+    by Mete Kemertas, Allan D. Jepson and Amir-massoud Farahmand.
+    URL: https://arxiv.org/abs/2307.08507
+
     Here, we use the Truncated Newton method for projection.
-    :param r: The first marginal.
-    :param c: The second marginal.
-    :param C: The cost matrix. Recommended use is to scale the entries to be in [0, 1].
-    :param gamma_f: The final temperature (inverse of the regularization weight).
-    :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: The dual variables u, v, the final temperature, the total number of O(n^2) primitive ops, and logs.
+
+    Args:
+        r: The first marginal of shape (n,).
+        c: The second marginal of shape (m,).
+        C: The cost matrix of shape (n, m). Recommended to scale entries to [0, 1].
+        gamma_f: The final temperature (inverse of the regularization weight).
+        gamma_i: The initial temperature.
+        p: The exponent for the epsilon function, used to determine the stopping
+           criterion for the dual gradient.
+        q: The temperature annealing (or mirror descent step size) schedule adjustment factor.
+
+    Returns:
+        A tuple containing:
+            - u: The row dual variables of shape (n,).
+            - v: The column dual variables of shape (m,).
+            - gamma: The final temperature achieved.
+            - k_total: The total number of O(n^2) primitive operations.
+            - logs: Dictionary with optimization statistics.
     """
     projector = TruncatedNewtonProjector(device=C.device, dtype=C.dtype)
 
     H_r = -(r * (r + 1e-30).log()).sum(-1)
     H_c = -(c * (c + 1e-30).log()).sum(-1)
     H_min = th.min(H_r, H_c)
-    eps_fn = lambda g_: H_min / (g_ ** p)
+    eps_fn = lambda g_: H_min / (g_**p)
 
-    logs = {
+    logs: Dict[str, Any] = {
         "proj_logs": [],
         "eps": [],
     }
@@ -104,7 +155,7 @@ def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0):
     t = 1
     done = False
     gamma = min(gamma_i, gamma_f)
-    gammas = [0., gamma]
+    gammas = [0.0, gamma]
 
     while not done:
         done = abs(gamma - gamma_f) < 1e-5  # Check if gamma == gamma_f (modulo rounding errors)
@@ -120,11 +171,14 @@ def mdot(r, c, C, gamma_f, gamma_i=16, p=1.5, q=2.0):
         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:
-            warnings.warn("Projection failed. Returning result at the last temperature: {:.4e}".format(1 / gammas[-2]))
+            warnings.warn(
+                f"Projection failed. Returning result at the last temperature: {1 / gammas[-2]:.4e}"
+            )
             u_cur = u_prev.clone()
             v_cur = v_prev.clone()
             gammas = gammas[:-1]

mdot_tnt/rounding.py

@@ -6,16 +6,22 @@ and c are the row and column marginals, respectively. The algorithm is used in t
 plan and compute the cost of the rounded plan. The implementation is based on the original paper.
 """
 
+from typing import Union
+
 import torch as th
 
 
-def round_altschuler(P, r, c):
+def round_altschuler(P: th.Tensor, r: th.Tensor, c: th.Tensor) -> th.Tensor:
     """
     Performs rounding given a transport plan and marginals.
-    :param P: the input transport plan
-    :param r: row marginal
-    :param c: column marginal
-    :return: rounded transport plan in feasible set U(r, c).
+
+    Args:
+        P: The input transport plan of shape (n, m).
+        r: Row marginal of shape (n,).
+        c: Column marginal of shape (m,).
+
+    Returns:
+        Rounded transport plan in feasible set U(r, c).
     """
     X = th.min(r / P.sum(-1), th.ones_like(r))
     P *= X.unsqueeze(-1)
@@ -25,19 +31,39 @@ def round_altschuler(P, r, c):
 
     err_r = (r - P.sum(-1)).clamp(min=0)
     err_c = (c - P.sum(-2)).clamp(min=0)
-    P += err_r.unsqueeze(-1) @ err_c.unsqueeze(-2) / (err_r.norm(p=1, dim=-1, keepdim=True) + 1e-30).unsqueeze(-1)
+    P += (
+        err_r.unsqueeze(-1)
+        @ err_c.unsqueeze(-2)
+        / (err_r.norm(p=1, dim=-1, keepdim=True) + 1e-30).unsqueeze(-1)
+    )
 
     return P
 
 
-def rounded_cost_altschuler(u, v, r, c, C, gamma):
-    """Performs rounding and cost computation in logdomain given dual variables, without storing n^2 matrices.
-    :param u: dual variable for rows
-    :param v: dual variable for columns
-    :param r: row marginal
-    :param c: column marginal
-    :param C: cost matrix
-    :param gamma: temperature, i.e., the inverse of the entropic regularization weight.
+def rounded_cost_altschuler(
+    u: th.Tensor,
+    v: th.Tensor,
+    r: th.Tensor,
+    c: th.Tensor,
+    C: th.Tensor,
+    gamma: Union[float, th.Tensor],
+) -> th.Tensor:
+    """
+    Performs rounding and cost computation in log-domain given dual variables.
+
+    This function computes the transport cost without storing the full n×m transport plan,
+    making it memory efficient.
+
+    Args:
+        u: Dual variable for rows of shape (n,).
+        v: Dual variable for columns of shape (m,).
+        r: Row marginal of shape (n,).
+        c: Column marginal of shape (m,).
+        C: Cost matrix of shape (n, m).
+        gamma: Temperature (inverse of the entropic regularization weight).
+
+    Returns:
+        The optimal transport cost as a scalar tensor.
     """
     r_P_log = u + th.logsumexp(v.unsqueeze(-2) - gamma * C, dim=-1)
     delta_u = th.min(r.log() - r_P_log, th.zeros_like(r))
@@ -50,7 +76,7 @@ def rounded_cost_altschuler(u, v, r, c, C, gamma):
     r_P_log = u + th.logsumexp(v.unsqueeze(-2) - gamma * C, dim=-1)
     r_P = r_P_log.exp()
     err_r = r - r_P
-    err_r /= (err_r.norm(p=1, dim=-1, keepdim=True) + 1e-30)
+    err_r /= err_r.norm(p=1, dim=-1, keepdim=True) + 1e-30
 
     c_P_log = v + th.logsumexp(u.unsqueeze(-1) - gamma * C, dim=-2)
     c_P = c_P_log.exp()

mdot_tnt/truncated_newton.py

@@ -1,27 +1,67 @@
+"""Truncated Newton projector for the MDOT algorithm."""
 
-import torch as th
 import warnings
+from typing import Any, Callable, Dict, Tuple, Union
+
+import torch as th
 
 
 class TruncatedNewtonProjector:
-    def __init__(self, device, dtype, **kwargs):
+    """
+    Truncated Newton projector for the MDOT algorithm.
+
+    Projects onto the set of couplings satisfying marginal constraints using
+    a preconditioned conjugate gradient method within a Newton framework.
+    """
+
+    def __init__(self, device: th.device, dtype: th.dtype, **kwargs: Any) -> None:
+        """
+        Initialize the projector.
+
+        Args:
+            device: PyTorch device for computations.
+            dtype: Data type for tensors.
+            **kwargs: Additional options (debug: bool for verbose output).
+        """
         self.device = device
         self.rho = th.zeros(1, device=device, dtype=dtype)
-        self.debug = kwargs.get('debug', False)
-
-    def project(self, gamma_C, log_r, log_c, eps_d, u, v):
+        self.debug = kwargs.get("debug", False)
+        self.LSE_r: Callable[[th.Tensor], th.Tensor]
+        self.LSE_c: Callable[[th.Tensor], th.Tensor]
+
+    def project(
+        self,
+        gamma_C: th.Tensor,
+        log_r: th.Tensor,
+        log_c: th.Tensor,
+        eps_d: Union[float, th.Tensor],
+        u: th.Tensor,
+        v: th.Tensor,
+    ) -> Tuple[th.Tensor, th.Tensor, Dict[str, Any], bool]:
         """
         Project onto the set of couplings that satisfy the marginal constraints.
-        :param gamma_C: The cost matrix scaled by gamma.
-        :param log_r:
+
+        Args:
+            gamma_C: The cost matrix scaled by gamma, shape (n, m).
+            log_r: Log of row marginals, shape (n,).
+            log_c: Log of column marginals, shape (m,).
+            eps_d: Convergence tolerance for the dual gradient norm.
+            u: Initial row dual variables, shape (n,).
+            v: Initial column dual variables, shape (m,).
+
+        Returns:
+            u: Updated row dual variables.
+            v: Updated column dual variables.
+            logs: Dictionary with optimization statistics.
+            success: Whether projection converged successfully.
         """
-        logs = {
+        logs: Dict[str, Any] = {
             "errs": [],
-            'ls_func_cnt': 0,
-            'chisinkhorn_steps': 0,
-            'newtonsolve_steps': 0,
+            "ls_func_cnt": 0,
+            "chisinkhorn_steps": 0,
+            "newtonsolve_steps": 0,
             "deltas": [],  # Ratios of actual to theoretically predicted (ideal) reduction in gradient norm.
-            "all_newtonsolve_steps": []
+            "all_newtonsolve_steps": [],
         }
         # In case of errors or issues, 10 times the tolerance level is considered
         # a good enough solution to keep MDOT going.
@@ -57,19 +97,20 @@ class TruncatedNewtonProjector:
             self.rho = th.max(th.zeros_like(self.rho), self.rho)
 
             P = th.exp(u.unsqueeze(-1) + v.unsqueeze(-2) - gamma_C)
-            diag_PPc = ((P ** 2) / c.unsqueeze(-2)).sum(-1)
+            diag_PPc = ((P**2) / c.unsqueeze(-2)).sum(-1)
             k += 8
             delta_u, delta_v, matmul_cnt, rho, pcg_success = self.newton_solve(
-                P, c, diag_PPc, grad_k, r_P, err, beta, eta_k, maxIter=5000)
+                P, c, diag_PPc, grad_k, r_P, err, beta, eta_k, maxIter=5000
+            )
             del P  # Free up memory
             if not pcg_success:
                 k += matmul_cnt
                 logs["n_iter"] = k
-                msg = "PCG did not converge. TruncatedNewton returning with success={}".format(success_fn(err))
+                msg = f"PCG did not converge. TruncatedNewton returning with success={success_fn(err)}"
                 warnings.warn(msg)
                 return u, v, logs, success_fn(err)
 
-            self.rho = th.max(th.zeros_like(self.rho), 1. - (1. - rho) * 4.)
+            self.rho = th.max(th.zeros_like(self.rho), 1.0 - (1.0 - rho) * 4.0)
             k += matmul_cnt
             logs["newtonsolve_steps"] += matmul_cnt
 
@@ -79,8 +120,7 @@ class TruncatedNewtonProjector:
             linear_decr = -(grad_k * delta_u).sum(-1, keepdim=True)
             if not linear_decr > 0:
                 logs["n_iter"] = k
-                msg = "Linear decrease condition not satisfied. TruncatedNewton returning with success={}".format(
-                    success_fn(err))
+                msg = f"Linear decrease condition not satisfied. TruncatedNewton returning with success={success_fn(err)}"
                 warnings.warn(msg)
                 return u, v, logs, success_fn(err)
 
@@ -89,8 +129,7 @@ class TruncatedNewtonProjector:
                 alpha *= 0.5
                 if alpha < 1e-9:
                     logs["n_iter"] = k
-                    msg = "Line search did not converge. TruncatedNewton returning with success={}".format(
-                        success_fn(err))
+                    msg = f"Line search did not converge. TruncatedNewton returning with success={success_fn(err)}"
                     warnings.warn(msg)
                     return u, v, logs, success_fn(err)
 
@@ -113,13 +152,17 @@ class TruncatedNewtonProjector:
             log_r_P = u + self.LSE_r(v)
             k += 4
 
-            u, v, log_r_P, err, k_ = self.chi_sinkhorn(u, v, log_r, log_c, log_r_P, eps_d ** (2 / 5))
+            u, v, log_r_P, err, k_ = self.chi_sinkhorn(
+                u, v, log_r, log_c, log_r_P, eps_d ** (2 / 5)
+            )
             r_P = log_r_P.exp()
             logs["chisinkhorn_steps"] += k_
             k += k_
 
             logs["errs"].append(err)
-            logs["deltas"].append(th.min((logs["errs"][-2] - err_before_sk) / ((1 - eta_k) * logs["errs"][-2])).item())
+            logs["deltas"].append(
+                th.min((logs["errs"][-2] - err_before_sk) / ((1 - eta_k) * logs["errs"][-2])).item()
+            )
 
         if u.isnan().any() or v.isnan().any():
             raise ValueError("NaNs encountered in u or v")
@@ -132,7 +175,16 @@ class TruncatedNewtonProjector:
 
         return u, v, logs, True
 
-    def chi_sinkhorn(self, u, v, log_r, log_c, log_r_P, eps_chi, maxOps=float('inf')):
+    def chi_sinkhorn(
+        self,
+        u: th.Tensor,
+        v: th.Tensor,
+        log_r: th.Tensor,
+        log_c: th.Tensor,
+        log_r_P: th.Tensor,
+        eps_chi: Union[float, th.Tensor],
+        maxOps: float = float("inf"),
+    ) -> Tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor, int]:
         k = 0
         r = log_r.exp()
         err = (r - log_r_P.exp()).norm(p=1, dim=-1)
@@ -154,27 +206,40 @@ class TruncatedNewtonProjector:
             k += 8
 
         if k >= maxOps:
-            raise ValueError("Chi-Sinkhorn did not converge in maxIter={} steps".format(maxOps))
+            raise ValueError(f"Chi-Sinkhorn did not converge in maxIter={maxOps} steps")
 
         return u, v, log_r_P, err, k
 
-    def newton_solve(self, P, c, diag_PPc, grad_k, r_P, err, beta=0.5, eta_k=0.5, maxIter=500):
+    def newton_solve(
+        self,
+        P: th.Tensor,
+        c: th.Tensor,
+        diag_PPc: th.Tensor,
+        grad_k: th.Tensor,
+        r_P: th.Tensor,
+        err: th.Tensor,
+        beta: float = 0.5,
+        eta_k: Union[float, th.Tensor] = 0.5,
+        maxIter: int = 500,
+    ) -> Tuple[th.Tensor, th.Tensor, int, th.Tensor, bool]:
         rho = self.rho
         tol = err * eta_k
 
-        matmul_PPc = lambda x_: (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
 
-        M = lambda rho_: (r_P - rho_ * diag_PPc)  # Diagonal preconditioner
+        M = lambda rho_: r_P - rho_ * diag_PPc  # Diagonal preconditioner
         M_rho = M(th.ones_like(self.rho))
         M_rho[M_rho <= 0] = M_rho[M_rho > 0].min()
 
         x0 = -grad_k / M_rho
         PPc_x0 = mml(x0)
         matmul_cnt = 2
-        r_P_x0 = (r_P * x0)
+        r_P_x0 = r_P * x0
 
         x = x0.clone()
         PPc_x = PPc_x0.clone()
@@ -197,7 +262,7 @@ class TruncatedNewtonProjector:
             best_sol[r_true_norm < best_r_true_norm] = x[r_true_norm < best_r_true_norm]
             best_r_true_norm = th.min(r_true_norm, best_r_true_norm)
 
-            rho[r_true_norm > tol] = 1. - (1. - rho[r_true_norm > tol]) * 0.25
+            rho[r_true_norm > tol] = 1.0 - (1.0 - rho[r_true_norm > tol]) * 0.25
             M_rho = M(rho)
 
             if matmul_cnt > 0:
@@ -227,8 +292,10 @@ class TruncatedNewtonProjector:
 
                 quad = (Fr_p * p).sum(-1, keepdim=True)
                 if (quad <= 0)[best_r_true_norm > tol].any():
-                    warnings.warn("Warning: negative curvature encountered in CG. Returning best solution. "
-                          "Residual norm less than error: {}".format((best_r_true_norm < err).item()))
+                    warnings.warn(
+                        "Warning: negative curvature encountered in CG. Returning best solution. "
+                        f"Residual norm less than error: {(best_r_true_norm < err).item()}"
+                    )
                     x = best_sol.clone()
                     done = True
                     success = best_r_true_norm < err
@@ -241,12 +308,15 @@ class TruncatedNewtonProjector:
                 res += alpha * Fr_p
                 r_norm = res.norm(p=1, dim=-1)
 
-                if th.isnan(r_norm)[best_r_true_norm > tol].any() or th.isinf(r_norm)[best_r_true_norm > tol].any():
+                if (
+                    th.isnan(r_norm)[best_r_true_norm > tol].any()
+                    or th.isinf(r_norm)[best_r_true_norm > tol].any()
+                ):
                     raise ValueError("NaNs or infs encountered in r_norm")
 
                 PPc_x += alpha * PPc_p
 
-                r_P_x = (r_P * x)
+                r_P_x = r_P * x
                 res_true = r_P_x - PPc_x + grad_k
                 r_true_norm = res_true.norm(p=1, dim=-1)
                 best_sol[r_true_norm < best_r_true_norm] = x[r_true_norm < best_r_true_norm]
@@ -279,4 +349,4 @@ class TruncatedNewtonProjector:
         Pc_x = ((x.unsqueeze(-2) @ P).transpose(-2, -1) / c.unsqueeze(-1)).squeeze(-1)
         matmul_cnt += 1
 
-        return x, -Pc_x, matmul_cnt, rho, success
+        return x, -Pc_x, matmul_cnt, rho, success