ott.geometry.grid.Grid.marginal_from_potentials#

Grid.marginal_from_potentials(f, g, axis=0)#

Output marginal of transportation matrix from potentials.

This applies first lse kernel in the standard way, removes the correction used to stabilize computations, and lifts this with an exp to recover either of the marginals corresponding to the transport map induced by potentials.

Parameters

  • f (Array) – jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix

  • g (Array) – jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix

  • axis (int) – axis along which to integrate, returns marginal on other axis.

Return type

Array

Returns

a vector of marginals of the transport matrix.