ott.geometry.graph.Graph#

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.

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

Methods

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_epsilon(other)

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.

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.

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.

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.

graph

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

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_sparse

Whether graph or laplacian is sparse.

is_squared_euclidean

Whether cost is computed by taking squared-Eucl.

is_symmetric

Whether geometry cost/kernel is a symmetric matrix.

kernel_matrix

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

laplacian

The graph Laplacian.

mean_cost_matrix

Mean of the cost_matrix.

median_cost_matrix

Median of the cost_matrix.

scale_epsilon

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

shape

Shape of the geometry.

solver

Instantiate the Cholesky solver and compute the factorization.

src_mask

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

t

Constant used when approximating the geodesic exponential kernel.

tgt_mask

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