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.

  • x (Optional[Sequence[Array]]) – list of arrays of varying sizes, describing the locations of the grid. Locations are provided as a list of arrays, that is \(d\) vectors of (possibly varying) size \(n_i\). The resulting grid is the Cartesian product of these vectors.

  • grid_size (Optional[Sequence[int]]) – tuple of integers describing grid sizes, namely \((n_1,...,n_d)\). This will only be used if x is None. In that case the grid will be assumed to lie in the hypercube \([0,1]^d\), with the \(d\) dimensions, described as points regularly sampled in \([0,1]\).

  • cost_fns (Optional[Sequence[CostFn]]) – a sequence of \(d\) cost functions, each being a cost taking two reals as inputs to output a real number.

  • num_a (Optional[int]) – total size of grid. This parameters will be computed from other inputs.

  • grid_dimension (Optional[int]) – dimension of grid. This parameters will be computed from other inputs.

  • kwargs (Any) – keyword arguments for Geometry.


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 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.


Compute dual potential vector from scaling vector.

prepare_divergences(*args[, static_b])

Instantiate the geometries used for a divergence computation.


Compute scaling vector from dual potential.


Modify how to rescale of the cost_matrix.

subset(src_ixs, tgt_ixs)

Not implemented.


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.



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.


Cost matrices along each dimension of the grid.


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.


Not implemented.


Shape of the geometry.


Mask of shape [num_a,] to compute cost_matrix statistics.


Mask of shape [num_b,] to compute cost_matrix statistics.