polystep 0.7.0__py3-none-any.whl

polystep/solvers/_prelude.py

@@ -0,0 +1,129 @@
+"""Shared entry-point preparation for the OT / weighting solvers.
+
+Every log-sum-exp / softmax solver needs the same preamble before it can run:
+promote half precision to FP32 (BF16's 7 mantissa bits collapse the row-max
+trick once the cost spread exceeds ~15 nats), replace non-finite costs with a
+finite penalty, default and device/dtype-align the marginals, and coerce
+warm-start duals onto the cost tensor. Centralizing it keeps the variants from
+drifting apart (e.g. one solver gaining FP32 promotion while another silently
+NaNs on the same input).
+
+Design note - no per-step host syncs: ``sanitize_cost`` is branch-free (no
+``.item()`` / ``.all()`` in a Python ``if``), and ``align_marginal`` /
+``align_dual`` only move tensors (``.to`` is a no-op when already aligned).
+Value checks that would force a device->host sync are intentionally omitted so
+the hot path (a fresh solve every optimizer step) stays GPU-resident.
+"""
+
+import warnings
+from typing import Optional
+
+import torch
+
+
+def validate_positive(value: float, name: str, context: str = "") -> None:
+    """Raise ValueError unless ``value > 0`` (a plain Python-float check).
+
+    Solvers store their temperature (``epsilon`` / ``tau``) as a mutable
+    attribute that schedules overwrite per step, so this is re-checked inside
+    ``solve()`` rather than only at construction.
+    """
+    if not value > 0:
+        msg = f"{name} must be > 0, got {value}."
+        if context:
+            msg += " " + context
+        raise ValueError(msg)
+
+
+def sanitize_cost(cost_matrix: torch.Tensor) -> torch.Tensor:
+    """Promote half precision to FP32 and replace non-finite costs, on-device.
+
+    Non-finite entries (a hard-constraint ``+inf`` or an upstream NaN) are
+    replaced with ``2 * max|finite| + 1`` (floored at ``1e6``) so a masked
+    vertex still gets near-zero weight without sending ``-C/eps`` to ``-inf``
+    and NaN-ing the whole row. Branch-free: no host sync on the finite path.
+    """
+    if cost_matrix.dtype in (torch.bfloat16, torch.float16):
+        cost_matrix = cost_matrix.to(torch.float32)
+    if cost_matrix.numel() == 0:
+        return cost_matrix
+    finite = torch.isfinite(cost_matrix)
+    max_finite = torch.where(finite, cost_matrix, cost_matrix.new_zeros(())).abs().amax()
+    penalty = torch.clamp(max_finite * 2.0 + 1.0, min=1e6)
+    return torch.where(finite, cost_matrix, penalty)
+
+
+def recenter_cost(cost_matrix: torch.Tensor):
+    """Subtract the per-matrix minimum so ``|C|`` stays bounded.
+
+    The entropic-OT plan is invariant to a constant shift of the cost:
+    ``softmax(-C/eps)`` is shift-invariant, and a shift only translates the
+    dual potentials, leaving the plan ``P`` unchanged. Keeping ``min(C)=0``
+    preserves FP32 precision in the log-sum-exp and the ``exp((f+g-C)/eps)``
+    reconstruction when ``|C|`` is much larger than ``eps``, and removes the
+    ``+inf`` logit that would NaN a softmax at tiny ``eps``.
+
+    Returns ``(shifted_cost, shift)`` with ``shift = C.min()``. Add
+    ``shift * a.sum()`` back to any reported ``<C, P>`` or dual value (plan
+    mass ``sum(P)`` equals ``a.sum()``). One reduction, no host sync.
+    """
+    if cost_matrix.numel() == 0:
+        return cost_matrix, cost_matrix.new_zeros(())
+    shift = cost_matrix.amin()
+    return cost_matrix - shift, shift
+
+
+def align_marginal(
+    a: Optional[torch.Tensor],
+    n: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    name: str = "a",
+) -> torch.Tensor:
+    """Return a length-``n`` marginal on ``(device, dtype)``.
+
+    ``None`` -> uniform ``1/n``. A provided marginal is moved onto the cost
+    tensor (a no-op when already aligned), shape-checked, and value-checked
+    (finite, nonnegative, positive total mass) so a malformed marginal fails
+    loudly instead of being silently clamped to an infeasible plan before the
+    ``log``. The value checks host-sync, so they run *only* on the user-supplied
+    path: the integrated optimizer passes ``a=None`` for its uniform marginal
+    (see :func:`solver.PolyStep.init_state`), keeping the per-step solve
+    sync-free.
+    """
+    if a is None:
+        return torch.full((n,), 1.0 / n, device=device, dtype=dtype)
+    a = a.to(device=device, dtype=dtype)
+    if a.shape != (n,):
+        raise ValueError(f"marginal {name} must have shape ({n},), got {tuple(a.shape)}.")
+    if not torch.isfinite(a).all():
+        raise ValueError(f"marginal {name} contains non-finite entries (NaN/Inf).")
+    if (a < 0).any():
+        raise ValueError(f"marginal {name} has negative entries; a transport marginal must be nonnegative.")
+    if not a.sum() > 0:
+        raise ValueError(f"marginal {name} has nonpositive total mass; it must sum to a positive value.")
+    return a
+
+
+def align_dual(
+    init: Optional[torch.Tensor],
+    n: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    name: str = "init",
+) -> Optional[torch.Tensor]:
+    """Move a warm-start dual onto ``(device, dtype)``, or ``None`` on mismatch.
+
+    Returns a fresh (cloned) tensor the caller may mutate in place, or ``None``
+    when the shape does not match (caller falls back to a zero init, after a
+    warning). The ``.to`` is a no-op when the dual is already aligned.
+    """
+    if init is None:
+        return None
+    if init.shape != (n,):
+        warnings.warn(
+            f"warm-start {name} shape mismatch: expected ({n},), got {tuple(init.shape)}. Falling back to zeros.",
+            stacklevel=2,
+        )
+        return None
+    return init.to(device=device, dtype=dtype).clone()

