ott.geometry.pointcloud.PointCloud

class ott.geometry.pointcloud.PointCloud(x, y=None, cost_fn=None, batch_size=None, scale_cost=1.0, **kwargs)

Defines geometry for 2 point clouds (possibly 1 vs itself).

When the number of points is large, setting the batch_size flag implies that cost and kernel matrices used to update potentials or scalings will be recomputed on the fly, rather than stored in memory.

Parameters:

• x (Array) – Array of shape [n, d].

• y (Optional[Array]) – Array of shape [m, d]. If None, use x.

• cost_fn (Optional[CostFn]) – Cost function between two points in dimension \(d\).

• batch_size (Optional[int]) – If None, the cost matrix corresponding to that point cloud is computed, stored and later re-used at each application of apply_lse_kernel(). When batch_size is a positive integer, computations are done in an online fashion, namely the cost matrix is recomputed at each call of the apply_lse_kernel() step, batch_size lines at a time, used on a vector and discarded. The online computation is particularly useful for big point clouds whose cost matrix does not fit in memory.

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

• kwargs (Any) – keyword arguments for Geometry.

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.
barycenter(weights)	Compute barycenter of points in self.x using weights.
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(x, y[, static_b])	Instantiate the geometries used for a divergence computation.
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([scale])	Convert point cloud to low-rank geometry.
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

batch_size	Batch size for online mode.
can_LRC	Check quickly if casting geometry as LRC makes sense.
cost_matrix	Cost matrix, recomputed from kernel if only kernel was specified.
cost_rank	Output rank of cost matrix, if any was provided.
dtype	The data type.
epsilon	Epsilon regularization value.
epsilon_scheduler	Epsilon scheduler.
inv_scale_cost	Compute and return inverse of scaling factor for cost matrix.
is_online	Whether the cost/kernel is computed on-the-fly.
is_squared_euclidean	Whether cost is computed by taking squared Euclidean distance.
is_symmetric	Whether geometry cost/kernel is a symmetric matrix.
kernel_matrix	Kernel matrix.
mean_cost_matrix	Mean of the cost_matrix.
median_cost_matrix	Median of the cost_matrix.
shape	Shape of the geometry.
std_cost_matrix	Standard deviation of all values stored in cost_matrix.