polystep 0.7.0__py3-none-any.whl

polystep/costs.py

@@ -0,0 +1,86 @@
+"""Cost matrix computation for Sinkhorn Step.
+
+Wraps an objective function into a cost matrix by evaluating it at
+probe points and averaging over the probe dimension.
+"""
+
+import math
+from typing import Callable, Optional, Union
+
+import torch
+
+
+def compute_cost_matrix(
+    objective_fn: Callable[[torch.Tensor], torch.Tensor],
+    X_probe: torch.Tensor,
+    chunk_size: Optional[int] = None,
+) -> torch.Tensor:
+    """Compute the cost matrix by evaluating objective at probe points.
+
+    Evaluates the objective function on all probe points and averages
+    over the probe dimension to get a (batch, num_vertices) cost matrix.
+
+    Args:
+        objective_fn: Function mapping (..., dim) -> (...) costs.
+        X_probe: Probe array of shape (batch, num_vertices, num_probe, dim).
+        chunk_size: If set, evaluate in chunks to limit memory.
+
+    Returns:
+        Cost matrix of shape (batch, num_vertices).
+    """
+    batch, num_verts, num_probe, dim = X_probe.shape
+
+    if chunk_size is not None and chunk_size > 0:
+        # Chunked evaluation for memory efficiency
+        flat_X = X_probe.reshape(-1, dim)
+        N = flat_X.shape[0]
+
+        results = []
+        for start in range(0, N, chunk_size):
+            end = min(start + chunk_size, N)
+            chunk = flat_X[start:end]
+            results.append(objective_fn(chunk))
+
+        raw_costs_flat = torch.cat(results, dim=0)
+    else:
+        flat_X = X_probe.reshape(-1, dim)
+        raw_costs_flat = objective_fn(flat_X)
+
+    raw_costs = raw_costs_flat.reshape(batch, num_verts, num_probe)
+    cost_matrix = raw_costs.mean(dim=-1)  # (batch, num_vertices)
+
+    return cost_matrix
+
+
+def scale_cost_matrix(
+    cost_matrix: torch.Tensor,
+    scale_cost: Optional[Union[str, float]] = None,
+) -> torch.Tensor:
+    """Apply cost scaling to a cost matrix.
+
+    Args:
+        cost_matrix: Raw cost matrix.
+        scale_cost: Scaling strategy ('mean', 'max_cost', or float).
+
+    Returns:
+        Scaled cost matrix.
+    """
+    if scale_cost is None:
+        return cost_matrix
+
+    if scale_cost == "mean":
+        s = torch.clamp(cost_matrix.abs().mean(), min=1e-10)
+        return cost_matrix / s
+    elif scale_cost == "max_cost":
+        s = torch.clamp(cost_matrix.abs().max(), min=1e-10)
+        return cost_matrix / s
+    elif isinstance(scale_cost, (int, float)):
+        s = float(scale_cost)
+        if not math.isfinite(s) or s <= 0.0:
+            raise ValueError(
+                f"Numeric scale_cost must be finite and positive, got {scale_cost!r}. "
+                "A negative divisor flips the cost sign and reverses the objective."
+            )
+        return cost_matrix / s
+    else:
+        raise ValueError(f"Unknown scale_cost: {scale_cost!r}. Expected 'mean', 'max_cost', or a float.")

polystep/dynamics.py

