ott.geometry.grid.Grid.apply_lse_kernel#

Grid.apply_lse_kernel(f, g, eps, vec=None, axis=0)[source]#

Apply grid kernel in log space. See notes in parent class for use case.

Reshapes vector inputs below as grids, applies kernels onto each slice, and then expands the outputs as vectors.

More implementation details in [Schmitz et al., 2018].

Parameters:

  • f (Array) – jnp.ndarray, a vector of potentials

  • g (Array) – jnp.ndarray, a vector of potentials

  • eps (float) – float, regularization strength

  • vec (Optional[Array]) – jnp.ndarray, if needed, a vector onto which apply the kernel weighted by f and g.

  • axis (int) – axis (0 or 1) along which summation should be carried out.

Return type:

Array

Returns:

a vector, the result of kernel applied in lse space onto vec.