ott.solvers.linear

ott.solvers.linear#

Linear solvers are the bread-and-butter of OT solvers. They can be called on their own, either the Sinkhorn Sinkhorn or Low-Rank LRSinkhorn solvers, to match two datasets. They also appear as subroutines for more advanced solvers in the ott.solvers module, notably ott.solvers.quadratic.

Sinkhorn Solvers#

solve(geom[, a, b, tau_a, tau_b, rank])

Solve linear regularized OT problem using Sinkhorn iterations.

solve_semidiscrete(geom[, b, rng])

Solve a (regularized) semidiscrete OT problem.

sinkhorn.Sinkhorn([lse_mode, threshold, ...])

Sinkhorn solver.

sinkhorn.SinkhornState(potentials[, errors, ...])

Holds the state variables used to solve OT with Sinkhorn.

sinkhorn.SinkhornOutput(potentials[, ...])

Holds the output of a Sinkhorn solver applied to a problem.

sinkhorn_lr.LRSinkhorn(rank[, gamma, ...])

Low-Rank Sinkhorn solver for linear reg-OT problems.

sinkhorn_lr.LRSinkhornState(q, r, g, gamma, ...)

State of the Low Rank Sinkhorn algorithm.

sinkhorn_lr.LRSinkhornOutput(q, r, g, costs, ...)

Transport interface for a low-rank Sinkhorn solution.

Barycenter Solvers#

continuous_barycenter.FreeWassersteinBarycenter(...)

Continuous Wasserstein barycenter solver [Cuturi and Doucet, 2014].

continuous_barycenter.FreeBarycenterState(x, a)

Holds the state of the Wasserstein barycenter solver.

discrete_barycenter.FixedBarycenter([...])

A Wasserstein barycenter solver for histograms on a common geometry.

discrete_barycenter.SinkhornBarycenterOutput(f, ...)

Univariate Solvers#

solve_univariate(geom[, a, b, ...])

Solve 1D OT problems between two \(d\)-dimensional point clouds.

univariate.uniform_solver(prob[, ...])

Univariate solver between two equally sized and uniformly weighted distributions.

univariate.quantile_solver(prob[, ...])

Univariate solver between quantile functions of distributions.

univariate.north_west_solver(prob)

Univariate solver that implements the north-west corner rule.

univariate.UnivariateOutput(prob, ot_costs)

Output of the univariate solver.

Semidiscrete Solvers#

semidiscrete.SemidiscreteSolver(*, ...[, ...])

Semidiscrete optimal transport solver.

semidiscrete.SemidiscreteState(it, epsilon, ...)

State of the SemidiscreteSolver.

semidiscrete.SemidiscreteOutput(g, prob[, ...])

Output of the SemidiscreteSolver.

semidiscrete.HardAssignmentOutput(ot_prob, ...)

Unregularized linear OT solution.

semidiscrete.constant_epsilon_scheduler(...)

Constant epsilon scheduler.

Sinkhorn Acceleration#

acceleration.Momentum([start, ...])

Momentum for Sinkhorn updates.

acceleration.AndersonAcceleration([memory, ...])

Implements Anderson acceleration for Sinkhorn.

Implicit Differentiation#

implicit_differentiation.ImplicitDiff([...])

Implicit differentiation of the Sinkhorn algorithm.

implicit_differentiation.solve_jax_cg(lin, b)

Wrapper around JAX native linear solvers.

lineax_implicit.solve_lineax(lin, b[, ...])

Solve a linear system using conjugate gradients.

Low-rank Sinkhorn Utilities#

lr_utils.unbalanced_dykstra_lse(c_q, c_r, ...)

Dykstra's algorithm for the unbalanced LRSinkhorn in LSE mode.

lr_utils.unbalanced_dykstra_kernel(k_q, k_r, ...)

Dykstra's algorithm for the unbalanced LRSinkhorn in kernel mode.