@@ -0,0 +1,138 @@
+"""Momentum and adaptive radius dynamics for Sinkhorn Step optimization.
+
+Pure functions that compute momentum coefficients, apply velocity updates,
+and adapt the step radius based on loss stagnation. These are composed by
+``PolyStepOptimizer`` (via ``use_momentum`` and ``use_adaptive_radius``
+parameters) but kept standalone for testability.
+
+**Momentum** accumulates a velocity vector across optimization steps, similar
+to SGD momentum but applied in the OT particle space. The OT barycentric
+projection gives a displacement; momentum blends this with accumulated
+velocity to smooth the trajectory and escape local minima.
+
+**Adaptive radius** automatically adjusts the step size based on optimization
+progress. When the loss stagnates (small relative change for several steps),
+the radius is increased to explore a wider region. When the loss is improving,
+the radius is decreased to exploit the current basin.
+"""
+
+import math
+from typing import Tuple
+
+import torch
+
+
+def compute_momentum_coefficient(
+    iteration: int,
+    max_iterations: int,
+    momentum_init: float = 0.5,
+    momentum_final: float = 0.95,
+) -> float:
+    """Linearly warm up momentum coefficient from init to final.
+
+    Args:
+        iteration: Current iteration index (0-based).
+        max_iterations: Total number of iterations.
+        momentum_init: Starting momentum value.
+        momentum_final: Final momentum value.
+
+    Returns:
+        Interpolated momentum coefficient.
+    """
+    progress = min(1.0, iteration / max(1, max_iterations - 1))
+    return momentum_init + progress * (momentum_final - momentum_init)
+
+
+@torch.inference_mode()
+def apply_momentum(
+    X_old: torch.Tensor,
+    X_barycentric: torch.Tensor,
+    velocity: torch.Tensor,
+    beta: float,
+    velocity_lr: float = 1.0,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply momentum update to particles.
+
+    Computes the displacement from the OT barycentric projection, blends it
+    with the accumulated velocity, and returns the updated positions and
+    velocity.
+
+    Args:
+        X_old: Particle positions before the OT step, shape (N, D).
+        X_barycentric: Particle positions from barycentric projection, shape (N, D).
+        velocity: Previous velocity tensor, shape (N, D).
+        beta: Momentum coefficient (0 = no momentum, 1 = full momentum).
+        velocity_lr: Learning rate applied to the velocity update.
+
+    Returns:
+        Tuple of (X_new, velocity_new) where X_new are the updated positions
+        and velocity_new is the updated velocity tensor.
+    """
+    displacement = X_barycentric - X_old
+    velocity_new = beta * velocity + displacement
+    X_new = X_old + velocity_lr * velocity_new
+    return X_new, velocity_new
+
+
+@torch.inference_mode()
+def update_adaptive_radius(
+    current_loss: float,
+    prev_loss: float,
+    stagnation_count: int,
+    radius_multiplier: float,
+    stagnation_threshold: float = 1e-4,
+    stagnation_patience: int = 10,
+    radius_increase: float = 1.5,
+    radius_decrease: float = 0.9,
+    radius_min: float = 0.5,
+    radius_max: float = 3.0,
+) -> Tuple[float, int, float]:
+    """Update adaptive radius based on loss progress.
+
+    Tracks stagnation (small relative loss change) and adjusts the radius
+    multiplier: increases on prolonged stagnation to explore more, decreases
+    on improvement to exploit.
+
+    Args:
+        current_loss: Loss at current iteration.
+        prev_loss: Loss at previous iteration.
+        stagnation_count: Current consecutive stagnation count.
+        radius_multiplier: Current radius multiplier.
+        stagnation_threshold: Relative change below which iteration is stagnating.
+        stagnation_patience: Stagnation iterations before radius boost.
+        radius_increase: Multiplicative factor for radius boost.
+        radius_decrease: Multiplicative factor for radius decay on improvement.
+        radius_min: Minimum allowed radius multiplier.
+        radius_max: Maximum allowed radius multiplier.
+
+    Returns:
+        Tuple of (radius_multiplier, stagnation_count, current_loss) where
+        current_loss is returned so the caller can store it as prev_loss.
+    """
+    # Guard against NaN/inf: if the loss is not finite, skip all radius
+    # adaptation and return the current state unchanged. This prevents a
+    # NaN loss from resetting the stagnation counter or corrupting the
+    # radius multiplier.
+    if not math.isfinite(current_loss):
+        return (radius_multiplier, stagnation_count, current_loss)
+
+    # Skip adaptation on the first step (prev_loss=inf) - no history to compare
+    if not math.isfinite(prev_loss):
+        return (radius_multiplier, stagnation_count, current_loss)
+
+    rel_change = abs(current_loss - prev_loss) / (abs(prev_loss) + 1e-10)
+
+    if rel_change < stagnation_threshold:
+        stagnation_count += 1
+    else:
+        stagnation_count = 0
+
+    if stagnation_count >= stagnation_patience:
+        radius_multiplier *= radius_increase
+        radius_multiplier = min(radius_multiplier, radius_max)
+        stagnation_count = 0
+    elif current_loss < prev_loss:
+        radius_multiplier *= radius_decrease
+        radius_multiplier = max(radius_multiplier, radius_min)
+
+    return (radius_multiplier, stagnation_count, current_loss)

polystep/epsilon.py

@@ -0,0 +1,217 @@
+"""Epsilon schedulers for entropic regularization decay.
+
+Entropic regularization (epsilon) controls the smoothness of the optimal
+transport plan. Higher epsilon gives a smoother, more diffuse plan (easier
+to solve but less precise). Lower epsilon gives a sharper plan (closer to
+exact OT but harder to solve numerically). Annealing from high to low
+epsilon allows coarse-to-fine optimization.
+
+Provides three schedulers:
+
+- ``LinearEpsilon``: Fixed linear decay: ``eps_t = max(init - decay * t, target)``.
+- ``CosineEpsilon``: Cosine annealing with optional SGDR-style warm restarts.
+- ``ProgressiveEpsilon``: Feedback-driven auto-adjustment based on Sinkhorn
+  solver convergence, inspired by ProgOT (Kassraie, Pooladian, Klein,
+  Thornton, Niles-Weed & Cuturi, *Progressive Entropic Optimal Transport
+  Solvers*, NeurIPS 2024, arXiv:2406.05061).
+"""
+
+import math
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class LinearEpsilon:
+    """Linearly decaying epsilon scheduler.
+
+    Computes epsilon at step t as:
+        epsilon(t) = max(init - decay * t, target)
+
+    The schedule starts at ``init`` and decreases by ``decay`` per iteration
+    until it reaches the ``target`` floor. This enables coarse-to-fine
+    optimization: early iterations use high epsilon for broad exploration,
+    later iterations use low epsilon for precise refinement.
+
+    Attributes:
+        target: Floor value for epsilon.
+        init: Starting value at iteration 0.
+        decay: Amount to subtract per iteration.
+    """
+
+    target: float = 1e-3
+    init: float = 1.0
+    decay: float = 0.01
+
+    def at(self, iteration: Optional[int] = 1) -> float:
+        """Compute epsilon at the given iteration.
+
+        Args:
+            iteration: Current iteration number.
+
+        Returns:
+            Epsilon value (float).
+        """
+        if iteration is None:
+            # Not yet started - return initial epsilon, not target
+            return self.init
+        eps = self.init - (self.decay * iteration)
+        return max(eps, self.target)
+
+
+@dataclass
+class ProgressiveEpsilon:
+    """Auto-adjusting epsilon based on Sinkhorn solver feedback.
+
+    Inspired by ProgOT (Kassraie et al., NeurIPS 2024, arXiv:2406.05061):
+    epsilon is adjusted from actual solver convergence behavior rather than
+    a fixed schedule.
+
+    When the solver converges quickly (few iterations), epsilon is decreased
+    to sharpen the transport plan. When the solver struggles (many iterations
+    or fails to converge), epsilon is increased to stabilize solving.
+
+    The ``at()`` method is a drop-in replacement for ``LinearEpsilon.at()``,
+    but ignores the ``iteration`` parameter. Instead, epsilon is driven by
+    explicit ``update()`` calls from the optimizer after each OT solve.
+
+    Attributes:
+        init: Starting epsilon value.
+        target: Floor value (minimum epsilon).
+        max_epsilon: Ceiling value (maximum epsilon).
+        increase_factor: Multiplicative increase when solver struggles.
+        decrease_factor: Multiplicative decrease when solver converges fast.
+        fast_threshold: n_iters below this fraction of max_iterations triggers decrease.
+        slow_threshold: n_iters above this fraction of max_iterations triggers increase.
+        ema_alpha: EMA smoothing for epsilon changes (0 = no smoothing, 1 = no change).
+    """
+
+    init: float = 1.0
+    target: float = 0.01
+    max_epsilon: float = 5.0
+    increase_factor: float = 1.2
+    decrease_factor: float = 0.95
+    fast_threshold: float = 0.1
+    slow_threshold: float = 0.5
+    ema_alpha: float = 0.7
+
+    # Internal state (not part of constructor signature for users)
+    _current: float = field(init=False, repr=False)
+    _smoothed: float = field(init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        self._current = self.init
+        self._smoothed = self.init
+
+    def at(self, iteration: Optional[int] = None) -> float:
+        """Return current epsilon value. Compatible with LinearEpsilon interface.
+
+        The ``iteration`` parameter is accepted for API compatibility with
+        ``LinearEpsilon`` but is ignored. ProgressiveEpsilon is driven by
+        explicit ``update()`` calls, not iteration count.
+
+        Args:
+            iteration: Ignored. Present for API compatibility.
+
+        Returns:
+            Current smoothed epsilon value (float).
+        """
+        return self._smoothed
+
+    def update(self, n_iters: int, max_iterations: int, converged: bool) -> None:
+        """Update epsilon based on solver feedback.
+
+        Args:
+            n_iters: Number of Sinkhorn iterations used in the last solve.
+            max_iterations: Maximum iterations the solver was allowed.
+            converged: Whether the solver converged.
+        """
+        ratio = n_iters / max(max_iterations, 1)
+
+        if not converged or ratio > self.slow_threshold:
+            # Solver struggling: increase epsilon
+            self._current = min(self._current * self.increase_factor, self.max_epsilon)
+        elif ratio < self.fast_threshold:
+            # Solver converging fast: decrease epsilon
+            self._current = max(self._current * self.decrease_factor, self.target)
+        # else: ratio in [fast_threshold, slow_threshold] -- keep current
+
+        # EMA smooth
+        self._smoothed = self.ema_alpha * self._smoothed + (1.0 - self.ema_alpha) * self._current
+        self._smoothed = max(self._smoothed, self.target)
+        self._smoothed = min(self._smoothed, self.max_epsilon)
+
+
+@dataclass
+class CosineEpsilon:
+    """Cosine-annealed epsilon scheduler.
+
+    Computes epsilon at step t using cosine annealing:
+        epsilon(t) = target + 0.5 * (init - target) * (1 + cos(pi * t / T))
+
+    where T = total_steps (computed from init/target/decay for API compatibility
+    with LinearEpsilon, or set explicitly via total_steps).
+
+    Cosine annealing keeps epsilon higher for longer in the middle of training
+    (encouraging exploration) then decays rapidly at the end (encouraging
+    exploitation). This typically improves convergence versus linear decay.
+
+    Supports optional warm restarts (SGDR-style): epsilon periodically resets
+    to init, with each period multiplied by ``restart_mult``.
+
+    Attributes:
+        target: Floor value for epsilon.
+        init: Starting value at iteration 0.
+        decay: Per-step linear decay rate (used to infer total_steps if
+            total_steps is not set: T = (init - target) / decay).
+        total_steps: Explicit total step count. Overrides decay-based inference.
+        restart_mult: Period multiplier for warm restarts (1.0 = no restarts).
+    """
+
+    target: float = 1e-3
+    init: float = 1.0
+    decay: float = 0.01
+    total_steps: int = 0
+    restart_mult: float = 1.0
+
+    def at(self, iteration: Optional[int] = 1) -> float:
+        """Compute epsilon at the given iteration.
+
+        Args:
+            iteration: Current iteration number.
+
+        Returns:
+            Epsilon value (float).
+        """
+        if iteration is None:
+            return self.init
+
+        T = (
+            self.total_steps
+            if self.total_steps > 0
+            else max(1, int((self.init - self.target) / max(self.decay, 1e-12)))
+        )
+
+        if self.restart_mult > 1.0:
+            # Warm restart: find which period we're in
+            period = T
+            t = iteration
+            # Guard: limit iterations to prevent unbounded loop when
+            # restart_mult is very close to 1.0 or period is tiny
+            max_restarts = 100
+            restarts = 0
+            while t >= period and period > 0 and restarts < max_restarts:
+                t -= period
+                period = int(period * self.restart_mult)
+                restarts += 1
+            T_local = max(period, 1)
+            t_local = t
+        else:
+            T_local = T
+            t_local = min(iteration, T)
+
+        # Clamp so a maxed-out restart loop can't push cos past pi (which would
+        # drift epsilon outside [target, init]).
+        t_local = min(max(t_local, 0), T_local)
+        cos_val = math.cos(math.pi * t_local / max(T_local, 1))
+        return self.target + 0.5 * (self.init - self.target) * (1.0 + cos_val)