'Sinkhorn-Knopp' optimal transport solver

Compute an entropy-regularized optimal transport plan using the 'Sinkhorn-Knopp' algorithm. Unlike other LAP solvers that return a hard 1-to-1 assignment, this returns a soft assignment (doubly stochastic matrix).

Usage

sinkhorn(
  cost,
  lambda = 10,
  tol = 1e-09,
  max_iter = 1000,
  r_weights = NULL,
  c_weights = NULL
)

Arguments

cost: Numeric matrix of transport costs. NA or Inf entries are treated as very high cost (effectively forbidden).
lambda: Regularization parameter (default 10). Higher values produce sharper (more deterministic) transport plans; lower values produce smoother distributions. Typical range: 1-100.
tol: Convergence tolerance (default 1e-9).
max_iter: Maximum iterations (default 1000).
r_weights: Optional numeric vector of row marginals (source distribution). Default is uniform. Will be normalized to sum to 1.
c_weights: Optional numeric vector of column marginals (target distribution). Default is uniform. Will be normalized to sum to 1.

Value

A list with elements:

transport_plan — numeric matrix, the optimal transport plan P. Row sums approximate r_weights, column sums approximate c_weights.
cost — the transport cost <C, P> (without entropy term).
u, v — scaling vectors (P = diag(u) * K * diag(v) where K = exp(-lambda*C)).
converged — logical, whether the algorithm converged.
iterations — number of iterations used.
lambda — the regularization parameter used.

Details

The 'Sinkhorn-Knopp' algorithm solves the entropy-regularized optimal transport problem:

$$P^* = \arg\min_P \langle C, P \rangle - \frac{1}{\lambda} H(P)$$

subject to row sums = r_weights and column sums = c_weights.

The entropy term H(P) encourages spread in the transport plan. As lambda -> Inf, the solution approaches the standard (unregularized) optimal transport.

Key differences from standard LAP solvers:

Returns a soft assignment (probabilities) not a hard 1-to-1 matching
Supports unequal marginals (weighted distributions)
Differentiable, making it useful in ML pipelines
Very fast: O(n^2) per iteration with typically O(1/tol^2) iterations

Use sinkhorn_to_assignment() to round the soft assignment to a hard matching.

References

Cuturi, M. (2013). 'Sinkhorn Distances': Lightspeed Computation of Optimal Transport. Advances in Neural Information Processing Systems, 26.

Examples

cost <- matrix(c(1, 2, 3, 4, 5, 6, 7, 8, 9), nrow = 3, byrow = TRUE)

# Soft assignment with default parameters
result <- sinkhorn(cost)
print(round(result$transport_plan, 3))

# Sharper assignment (higher lambda)
result_sharp <- sinkhorn(cost, lambda = 50)
print(round(result_sharp$transport_plan, 3))

# With custom marginals (more mass from row 1)
result_weighted <- sinkhorn(cost, r_weights = c(0.5, 0.25, 0.25))
print(round(result_weighted$transport_plan, 3))

# Round to hard assignment
hard_match <- sinkhorn_to_assignment(result)
print(hard_match)