# ott.geometry.grid.Grid#

class ott.geometry.grid.Grid(x=None, grid_size=None, cost_fns=None, num_a=None, grid_dimension=None, **kwargs)[source]#

Class describing the geometry of points taken in a Cartesian product.

This class implements a geometry in which probability measures are supported on a $$d$$-dimensional Cartesian grid, a Cartesian product of $$d$$ lists of values, each list being itself of size $$n_i$$.

The transportation cost between points in the grid is assumed to be separable, namely a sum of coordinate-wise cost functions, as in:

$cost(x,y) = \sum_{i=1}^d cost_i(x_i, y_i)$

where $$cost_i$$: R x R → R.

In such a regime, and despite the fact that the total number $$n_{total}$$ of points in the grid is exponential $$d$$ (namely $$\prod_i n_i$$), applying a kernel in the context of regularized optimal transport can be carried out in time that is of the order of $$n_{total}^{(1+1/d)}$$ using convolutions, either in the original domain or log-space domain. This class precomputes $$d$$ $$n_i$$ x $$n_i$$ cost matrices (one per dimension) and implements these two operations by carrying out these convolutions one dimension at a time.

Parameters:

Methods

 apply_cost(arr[, axis, fn]) Apply cost_matrix to array (vector or matrix). apply_kernel(scaling[, eps, axis]) Apply grid kernel on scaling vector. apply_lse_kernel(f, g, eps[, vec, axis]) Apply grid kernel in log space. 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]) Not implemented. potential_from_scaling(scaling) Compute dual potential vector from scaling vector. prepare_divergences(*args[, 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(src_ixs, tgt_ixs) Not implemented. to_LRCGeometry([scale]) Converts grid to low-rank geometry. transport_from_potentials(f, g[, axis]) Not implemented, use apply_transport_from_potentials() instead. transport_from_scalings(f, g[, axis]) Not implemented, use apply_transport_from_scalings() instead. 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

 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. geometries Cost matrices along each dimension of the grid. 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 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 Not implemented. shape Shape of the geometry. src_mask Mask of shape [num_a,] to compute cost_matrix statistics. tgt_mask Mask of shape [num_b,] to compute cost_matrix statistics.