polystep/solvers/base.py

@@ -0,0 +1,77 @@
+"""Solver abstraction layer: Protocol and base result type.
+
+Defines the ``Solver`` protocol for pluggable OT/weighting solvers and the
+``SolverResult`` dataclass that all solvers return. Uses structural typing
+(Protocol) so any class with a matching ``.solve()`` signature qualifies
+without explicit inheritance.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+
+
+@dataclass
+class SolverResult:
+    """Base result from any solver.
+
+    Attributes:
+        matrix: Transport plan / weight matrix of shape (P, V).
+        cost: Regularized cost scalar.
+        f: First dual potential (Sinkhorn) or None (softmax).
+        g: Second dual potential (Sinkhorn) or None (softmax).
+        converged: Whether the solver converged within tolerance.
+        n_iters: Number of iterations actually run.
+        ent_reg_cost: Entropic regularized cost of the transport plan. For the
+            Sinkhorn dual this is <f, a> + <g, b> - eps * sum_ij P_ij (true plan
+            mass; equals the converged value when sum(P) == sum(a)).
+    """
+
+    matrix: torch.Tensor
+    cost: float
+    f: Optional[torch.Tensor] = None
+    g: Optional[torch.Tensor] = None
+    converged: bool = True
+    n_iters: int = 1
+    ent_reg_cost: float = 0.0
+
+
+class Solver:
+    """Protocol for pluggable solvers.
+
+    Any class that implements ``solve()`` with the following signature
+    qualifies as a ``Solver`` via structural typing (duck typing). This
+    is intentionally NOT a ``typing.Protocol`` subclass to avoid runtime
+    metaclass conflicts with ``@dataclass``; conformance is checked by
+    the type checker via structural subtyping.
+
+    Attributes:
+        epsilon: Entropic regularization / temperature parameter.
+    """
+
+    epsilon: float
+
+    def solve(
+        self,
+        cost_matrix: torch.Tensor,
+        a: Optional[torch.Tensor] = None,
+        b: Optional[torch.Tensor] = None,
+        init_f: Optional[torch.Tensor] = None,
+        init_g: Optional[torch.Tensor] = None,
+        scale_cost: Optional[Union[str, float]] = None,
+    ) -> SolverResult:
+        """Solve the weighting / transport problem.
+
+        Args:
+            cost_matrix: Cost matrix C of shape (n, m).
+            a: Source marginal of shape (n,). Defaults to uniform 1/n.
+            b: Target marginal of shape (m,). Defaults to uniform 1/m.
+            init_f: Warm-start first dual potential of shape (n,).
+            init_g: Warm-start second dual potential of shape (m,).
+            scale_cost: Cost scaling: 'mean', 'max_cost', or float divisor.
+
+        Returns:
+            SolverResult with transport plan and cost.
+        """
+        raise NotImplementedError

polystep/solvers/greedy.py

@@ -0,0 +1,159 @@
+"""Greedy solvers: deterministic update rules for ablation studies.
+
+Provides two non-iterative solvers that bypass entropic OT entirely:
+
+- ``MinCostGreedySolver``: Each particle moves to its single lowest-cost vertex.
+- ``TopKMeanSolver``: Each particle moves to the uniform average of its k
+  lowest-cost vertices.
+
+Both produce transport matrices compatible with the standard barycentric
+projection (Eq. 7) and conform to the ``Solver`` protocol.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+
+from ..costs import scale_cost_matrix
+from .base import SolverResult
+
+
+@dataclass
+class MinCostGreedySolver:
+    """Greedy argmin solver: each particle moves to its lowest-cost vertex.
+
+    For each row i of the cost matrix, assigns all mass ``a[i]`` to the
+    single vertex with minimum cost:  ``T[i, argmin_v C[i,v]] = a[i]``,
+    zero elsewhere.
+
+    Attributes:
+        epsilon: Accepted for API compatibility; unused by greedy assignment.
+        compile: Accepted for API compatibility; unused.
+    """
+
+    epsilon: float = 0.1
+    compile: bool = False
+
+    def solve(
+        self,
+        cost_matrix: torch.Tensor,
+        a: Optional[torch.Tensor] = None,
+        b: Optional[torch.Tensor] = None,
+        init_f: Optional[torch.Tensor] = None,
+        init_g: Optional[torch.Tensor] = None,
+        scale_cost: Optional[Union[str, float]] = None,
+    ) -> SolverResult:
+        """Compute greedy argmin assignment from cost matrix.
+
+        Args:
+            cost_matrix: Cost matrix C of shape (P, V).
+            a: Source marginal of shape (P,). Defaults to uniform 1/P.
+            b: Target marginal (accepted but ignored).
+            init_f: Warm-start dual potential (accepted but ignored).
+            init_g: Warm-start dual potential (accepted but ignored).
+            scale_cost: Cost scaling: 'mean', 'max_cost', or float divisor.
+
+        Returns:
+            SolverResult with sparse transport matrix (one non-zero per row).
+        """
+        P, V = cost_matrix.shape
+        device = cost_matrix.device
+        dtype = cost_matrix.dtype
+
+        if a is None:
+            a = torch.ones(P, device=device, dtype=dtype) / P
+
+        C = scale_cost_matrix(cost_matrix.clone(), scale_cost)
+
+        # Greedy: each particle picks the single lowest-cost vertex
+        min_indices = C.argmin(dim=-1)  # (P,)
+        transport = torch.zeros(P, V, device=device, dtype=dtype)
+        transport.scatter_(1, min_indices.unsqueeze(1), a.unsqueeze(1))
+
+        ent_cost = (C * transport).sum().item()
+
+        return SolverResult(
+            matrix=transport,
+            cost=ent_cost,
+            f=None,
+            g=None,
+            converged=True,
+            n_iters=1,
+            ent_reg_cost=ent_cost,
+        )
+
+
+@dataclass
+class TopKMeanSolver:
+    """Top-K mean solver: each particle moves to the uniform average of its k best vertices.
+
+    For each row i, finds the k vertices with lowest cost and assigns
+    equal mass ``a[i] / k`` to each: ``T[i,v] = a[i] / k`` for the top-k
+    vertices, zero elsewhere.
+
+    When V < k, gracefully falls back to using all V vertices.
+
+    Attributes:
+        epsilon: Accepted for API compatibility; unused.
+        compile: Accepted for API compatibility; unused.
+        k: Number of lowest-cost vertices to average over. Default 3.
+    """
+
+    epsilon: float = 0.1
+    compile: bool = False
+    k: int = 3
+
+    def solve(
+        self,
+        cost_matrix: torch.Tensor,
+        a: Optional[torch.Tensor] = None,
+        b: Optional[torch.Tensor] = None,
+        init_f: Optional[torch.Tensor] = None,
+        init_g: Optional[torch.Tensor] = None,
+        scale_cost: Optional[Union[str, float]] = None,
+    ) -> SolverResult:
+        """Compute top-k uniform assignment from cost matrix.
+
+        Args:
+            cost_matrix: Cost matrix C of shape (P, V).
+            a: Source marginal of shape (P,). Defaults to uniform 1/P.
+            b: Target marginal (accepted but ignored).
+            init_f: Warm-start dual potential (accepted but ignored).
+            init_g: Warm-start dual potential (accepted but ignored).
+            scale_cost: Cost scaling: 'mean', 'max_cost', or float divisor.
+
+        Returns:
+            SolverResult with transport matrix (k non-zeros per row).
+        """
+        P, V = cost_matrix.shape
+        device = cost_matrix.device
+        dtype = cost_matrix.dtype
+
+        if a is None:
+            a = torch.ones(P, device=device, dtype=dtype) / P
+
+        C = scale_cost_matrix(cost_matrix.clone(), scale_cost)
+
+        # Graceful fallback when fewer vertices than k
+        k_eff = min(self.k, V)
+
+        # Top-k: find k lowest-cost vertex indices per particle
+        _, topk_indices = C.topk(k_eff, dim=-1, largest=False)  # (P, k_eff)
+
+        # Uniform mass over top-k vertices
+        transport = torch.zeros(P, V, device=device, dtype=dtype)
+        mass_per_vertex = a.unsqueeze(1) / k_eff  # (P, 1)
+        transport.scatter_(1, topk_indices, mass_per_vertex.expand_as(topk_indices))
+
+        ent_cost = (C * transport).sum().item()
+
+        return SolverResult(
+            matrix=transport,
+            cost=ent_cost,
+            f=None,
+            g=None,
+            converged=True,
+            n_iters=1,
+            ent_reg_cost=ent_cost,
+        )

