class ott.geometry.graph.Graph(graph=None, laplacian=None, t=0.001, n_steps=100, numerical_scheme='backward_euler', directed=False, tol=- 1.0, **kwargs)[source]#

Graph distance approximation using heat kernel [Crane et al., 2013, Heitz et al., 2021].

Approximates the heat kernel for large n_steps, which for small t approximates the geodesic exponential kernel \(e^{\frac{-d(x, y)^2}{t}}\).

For sparse graphs, sksparse.cholmod is required to compute the Cholesky decomposition. Differentiating w.r.t. the edge weights is currently possible only when the graph is represented as a dense adjacency matrix.

  • graph (Union[Array, BCOO, None]) – Graph represented as an adjacency matrix of shape [n, n]. If None, the symmetric graph Laplacian has to be specified.

  • laplacian (Union[Array, CSR, CSC, COO, BCOO, None]) – Symmetric graph Laplacian. The check for symmetry is NOT performed. If None, the graph has to be specified instead.

  • t (Optional[float]) – Constant used when approximating the geodesic exponential kernel. If None, use \(\frac{1}{|E|} \sum_{(u, v) \in E} weight(u, v)\) [Crane et al., 2013]. In this case, the graph must be specified and the edge weights are all assumed to be positive.

  • n_steps (int) – Maximum number of steps used to approximate the heat kernel.

  • numerical_scheme (Literal[‘backward_euler’, ‘crank_nicolson’]) – Numerical scheme used to solve the heat diffusion.

  • directed (bool) – Whether the graph is directed. If not, it will be made undirected as \(G + G^T\). This parameter is ignored when directly passing the Laplacian, which is assumed to be symmetric.

  • tol (float) – Relative tolerance with respect to the Hilbert metric, see [Peyré and Cuturi, 2019], Remark 4.12. Used when iteratively updating scalings. If negative, this option is ignored and only n_steps is used.

  • 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 kernel_matrix on positive scaling vector.

apply_lse_kernel(f, g, eps[, vec, axis])

Apply kernel_matrix in log domain on a pair of dual potential variables.

apply_square_cost(arr[, axis])

Apply elementwise-square of cost matrix to array (vector or matrix).

apply_transport_from_potentials(f, g, vec[, ...])

Since applying from potentials is not feasible in grids, use scalings.

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])

Not implemented.

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.


Compute dual potential vector from scaling vector.

prepare_divergences(*args[, static_b])

Instantiate 2 (or 3) geometries to compute a Sinkhorn divergence.


Compute scaling vector from dual potential.

subset(src_ixs, tgt_ixs, **kwargs)

Subset rows or columns of a geometry.

to_LRCGeometry([rank, tol, seed, scale])

Factorize the cost matrix using either SVD (full) or [Indyk et al., 2019].

transport_from_potentials(f, g)

Not implemented.

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.



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.


The underlying undirected graph as an adjacency matrix, if provided.


Compute and return inverse of scaling factor for cost matrix.


Whether geometry cost/kernel should be recomputed on the fly.


Whether graph or laplacian is sparse.


Whether cost is computed by taking squared-Eucl.


Whether geometry cost/kernel is a symmetric matrix.


Kernel matrix, either provided by user or recomputed from cost_matrix.


The graph Laplacian.


Mean of the cost_matrix.


Median of the cost_matrix.


Compute the scale of the epsilon, potentially based on data.


Shape of the geometry.


Instantiate the Cholesky solver and compute the factorization.


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


Constant used when approximating the geodesic exponential kernel.


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