ott.geometry.geometry.Geometry

class ott.geometry.geometry.Geometry(cost_matrix=None, kernel_matrix=None, epsilon=None, relative_epsilon=None, scale_epsilon=None, scale_cost=None, **kwargs)

Base class to define ground costs/kernels used in optimal transport.

Optimal transport problems are intrinsically geometric: they compute an optimal way to transport mass from one configuration onto another. To define what is meant by optimality of a transport requires defining a cost, of moving mass from one among several sources, towards one out of multiple targets. These sources and targets can be provided as points in vectors spaces, grids, or more generally exclusively described through a (dissimilarity) cost matrix, or almost equivalently, a (similarity) kernel matrix.

Once that cost or kernel matrix is set, the Geometry class provides a basic operations to be run with the Sinkhorn algorithm.

Parameters

• cost_matrix (Optional[ndarray]) – jnp.ndarray<float>[num_a, num_b]: a cost matrix storing n x m costs.

• kernel_matrix (Optional[ndarray]) – jnp.ndarray<float>[num_a, num_b]: a kernel matrix storing n x m kernel values.

• epsilon (Union[Epsilon, float, None]) – a regularization parameter. If a ott.geometry.epsilon_scheduler.Epsilon object is passed, other parameters below are ignored in practice. If the parameter is a float, then this is understood to be the regularization that is needed, unless relative_epsilon below is True, in which case epsilon is understood as a normalized quantity, to be scaled by the mean value of the cost_matrix.

• relative_epsilon (Optional[bool]) – whether epsilon is passed relative to scale of problem, here understood as mean value of cost_matrix.

• scale_epsilon (Optional[float]) – the scale multiplier for epsilon.

• scale_cost (Union[Literal[‘mean’, ‘max_cost’, ‘median’], bool, float, None]) – option to rescale the cost matrix. Implemented scalings are ‘median’, ‘mean’ and ‘max_cost’. Alternatively, a float factor can be given to rescale the cost such that cost_matrix /= scale_cost. If True, use ‘mean’.

• kwargs (Any) – additional kwargs to epsilon.

Note

When defining a Geometry through a cost_matrix, it is important to select an epsilon regularization parameter that is meaningful. That parameter can be provided by the user, or assigned a default value through a simple rule, using the mean cost value implied by the cost_matrix.

Methods

apply_cost(arr[, axis, fn])	Apply cost matrix to array (vector or matrix).
apply_kernel(scaling[, eps, axis])	Apply kernel on positive scaling vector.
apply_lse_kernel(f, g, eps[, vec, axis])	Apply kernel in log domain on pair of dual potential variables.
apply_square_cost(arr[, axis])	Apply elementwise-square of cost matrix to array (vector or matrix).
apply_transport_from_potentials(f, g, vec[, ...])	Apply transport matrix computed from potentials to a (batched) vec.
apply_transport_from_scalings(u, v, vec[, axis])	Apply transport matrix computed from scalings to a (batched) vec.
copy_epsilon(other)	Copy the epsilon parameters from another geometry.
marginal_from_potentials(f, g[, axis])	Output marginal of transportation matrix from potentials.
marginal_from_scalings(u, v[, axis])	Output marginal of transportation matrix from scalings.
potential_from_scaling(scaling)	Compute dual potential vector from scaling vector.
prepare_divergences(*args[, static_b])	Instantiate 2 (or 3) geometries to compute a Sinkhorn divergence.
rescale_cost_fn(factor)	Rescale the cost or kernel matrix using a factor.
scaling_from_potential(potential)	Compute scaling vector from dual potential.
transport_from_potentials(f, g)	Output transport matrix from potentials.
transport_from_scalings(u, v)	Output transport matrix from pair of scalings.
update_potential(f, g, log_marginal[, ...])	Carry out one Sinkhorn update for potentials, i.e. in log space.
update_scaling(scaling, marginal[, ...])	Carry out one Sinkhorn update for scalings, using kernel directly.

Attributes

cost_matrix	Cost matrix, recomputed from kernel if only kernel was specified.
cost_rank	Output rank of cost matrix, if any was provided.
epsilon	Epsilon regularization value.
inv_scale_cost	Compute and return inverse of scaling factor for cost matrix.
is_online	Whether geometry cost/kernel should be recomputed on the fly.
is_squared_euclidean	Whether cost is computed by taking squared-Eucl.
is_symmetric	Whether geometry cost/kernel is a symmetric matrix.
kernel_matrix	Kernel matrix, either provided by user or recomputed from cost.
mean_cost_matrix	Mean of cost matrix.
median_cost_matrix	Median of cost matrix.
scale_epsilon	Compute the scale of the epsilon, potentially based on data.
shape	Shape of cost or kernel matrix.