ott.tools.progot.get_epsilon_schedule

ott.tools.progot.get_epsilon_schedule#

ott.tools.progot.get_epsilon_schedule(geom, *, alphas, epsilon_scales, y_eval, start_epsilon_scale=1.0, **kwargs)[source]#

Get the epsilon regularization schedule.

See Algorithm 4 in [Kassraie et al., 2024] for more information.

Parameters:
  • geom (PointCloud) – Point cloud geometry.

  • alphas (Array) – Stepsize schedule of shape [num_steps,].

  • epsilon_scales (Array) – Array of shape [num_scales,] from which to select the best scale of the default epsilon in the (y, y) point cloud.

  • y_eval (Array) – Array of shape [k, d] from the target distribution used to compute the error.

  • start_epsilon_scale (float) – Constant by which to scale the initial epsilon.

  • kwargs (Any) – Keyword arguments for Sinkhorn.

Return type:

Array

Returns:

The epsilon regularization schedule of shape [num_steps,].