tessl/pypi-pot

Python Optimal Transport Library providing solvers for optimization problems related to Optimal Transport for signal, image processing and machine learning

—

Pending

Quality

Pending

Does it follow best practices?

Impact

Pending

No eval scenarios have been run

Overview

Eval results

Files

Stochastic Solvers

Name: tessl/pypi-pot
Author: tessl

Stochastic algorithms for large-scale optimal transport problems using Stochastic Average Gradient (SAG) and Stochastic Gradient Descent (SGD) methods. These solvers are particularly effective for problems with large numbers of samples where exact methods become computationally prohibitive.

Capabilities

SAG-based Entropic Transport

Stochastic Average Gradient methods for solving regularized optimal transport problems.

def sag_entropic_transport(a, b, M, reg, numItermax=10000, lr=None, random_state=None):
    """
    Solve entropic regularized OT with Stochastic Average Gradient algorithm.
    
    Parameters:
    - a: array-like, source distribution
    - b: array-like, target distribution
    - M: array-like, cost matrix
    - reg: float, regularization parameter
    - numItermax: int, maximum number of iterations
    - lr: float, learning rate (auto-computed if None)
    - random_state: int, random seed for reproducibility
    
    Returns:
    - transport plan matrix
    """

def averaged_sgd_entropic_transport(a, b, M, reg, numItermax=300000, lr=1e-4, random_state=None):
    """
    Solve entropic regularized OT with averaged SGD.
    
    Parameters:
    - a: array-like, source distribution
    - b: array-like, target distribution
    - M: array-like, cost matrix
    - reg: float, regularization parameter
    - numItermax: int, maximum number of iterations
    - lr: float, learning rate
    - random_state: int, random seed
    
    Returns:
    - transport plan matrix
    """

SGD-based Dual Solvers

Stochastic gradient descent methods for dual optimal transport formulations.

def sgd_entropic_regularization(a, b, M, reg, batch_size, numItermax=10000, lr=0.1, log=False):
    """
    Solve entropic regularized OT using SGD on dual formulation.
    
    Parameters:
    - a: array-like, source distribution
    - b: array-like, target distribution
    - M: array-like, cost matrix
    - reg: float, regularization parameter
    - batch_size: int, batch size for stochastic updates
    - numItermax: int, maximum iterations
    - lr: float, learning rate
    - log: bool, return optimization log
    
    Returns:
    - dual variables (u, v) or ((u, v), log) if log=True
    """

def solve_dual_entropic(a, b, M, reg, batch_size, numItermax=10000, lr=1, log=False):
    """
    Solve dual entropic regularized OT with SGD.
    
    Parameters:
    - a: array-like, source distribution
    - b: array-like, target distribution
    - M: array-like, cost matrix
    - reg: float, regularization parameter
    - batch_size: int, batch size for gradient computation
    - numItermax: int, maximum iterations
    - lr: float, learning rate
    - log: bool, return optimization log
    
    Returns:
    - dual variables (u, v) or ((u, v), log) if log=True
    """

Semi-Dual Entropic Solvers

Semi-dual formulation solvers for enhanced computational efficiency.

def solve_semi_dual_entropic(a, b, M, reg, method='SLSQP', numItermax=10000, lr=1, log=False):
    """
    Solve semi-dual entropic regularized optimal transport.
    
    Parameters:
    - a: array-like, source distribution
    - b: array-like, target distribution
    - M: array-like, cost matrix
    - reg: float, regularization parameter
    - method: str, optimization method ('SLSQP', 'L-BFGS-B')
    - numItermax: int, maximum iterations
    - lr: float, learning rate
    - log: bool, return optimization log
    
    Returns:
    - dual variable u or (u, log) if log=True
    """

Empirical Loss Functions

Loss and plan computation for empirical optimal transport with dual variables.

def loss_dual_entropic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidean"):
    """
    Compute dual entropic loss for empirical OT.
    
    Parameters:
    - u: array-like, dual variable for source samples
    - v: array-like, dual variable for target samples
    - xs: array-like, source samples
    - xt: array-like, target samples
    - reg: float, regularization parameter
    - ws: array-like, source weights (uniform if None)
    - wt: array-like, target weights (uniform if None)
    - metric: str, distance metric ('sqeuclidean', 'euclidean', etc.)
    
    Returns:
    - dual loss value
    """

def plan_dual_entropic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidean"):
    """
    Compute transport plan from dual variables for empirical OT.
    
    Parameters:
    - u: array-like, dual variable for source samples
    - v: array-like, dual variable for target samples
    - xs: array-like, source samples
    - xt: array-like, target samples
    - reg: float, regularization parameter
    - ws: array-like, source weights
    - wt: array-like, target weights
    - metric: str, distance metric
    
    Returns:
    - transport plan matrix
    """

def loss_dual_quadratic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidean"):
    """
    Compute quadratic dual loss for empirical OT.
    
    Parameters:
    - u: array-like, dual variable for source samples
    - v: array-like, dual variable for target samples
    - xs: array-like, source samples
    - xt: array-like, target samples
    - reg: float, regularization parameter
    - ws: array-like, source weights
    - wt: array-like, target weights
    - metric: str, distance metric
    
    Returns:
    - quadratic dual loss value
    """

def plan_dual_quadratic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidean"):
    """
    Compute transport plan from dual variables with quadratic regularization.
    
    Parameters:
    - Same as loss_dual_quadratic
    
    Returns:
    - transport plan matrix
    """

Transform Functions

Utility functions for coordinate transformations in stochastic algorithms.

def c_transform_entropic(b, M, reg, beta):
    """
    Compute c-transform for entropic regularization.
    
    Parameters:
    - b: array-like, target distribution
    - M: array-like, cost matrix
    - reg: float, regularization parameter
    - beta: array-like, dual variable
    
    Returns:
    - c-transform result
    """

Usage Example

import ot
import ot.stochastic
import numpy as np

# Generate large-scale data
n_samples = 10000
n_features = 50
X_s = np.random.randn(n_samples, n_features)
X_t = np.random.randn(n_samples, n_features)

# Create distributions
a = ot.utils.unif(n_samples)
b = ot.utils.unif(n_samples)
M = ot.dist(X_s, X_t)

# Solve with SAG algorithm
reg = 0.1
plan_sag = ot.stochastic.sag_entropic_transport(a, b, M, reg, numItermax=1000)

# Solve with SGD on dual formulation
batch_size = 100
u, v = ot.stochastic.sgd_entropic_regularization(
    a, b, M, reg, batch_size, numItermax=5000, lr=0.01
)

# Compute loss and plan from dual variables
loss = ot.stochastic.loss_dual_entropic(u, v, X_s, X_t, reg)
plan = ot.stochastic.plan_dual_entropic(u, v, X_s, X_t, reg)

Import Statements

import ot.stochastic
from ot.stochastic import sag_entropic_transport, sgd_entropic_regularization
from ot.stochastic import loss_dual_entropic, plan_dual_entropic

Install with Tessl CLI