Entropic Regularized Transport

def ot.bregman.sinkhorn(a, b, M, reg, method='sinkhorn', numItermax=1000, stopThr=1e-9, verbose=False, log=False, warn=True, warmstart=None, **kwargs):
    """
    Solve entropic regularized optimal transport using Sinkhorn algorithm.
    
    The Sinkhorn algorithm iteratively projects onto marginal constraints to find
    the optimal transport plan for the entropy-regularized problem:
    min <P,M> + reg * KL(P|K) subject to P1=a, P^T1=b
    where K = exp(-M/reg).
    
    Parameters:
    - a: array-like, shape (n_samples_source,)
         Source distribution (histogram). Must be positive and sum to 1.
    - b: array-like, shape (n_samples_target,)
         Target distribution (histogram). Must be positive and sum to 1.
    - M: array-like, shape (n_samples_source, n_samples_target)
         Ground cost matrix.
    - reg: float
         Regularization parameter (>0). Lower values give solutions closer to EMD.
    - method: str, default='sinkhorn'
         Algorithm variant. Options: 'sinkhorn', 'sinkhorn_log', 'sinkhorn_stabilized',
         'sinkhorn_epsilon_scaling', 'greenkhorn', 'screenkhorn'
    - numItermax: int, default=1000
         Maximum number of iterations.
    - stopThr: float, default=1e-9
         Convergence threshold on marginal difference.
    - verbose: bool, default=False
         Print iteration information.
    - log: bool, default=False
         Return optimization log with convergence details.
    - warn: bool, default=True
         Warn if algorithm doesn't converge.
    - warmstart: tuple, default=None
         Tuple (u, v) of dual variables for warm start initialization.
    
    Returns:
    - transport_plan: ndarray, shape (n_samples_source, n_samples_target)
         Entropic optimal transport plan.
    - log: dict (if log=True)
         Contains 'err': convergence errors, 'niter': iterations used,
         'u': source scaling, 'v': target scaling.
    """

def ot.bregman.sinkhorn2(a, b, M, reg, method='sinkhorn', numItermax=1000, stopThr=1e-9, verbose=False, log=False, warn=True, **kwargs):
    """
    Solve entropic regularized OT and return transport cost only.
    
    More efficient than sinkhorn() when only the optimal value is needed.
    
    Parameters: Same as sinkhorn()
    
    Returns:
    - cost: float
         Entropic regularized transport cost.
    - log: dict (if log=True)
    """

def ot.bregman.sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9, verbose=False, log=False, **kwargs):
    """
    Sinkhorn-Knopp algorithm for entropic optimal transport.
    
    Classic formulation using multiplicative updates with diagonal scaling matrices.
    
    Parameters: Similar to sinkhorn()
    
    Returns:
    - transport_plan: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.sinkhorn_log(a, b, M, reg, numItermax=1000, stopThr=1e-9, verbose=False, log=False, **kwargs):
    """
    Sinkhorn algorithm in log-domain for numerical stability.
    
    Performs computations in log space to avoid numerical overflow/underflow
    issues when regularization parameter is small or cost matrix has large values.
    
    Parameters: Same as sinkhorn()
    
    Returns:
    - transport_plan: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.sinkhorn_stabilized(a, b, M, reg, numItermax=1000, stopThr=1e-9, verbose=False, log=False, tau=1e3, **kwargs):
    """
    Stabilized Sinkhorn algorithm with absorption technique.
    
    Uses tau-absorption to prevent numerical overflow while maintaining
    precision. Automatically switches between normal and log computations.
    
    Parameters:
    - Additional parameter:
    - tau: float, default=1e3
         Absorption threshold. When scaling factors exceed tau, algorithm
         absorbs them into the dual variables.
    
    Returns:
    - transport_plan: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.sinkhorn_epsilon_scaling(a, b, M, reg, numItermax=100, epsilon0=1e4, numInnerItermax=100, stopThr=1e-9, verbose=False, log=False, **kwargs):
    """
    Epsilon-scaling Sinkhorn for better convergence with small regularization.
    
    Starts with large regularization parameter and progressively decreases it
    to the target value, using warm-start between scales.
    
    Parameters:
    - epsilon0: float, default=1e4
         Initial (large) regularization parameter.
    - numInnerItermax: int, default=100
         Maximum iterations per epsilon scale.
    - Other parameters same as sinkhorn()
    
    Returns:
    - transport_plan: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.greenkhorn(a, b, M, reg, numItermax=10000, stopThr=1e-9, verbose=False, log=False):
    """
    Greenkhorn algorithm for sparse optimal transport.
    
    Coordinate-wise variant of Sinkhorn that updates one row/column at a time,
    leading to sparse solutions suitable for large-scale problems.
    
    Parameters: Same as sinkhorn() with typically larger numItermax
    
    Returns:
    - transport_plan: ndarray (often sparse)
    - log: dict (if log=True)
    """

