ott.geometry.low_rank.LRCGeometry#
- class ott.geometry.low_rank.LRCGeometry(cost_1, cost_2, bias=0.0, scale_factor=1.0, scale_cost=1.0, **kwargs)[source]#
Geometry whose cost is defined by product of two low-rank matrices.
Implements geometries that are defined as low rank products, i.e. for which there exists two matrices \(A\) and \(B\) of \(r\) columns such that the cost of the geometry equals \(AB^T\). Apart from being faster to apply to a vector, these geometries are characterized by the fact that adding two such geometries should be carried out by concatenating factors, i.e. if \(C = AB^T\) and \(D = EF^T\) then \(C + D = [A,E][B,F]^T\)
- Parameters:
cost_1 (
Array
) – Array of shape[num_a, r]
.cost_2 (
Array
) – Array of shape[num_b, r]
.bias (
float
) – constant added to entire cost matrix.scale – Value used to rescale the factors of the low-rank geometry.
scale_cost (
Union
[float
,Literal
['mean'
,'max_bound'
,'max_cost'
]]) – option to rescale the cost matrix. Implemented scalings are ‘max_bound’, ‘mean’ and ‘max_cost’. Alternatively, a float factor can be given to rescale the cost such thatcost_matrix /= scale_cost
.scale_factor (
float
)
Methods
apply_cost
(arr[, axis, fn, is_linear])Apply
cost_matrix
to array (vector or matrix).apply_kernel
(vec[, 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.
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
([row_ixs, col_ixs])Subset rows or columns of a geometry.
to_LRCGeometry
([rank, tol, rng, scale])Return self.
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
Constant offset added to the entire
cost_matrix
.Check quickly if casting geometry as LRC makes sense.
First factor of the
cost_matrix
.Second factor of the
cost_matrix
.Materialize the cost matrix.
Output rank of cost matrix, if any was provided.
The data type.
Epsilon regularization value.
Epsilon scheduler.
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.
Standard deviation of all values stored in
cost_matrix
.