polystep/solvers/kl_softmax.py

@@ -0,0 +1,200 @@
+"""KL-penalized one-sided OT solver (interpolates softmax ↔ Sinkhorn).
+
+Implements the soft-target-marginal formulation:
+
+    min_P  <C, P> + epsilon * H(P) + lam * KL(P^T 1 || b)
+    s.t.   P 1 = a
+
+with `lam ∈ [0, ∞]`. The two limits are exact:
+
+- `lam = 0`  ≡ ``SoftmaxSolver`` (only row marginal enforced).
+- `lam -> ∞` ≡ ``SinkhornSolver`` (both row and column marginals).
+
+Algorithm (log-domain alternating updates):
+
+    α = lam / (lam + epsilon)            ∈ [0, 1]
+    f_i = epsilon * (log a_i - LSE_j((g_j - C_ij) / epsilon))   # exact
+    g_j = α * epsilon * (log b_j - LSE_i((f_i - C_ij) / epsilon))   # soft
+
+Setting α = 0 freezes g at zero, recovering ``SoftmaxSolver`` (one
+iteration suffices). Setting α = 1 (lam = ∞) recovers standard
+Sinkhorn alternating projections. Intermediate α produces a smooth
+interpolation, with `KL(P^T 1 || b)` decreasing monotonically as α
+grows.
+
+The α-scaling matches the scaling-algorithm form for unbalanced OT in
+Chizat, Peyré, Schmitzer & Vialard, *Scaling Algorithms for Unbalanced
+Optimal Transport Problems*, Math. Comp. 87 (2018), arXiv:1607.05816.
+We also support `lam = inf` explicitly (the user-facing default for
+"go to full Sinkhorn") so downstream code can pass `float('inf')`
+without arithmetic on infinity.
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+import torch
+
+from ..costs import scale_cost_matrix
+from ._prelude import align_marginal, recenter_cost, sanitize_cost
+from .base import SolverResult
+
+
+@dataclass
+class KLSoftmaxSolver:
+    """KL-penalized one-sided entropic OT solver.
+
+    Attributes:
+        epsilon: Entropic regularization (temperature). Must be > 0.
+        lam: KL penalty weight on the column marginal.
+            `0` reduces to ``SoftmaxSolver``; `inf` reduces to
+            ``SinkhornSolver``. Must be >= 0.
+        max_iterations: Maximum dual-update iterations.
+        threshold: Convergence tolerance on max(|Δf|, |Δg|).
+        compile: Placeholder for API compatibility (unused).
+    """
+
+    epsilon: float = 0.1
+    lam: float = float("inf")
+    max_iterations: int = 2000
+    threshold: float = 1e-6
+    compile: bool = False
+
+    # KL(P^T 1 || b) recorded on every solve() call; None until the first solve.
+    last_marginal_violation: Optional[float] = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        if self.epsilon <= 0:
+            raise ValueError(
+                f"epsilon must be > 0, got {self.epsilon}. epsilon is the entropic temperature and must be positive."
+            )
+        if self.lam < 0:
+            raise ValueError(f"lam must be >= 0, got {self.lam}. lam is the KL penalty on the column marginal.")
+        if self.threshold < 0:
+            raise ValueError(f"threshold must be >= 0, got {self.threshold}")
+        if self.max_iterations < 1:
+            raise ValueError(f"max_iterations must be >= 1, got {self.max_iterations}")
+
+    @property
+    def alpha(self) -> float:
+        """The KL-scaling coefficient α = lam / (lam + epsilon) ∈ [0, 1]."""
+        if self.lam == 0.0:
+            return 0.0
+        if math.isinf(self.lam):
+            return 1.0
+        return float(self.lam / (self.lam + self.epsilon))
+
+    def solve(
+        self,
+        cost_matrix: torch.Tensor,
+        a: Optional[torch.Tensor] = None,
+        b: Optional[torch.Tensor] = None,
+        init_f: Optional[torch.Tensor] = None,
+        init_g: Optional[torch.Tensor] = None,
+        scale_cost: Optional[Union[str, float]] = None,
+    ) -> SolverResult:
+        # Shared prelude: FP32 promotion for LSE stability plus device-side
+        # non-finite handling, then device/dtype-aligned marginals. Sanitizing
+        # before scaling keeps a +Inf penalty out of the 'mean'/'max_cost' scale.
+        cost_matrix = sanitize_cost(cost_matrix)
+        n, m = cost_matrix.shape
+        device, dtype = cost_matrix.device, cost_matrix.dtype
+        a = align_marginal(a, n, device, dtype, "a")
+        b = align_marginal(b, m, device, dtype, "b")
+
+        # Optional cost rescaling (on the already-sanitized, finite cost).
+        C = scale_cost_matrix(cost_matrix, scale_cost)
+        # Recenter to min(C)=0 so the log-sum-exp updates and the exp((f+g-C)/eps)
+        # plan keep FP32 precision when |C| is much larger than eps.
+        C, cost_shift = recenter_cost(C)
+
+        eps = float(self.epsilon)
+        alpha = self.alpha
+
+        log_a = a.clamp(min=1e-30).log()
+        log_b = b.clamp(min=1e-30).log()
+
+        # Disable any outer mixed-precision autocast inside the iteration -
+        # downcast LSE to BF16 collapses the dual potentials.
+        with torch.amp.autocast("cuda", enabled=False), torch.amp.autocast("cpu", enabled=False):
+            f = (
+                init_f.to(dtype=dtype, device=device).clone()
+                if init_f is not None
+                else torch.zeros(n, device=device, dtype=dtype)
+            )
+            g = (
+                init_g.to(dtype=dtype, device=device).clone()
+                if init_g is not None
+                else torch.zeros(m, device=device, dtype=dtype)
+            )
+
+            # Special-case alpha == 0 (softmax limit): single closed-form.
+            if alpha == 0.0:
+                f = eps * (log_a - torch.logsumexp(-C / eps, dim=1))
+                g = torch.zeros_like(g)
+                converged = True
+                n_iters = 1
+            else:
+                converged = False
+                n_iters = self.max_iterations
+                # Only sync the convergence flag once per ``check_every``
+                # iterations to keep the dual updates GPU-resident.
+                check_every = max(1, self.max_iterations // 20)
+                threshold = float(self.threshold)
+                for it in range(self.max_iterations):
+                    # f-update: exact row-marginal enforcement.
+                    f_new = eps * (log_a - torch.logsumexp((g.unsqueeze(0) - C) / eps, dim=1))
+                    # g-update: α-fraction of full-Sinkhorn target.
+                    g_target = eps * (log_b - torch.logsumexp((f_new.unsqueeze(1) - C) / eps, dim=0))
+                    g_new = alpha * g_target
+
+                    if (it + 1) % check_every == 0 or it == self.max_iterations - 1:
+                        delta = torch.maximum(
+                            (f_new - f).abs().amax(),
+                            (g_new - g).abs().amax(),
+                        )
+                        if delta.item() < threshold:
+                            f, g = f_new, g_new
+                            converged = True
+                            n_iters = it + 1
+                            break
+                    f, g = f_new, g_new
+
+            log_P = (f.unsqueeze(1) + g.unsqueeze(0) - C) / eps
+            P = log_P.exp()
+
+            # Numerical hygiene
+            if not torch.isfinite(P).all():
+                warnings.warn(
+                    "KLSoftmaxSolver produced non-finite transport entries; consider raising epsilon or lowering lam.",
+                    stacklevel=2,
+                )
+                P = torch.where(torch.isfinite(P), P, torch.zeros_like(P))
+
+            # Undo the recenter shift so cost is <C_scaled, P> (sum(P) == a.sum()).
+            cost = ((C * P).sum() + cost_shift * a.sum()).item()
+
+            # Theorem 4.1 instrumentation: KL(P^T 1 || b).
+            # P^T 1 is the realized column marginal; b is the target.
+            # KL = sum_j q_j * log(q_j / b_j) with q = P^T 1.
+            try:
+                col_marginal = P.sum(dim=0).clamp(min=1e-30)
+                b_safe = b.clamp(min=1e-30)
+                kl = (col_marginal * (col_marginal.log() - b_safe.log())).sum().item()
+                self.last_marginal_violation = float(kl)
+            except Exception:  # noqa: BLE001
+                self.last_marginal_violation = float("nan")
+
+        return SolverResult(
+            matrix=P,
+            cost=cost,
+            f=f,
+            g=g,
+            converged=converged,
+            n_iters=n_iters,
+            ent_reg_cost=cost,
+        )