def ot.bregman.screenkhorn(a, b, M, reg, ns_budget=None, nt_budget=None, uniform=False, restricted=True, maxiter=10000, maxfun=10000, pgtol=1e-09, verbose=False, log=False):
    """
    Screenkhorn algorithm for large-scale optimal transport.
    
    Uses screening techniques to identify and ignore negligible entries in the
    transport matrix, significantly reducing computational cost for large problems.
    
    Parameters:
    - ns_budget: int, optional
         Maximum number of active source samples.
    - nt_budget: int, optional
         Maximum number of active target samples.
    - uniform: bool, default=False
         Use uniform sampling for screening.
    - restricted: bool, default=True
         Use restricted Sinkhorn on screened samples.
    - maxiter: int, default=10000
    - maxfun: int, default=10000
         Maximum function evaluations.
    - pgtol: float, default=1e-09
         Projected gradient tolerance.
    
    Returns:
    - transport_plan: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.barycenter(A, M, reg, weights=None, method="sinkhorn", numItermax=10000, stopThr=1e-4, verbose=False, log=False, **kwargs):
    """
    Compute Wasserstein barycenter using entropic regularization.
    
    Solves the multi-marginal optimal transport problem to find the barycenter
    that minimizes the sum of regularized transport costs to all input distributions.
    
    Parameters:
    - A: array-like, shape (n_samples, n_distributions)
         Input distributions as columns of matrix A.
    - M: array-like, shape (n_samples, n_samples)
         Ground cost matrix on barycenter support.
    - reg: float
         Entropic regularization parameter.
    - weights: array-like, shape (n_distributions,), optional
         Weights for barycenter combination. Default is uniform.
    - method: str, default="sinkhorn"
         Algorithm to use for transport computation.
    - numItermax: int, default=10000
         Maximum iterations for barycenter computation.
    - stopThr: float, default=1e-4
         Convergence threshold.
    - verbose: bool, default=False
    - log: bool, default=False
    
    Returns:
    - barycenter: ndarray, shape (n_samples,)
         Wasserstein barycenter distribution.
    - log: dict (if log=True)
         Contains convergence information and transport plans.
    """

def ot.bregman.barycenter_sinkhorn(A, M, reg, weights=None, numItermax=1000, stopThr=1e-4, verbose=False, log=False):
    """
    Compute barycenter using Sinkhorn algorithm (alternative implementation).
    
    Parameters: Same as barycenter()
    
    Returns:
    - barycenter: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.barycenter_stabilized(A, M, reg, tau=1e3, weights=None, numItermax=1000, stopThr=1e-4, verbose=False, log=False):
    """
    Compute barycenter using stabilized Sinkhorn algorithm.
    
    Parameters:
    - tau: float, default=1e3
         Stabilization parameter.
    - Other parameters same as barycenter()
    
    Returns:
    - barycenter: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.barycenter_debiased(A, M, reg, weights=None, numItermax=1000, stopThr=1e-4, verbose=False, log=False):
    """
    Compute debiased Wasserstein barycenter.
    
    Applies debiasing correction to reduce bias introduced by entropic regularization.
    
    Parameters: Same as barycenter()
    
    Returns:
    - barycenter: ndarray
    - log: dict (if log=True)
    """

def ot.bregman.free_support_sinkhorn_barycenter(measures_locations, measures_weights, X_init, reg, b=None, weights=None, numItermax=100, stopThr=1e-7, verbose=False, log=False, **kwargs):
    """
    Compute free-support Wasserstein barycenter using Sinkhorn algorithm.
    
    Optimizes both barycenter weights and support locations simultaneously,
    unlike fixed-support methods that only optimize weights.
    
    Parameters:
    - measures_locations: list of arrays
         Support points for each input measure.
    - measures_weights: list of arrays
         Weights for each input measure.
    - X_init: array-like, shape (k, d)
         Initial barycenter support points.
    - reg: float
         Entropic regularization parameter.
    - b: array-like, shape (k,), optional
         Barycenter weights (optimized if None).
    - weights: array-like, optional
         Weights for combining input measures.
    - numItermax: int, default=100
    - stopThr: float, default=1e-7
    - verbose: bool, default=False
    - log: bool, default=False
    
    Returns:
    - X: ndarray, shape (k, d)
         Optimal barycenter support points.
    - b: ndarray, shape (k,)
         Optimal barycenter weights.
    - log: dict (if log=True)
    """

def ot.bregman.jcpot_barycenter(Xs, Ys, Ps, lambdas, reg, metric='sqeuclidean', numItermax=100, stopThr=1e-6, verbose=False, log=False, **kwargs):
    """
    Compute Joint Characteristic-Optimal-Transport (JCPOT) barycenter.
    
    Specialized barycenter for joint distributions, commonly used in
    domain adaptation scenarios.
    
    Parameters:
    - Xs: list of arrays
         Source feature matrices for each domain.
    - Ys: list of arrays
         Target feature matrices for each domain.
    - Ps: list of arrays
         Initial transport plans for each domain.
    - lambdas: array-like
         Weights for domain combination.
    - reg: float
         Entropic regularization.
    - metric: str, default='sqeuclidean'
         Ground metric for cost computation.
    - numItermax: int, default=100
    - stopThr: float, default=1e-6
    - verbose: bool, default=False
    - log: bool, default=False
    
    Returns:
    - X_barycenter: ndarray
         Barycenter in source space.
    - Y_barycenter: ndarray
         Barycenter in target space.
    - log: dict (if log=True)
    """

def ot.bregman.convolutional_barycenter2d(A, reg, weights=None, numItermax=10000, stopThr=1e-9, stabThr=1e-30, verbose=False, log=False):
    """
    Compute 2D convolutional Wasserstein barycenter.
    
    Specialized algorithm for 2D images using convolutional structure for
    efficiency. Exploits translation invariance of the ground metric.
    
    Parameters:
    - A: array-like, shape (h, w, n_images)
         Stack of 2D images/distributions.
    - reg: float
         Entropic regularization parameter.
    - weights: array-like, shape (n_images,), optional
         Barycenter weights.
    - numItermax: int, default=10000
    - stopThr: float, default=1e-9
    - stabThr: float, default=1e-30
         Numerical stability threshold.
    - verbose: bool, default=False
    - log: bool, default=False
    
    Returns:
    - barycenter: ndarray, shape (h, w)
         2D convolutional Wasserstein barycenter.
    - log: dict (if log=True)
    """

def ot.bregman.convolutional_barycenter2d_debiased(A, reg, weights=None, numItermax=10000, stopThr=1e-9, stabThr=1e-30, verbose=False, log=False):
    """
    Compute debiased 2D convolutional Wasserstein barycenter.
    
    Applies debiasing to reduce regularization bias in convolutional barycenters.
    
    Parameters: Same as convolutional_barycenter2d()
    
    Returns:
    - barycenter: ndarray, shape (h, w)
    - log: dict (if log=True)
    """

def ot.bregman.empirical_sinkhorn(X_s, X_t, reg, a=None, b=None, metric='sqeuclidean', numItermax=10000, stopThr=1e-9, verbose=False, log=False, **kwargs):
    """
    Compute Sinkhorn transport between empirical distributions.
    
    Convenient wrapper that computes cost matrix from sample coordinates
    and applies Sinkhorn algorithm.
    
    Parameters:
    - X_s: array-like, shape (n_samples_source, n_features)
         Source samples.
    - X_t: array-like, shape (n_samples_target, n_features)
         Target samples.
    - reg: float
         Entropic regularization parameter.
    - a: array-like, shape (n_samples_source,), optional
         Source sample weights. Default is uniform.
    - b: array-like, shape (n_samples_target,), optional
         Target sample weights. Default is uniform.
    - metric: str, default='sqeuclidean'
         Ground metric for cost matrix computation.
    - numItermax: int, default=10000
    - stopThr: float, default=1e-9
    - verbose: bool, default=False
    - log: bool, default=False
    
    Returns:
    - transport_plan: ndarray, shape (n_samples_source, n_samples_target)
    - log: dict (if log=True)
    """

def ot.bregman.empirical_sinkhorn2(X_s, X_t, reg, a=None, b=None, metric='sqeuclidean', numItermax=10000, stopThr=1e-9, verbose=False, log=False, **kwargs):
    """
    Compute empirical Sinkhorn transport cost only.
    
    Parameters: Same as empirical_sinkhorn()
    
    Returns:
    - cost: float
         Empirical Sinkhorn transport cost.
    - log: dict (if log=True)
    """

def ot.bregman.empirical_sinkhorn_divergence(X_s, X_t, reg, a=None, b=None, metric='sqeuclidean', numItermax=10000, stopThr=1e-9, verbose=False, log=False, **kwargs):
    """
    Compute empirical Sinkhorn divergence (debiased).
    
    Computes the Sinkhorn divergence: W_reg(a,b) - 0.5*W_reg(a,a) - 0.5*W_reg(b,b)
    which removes the regularization bias for better approximation of Wasserstein distance.
    
    Parameters: Same as empirical_sinkhorn()
    
    Returns:
    - divergence: float
         Sinkhorn divergence value.
    - log: dict (if log=True)
    """

def ot.bregman.empirical_sinkhorn2_geomloss(X_s, X_t, reg, a=None, b=None, metric='sqeuclidean', numItermax=10000, stopThr=1e-9, verbose=False, log=False):
    """
    GeomLoss-compatible implementation of empirical Sinkhorn.
    
    Parameters: Same as empirical_sinkhorn()
    
    Returns:
    - cost: float
    - log: dict (if log=True)
    """

def ot.bregman.geomloss(X_s, X_t, a=None, b=None, loss='sinkhorn', p=2, blur=0.1, reach=1.0, diameter=1.0, scaling=0.5, truncate=5, cost=None, kernel=None, cluster_scale=None, debias=True, potentials=False, verbose=False, backend='auto'):
    """
    GeomLoss wrapper for various optimal transport losses.
    
    Unified interface supporting Wasserstein, energy, Hausdorff and Sinkhorn losses
    with automatic differentiation support.
    
    Parameters:
    - X_s: array-like, shape (n_s, d)
         Source samples.
    - X_t: array-like, shape (n_t, d)
         Target samples.
    - a: array-like, shape (n_s,), optional
         Source weights.
    - b: array-like, shape (n_t,), optional
         Target weights.
    - loss: str, default='sinkhorn'
         Loss type: 'wasserstein', 'sinkhorn', 'energy', 'hausdorff'
    - p: int, default=2
         Ground metric exponent.
    - blur: float, default=0.1
         Regularization/smoothing parameter.
    - reach: float, default=1.0
         Kernel reach parameter.
    - diameter: float, default=1.0
         Point cloud diameter estimate.
    - scaling: float, default=0.5
         Multi-scale algorithm parameter.
    - truncate: int, default=5
         Kernel truncation parameter.
    - cost: callable, optional
         Custom cost function.
    - kernel: callable, optional
         Custom kernel function.
    - cluster_scale: float, optional
         Clustering scale for acceleration.
    - debias: bool, default=True
         Apply debiasing for Sinkhorn loss.
    - potentials: bool, default=False
         Return dual potentials.
    - verbose: bool, default=False
    - backend: str, default='auto'
         Computation backend.
    
    Returns:
    - loss_value: float
         Computed loss value.
    - potentials: tuple (if potentials=True)
         Dual potentials (f, g).
    """

def ot.bregman.geometricBar(weights, alldistribT):
    """
    Compute geometric barycenter in Bregman divergence sense.
    
    Parameters:
    - weights: array-like
         Barycenter combination weights.
    - alldistribT: array-like
         Matrix of input distributions (columns).
    
    Returns:
    - barycenter: ndarray
         Geometric barycenter.
    """

def ot.bregman.geometricMean(alldistribT):
    """
    Compute geometric mean of distributions.
    
    Parameters:
    - alldistribT: array-like
         Matrix of distributions (columns).
    
    Returns:
    - geometric_mean: ndarray
    """

def ot.bregman.projR(gamma, p):
    """
    Project transport matrix onto row constraints.
    
    Parameters:
    - gamma: array-like
         Transport matrix.
    - p: array-like
         Row marginal constraints.
    
    Returns:
    - projected_gamma: ndarray
    """

def ot.bregman.projC(gamma, q):
    """
    Project transport matrix onto column constraints.
    
    Parameters:
    - gamma: array-like
         Transport matrix.
    - q: array-like
         Column marginal constraints.
    
    Returns:
    - projected_gamma: ndarray
    """

def ot.bregman.unmix(a, D, M, M0, h0, reg, reg0, alpha, numItermax=1000, stopThr=1e-3, verbose=False, log=False):
    """
    Solve optimal transport unmixing problem with regularization.
    
    Decompose a distribution as a convex combination of dictionary atoms
    using optimal transport as the fidelity term.
    
    Parameters:
    - a: array-like
         Distribution to unmix.
    - D: array-like
         Dictionary of atoms (columns).
    - M: array-like
         Cost matrix for transport.
    - M0: array-like
         Cost matrix for dictionary regularization.
    - h0: array-like
         Prior on dictionary coefficients.
    - reg: float
         Transport regularization.
    - reg0: float
         Dictionary regularization.
    - alpha: float
         Fidelity vs regularization trade-off.
    - numItermax: int, default=1000
    - stopThr: float, default=1e-3
    - verbose: bool, default=False
    - log: bool, default=False
    
    Returns:
    - h: ndarray
         Dictionary coefficients.
    - log: dict (if log=True)
    """

import ot
import numpy as np

# Define distributions
a = np.array([0.5, 0.5])
b = np.array([0.3, 0.7])

# Cost matrix
M = np.array([[0.0, 1.0],
              [1.0, 0.0]])

# Regularization parameter
reg = 0.1

# Compute regularized transport
plan_sinkhorn = ot.bregman.sinkhorn(a, b, M, reg)
cost_sinkhorn = ot.bregman.sinkhorn2(a, b, M, reg)

print("Sinkhorn plan:", plan_sinkhorn)
print("Sinkhorn cost:", cost_sinkhorn)

# Multiple distributions
A = np.array([[0.6, 0.2, 0.4],
              [0.4, 0.8, 0.6]])  # 3 distributions

# Cost matrix  
M = ot.dist(np.arange(2).reshape(-1, 1))

# Regularization
reg = 0.05

# Compute barycenter
barycenter = ot.bregman.barycenter(A, M, reg)
print("Barycenter:", barycenter)

# With custom weights
weights = np.array([0.5, 0.3, 0.2])
weighted_barycenter = ot.bregman.barycenter(A, M, reg, weights=weights)
print("Weighted barycenter:", weighted_barycenter)

# Generate sample data
np.random.seed(42)
X_s = np.random.randn(100, 2)
X_t = np.random.randn(80, 2) + 1

# Regularization
reg = 0.1

# Compute empirical transport
plan = ot.bregman.empirical_sinkhorn(X_s, X_t, reg)
cost = ot.bregman.empirical_sinkhorn2(X_s, X_t, reg)

print("Empirical transport cost:", cost)
print("Transport plan shape:", plan.shape)

# Sinkhorn divergence (debiased)
divergence = ot.bregman.empirical_sinkhorn_divergence(X_s, X_t, reg)
print("Sinkhorn divergence:", divergence)

# Small regularization parameter
reg_small = 1e-3

# Use stabilized version to avoid numerical issues
plan_stable = ot.bregman.sinkhorn_stabilized(a, b, M, reg_small, verbose=True)
print("Stabilized Sinkhorn plan:", plan_stable)

# Or use epsilon scaling
plan_eps = ot.bregman.sinkhorn_epsilon_scaling(a, b, M, reg_small, verbose=True)
print("Epsilon-scaled plan:", plan_eps)