ott.geometry.geometry.Geometry#
- class ott.geometry.geometry.Geometry(cost_matrix=None, kernel_matrix=None, epsilon=None, relative_epsilon=None, scale_cost=1.0, src_mask=None, tgt_mask=None)[source]#
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 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
[Array
]) – Cost matrix of shape[n, m]
.kernel_matrix (
Optional
[Array
]) – Kernel matrix of shape[n, m]
.epsilon (
Union
[float
,Epsilon
,None
]) – Regularization parameter. IfNone
and eitherrelative_epsilon = True
orrelative_epsilon = None
, this defaults to the value computed inmean_cost_matrix
/ 20. If passed as afloat
, then the regularizer that is ultimately used is either thatfloat
value (ifrelative_epsilon = False
orNone
) or thatfloat
times themean_cost_matrix
(ifrelative_epsilon = True
). Look forEpsilon
when passed as a scheduler.relative_epsilon (
Optional
[bool
]) – when False, the parameterepsilon
specifies the value of the entropic regularization parameter. When True,epsilon
refers to a fraction of themean_cost_matrix
, which is computed adaptively from data.scale_cost (
Union
[bool
,int
,float
,Literal
['mean'
,'max_cost'
,'median'
]]) – 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 thatcost_matrix /= scale_cost
. If True, use ‘mean’.src_mask (
Optional
[Array
]) – Mask specifying valid rows when computing some statistics ofcost_matrix
, seesrc_mask
.tgt_mask (
Optional
[Array
]) – Mask specifying valid columns when computing some statistics ofcost_matrix
, seetgt_mask
.
Note
When defining a
Geometry
through acost_matrix
, it is important to select anepsilon
regularization parameter that is meaningful. That parameter can be provided by the user, or assigned a default value through a simple rule, using themean_cost_matrix
.Methods
apply_cost
(arr[, axis, fn])Apply
cost_matrix
to array (vector or matrix).apply_kernel
(scaling[, eps, axis])Apply
kernel_matrix
on positive scaling vector.apply_lse_kernel
(f, g, eps[, vec, axis])Apply
kernel_matrix
in log domain.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.
mask
(src_mask, tgt_mask[, mask_value])Mask rows or columns of a geometry.
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.
scaling_from_potential
(potential)Compute scaling vector from dual potential.
set_scale_cost
(scale_cost)Modify how to rescale of the
cost_matrix
.subset
(src_ixs, tgt_ixs, **kwargs)Subset rows or columns of a geometry.
to_LRCGeometry
([rank, tol, rng, scale])Factorize the cost matrix using either SVD (full) or [Indyk et al., 2019].
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
Check quickly if casting geometry as LRC makes sense.
Cost matrix, recomputed from kernel if only kernel was specified.
Output rank of cost matrix, if any was provided.
The data type.
Epsilon regularization value.
Compute and return inverse of scaling factor for cost matrix.
Whether geometry cost/kernel should be recomputed on the fly.
Whether cost is computed by taking squared Euclidean distance.
Whether geometry cost/kernel is a symmetric matrix.
Kernel matrix.
Mean of the
cost_matrix
.Median of the
cost_matrix
.Shape of the geometry.
Mask of shape
[num_a,]
to computecost_matrix
statistics.Mask of shape
[num_b,]
to computecost_matrix
statistics.