ott.tools

ott.tools#

The tools package contains high level functions that build on outputs produced by lower-level components in the toolbox, such as solvers.

In particular, we provide user-friendly APIs to unregularized OT quantities, such as the Wasserstein distance for two point clouds of the same size. We also provide functions to pad efficiently point clouds when doing large scale OT between them in parallel, implementations of the Sinkhorn divergence [Genevay et al., 2018, Séjourné et al., 2019], sliced Wasserstein distances [Rabin et al., 2012], differentiable approximations to ranks and quantile functions [Cuturi et al., 2019], and various tools to study Gaussians with the 2-Wasserstein distance [Delon and Desolneux, 2020, Gelbrich, 1990].

Unregularized Optimal Transport#

unreg.hungarian(geom)

Solve matching problem using the Hungarian algorithm.

unreg.wassdis_p(x, y, *[, p])

Compute the Wasserstein distance, uses Hungarian algorithm.

Segmented Sinkhorn#

segment_sinkhorn.segment_sinkhorn(x, y[, ...])

Compute regularized OT cost between subsets of vectors in x and y.

Sinkhorn Divergence#

sinkhorn_divergence.sinkdiv(x, y, *[, ...])

Wrapper to get the Sinkhorn divergence between two point clouds.

sinkhorn_divergence.sinkhorn_divergence(...)

Compute Sinkhorn divergence between two measures.

sinkhorn_divergence.SinkhornDivergenceOutput(...)

Holds the outputs of a call to sinkhorn_divergence().

sinkhorn_divergence.segment_sinkhorn_divergence(x, y)

Compute Sinkhorn divergence between subsets of vectors in x and y.

Sliced Wasserstein Distance#

sliced.random_proj_sphere(x[, n_proj, rng])

Project data on directions sampled randomly from sphere.

sliced.sliced_wasserstein(x, y[, a, b, ...])

Compute the Sliced Wasserstein distance between two weighted point clouds.

ProgOT#

progot.ProgOT(alphas, *[, epsilons, ...])

Progressive Entropic Optimal Transport solver [Kassraie et al., 2024].

progot.ProgOTOutput(prob, alphas, epsilons, ...)

Output of the ProgOT solver.

progot.get_alpha_schedule(kind, *, num_steps)

Get the step size schedule.

progot.get_epsilon_schedule(geom, *, alphas, ...)

Get the epsilon regularization schedule.

Conformal Prediction#

conformal.OTCP(model[, nonconformity_fn, ...])

Optimal transport conformal prediction [Klein et al., 2025].

conformal.sobol_ball_sampler(rng, shape[, ...])

Sample target measure for OTCP.

Soft Sorting Algorithms#

soft_sort.multivariate_cdf_quantile_maps(inputs)

Returns multivariate CDF and quantile maps, given input samples.

soft_sort.quantile(inputs, q[, axis, weight])

Apply the soft quantiles operator on the input tensor.

soft_sort.quantile_normalization(inputs, targets)

Re-normalize inputs so that its quantiles match those of targets/weights.

soft_sort.quantize(inputs[, num_levels, axis])

Soft quantizes an input according using num_levels values along axis.

soft_sort.ranks(inputs[, axis, num_targets, ...])

Apply the soft rank operator on input tensor.

soft_sort.sort(inputs[, axis, topk, num_targets])

Apply the soft sort operator on a given axis of the input.

soft_sort.sort_with(inputs, criterion[, topk])

Sort a multidimensional array according to a real valued criterion.

soft_sort.topk_mask(inputs[, axis, k])

Soft \(\text{top-}k\) selection mask.

Clustering#

k_means.k_means(geom, k[, weights, init, ...])

K-means clustering using Lloyd's algorithm [Lloyd, 1982].

k_means.KMeansOutput(centroids, assignment, ...)

Output of the k_means() algorithm.

Plotting#

plot.Plot([fig, ax, threshold, scale, ...])

Plot an optimal transport map between two point clouds.

plot.PlotMM([fig, ax, fix_axes_lim, cmap, ...])

Plots an optimal transport map for MMSinkhorn.

plot.transport_animation(n_frames, ...[, ...])

Create animated visualizations of optimal transport and flow matching.

ott.tools.gaussian_mixture package#

This package implements various tools to manipulate Gaussian mixtures with a slightly modified Wasserstein geometry: here a Gaussian mixture is no longer strictly regarded as a density \(\mathbb{R}^d\), but instead as a point cloud in the space of Gaussians in \(\mathbb{R}^d\). This viewpoint provides a new approach to compare, and fit Gaussian mixtures, as described for instance in [Delon and Desolneux, 2020] and references therein.

Gaussian Mixtures#

gaussian.Gaussian(loc, scale)

Normal distribution.

gaussian_mixture.GaussianMixture(loc, ...)

Gaussian Mixture model.

gaussian_mixture_pair.GaussianMixturePair(...)

Coupled pair of Gaussian mixture models.

fit_gmm.initialize(rng, points, ...[, ...])

Initialize a GMM via K-means++ with retries on failure.

fit_gmm.fit_model_em(gmm, points, ...[, ...])

Fit a GMM using the EM algorithm.

fit_gmm_pair.get_fit_model_em_fn(...[, ...])

Get a function that performs penalized EM.