ott.solvers.linear.sinkhorn.SinkhornOutput

Contents

ott.solvers.linear.sinkhorn.SinkhornOutput#

class ott.solvers.linear.sinkhorn.SinkhornOutput(potentials, errors=None, reg_ot_cost=None, ot_prob=None, threshold=None, converged=None, inner_iterations=None)[source]#

Holds the output of a Sinkhorn solver applied to a problem.

Objects of this class contain both solutions and problem definition of a regularized OT problem, along several methods that can be used to access its content, to, for instance, materialize an OT matrix or apply it to a vector (without having to materialize it when not needed).

Parameters:
  • f – dual variables vector of size ot.prob.shape[0] returned by Sinkhorn

  • g – dual variables vector of size ot.prob.shape[1] returned by Sinkhorn

  • errors (Array | None) – vector or errors, along iterations. This vector is of size max_iterations // inner_iterations where those were the parameters passed on to the Sinkhorn solver. For each entry indexed at i, errors[i] can be either a real non-negative value (meaning the algorithm recorded that error at the i * inner_iterations iteration), a jnp.inf value (meaning the algorithm computed that iteration but did not compute its error, because, for instance, i < min_iterations // inner_iterations), or a -1, meaning that execution was terminated before that iteration, because the criterion was found to be smaller than threshold.

  • reg_ot_cost (float | None) – the regularized optimal transport cost. By default this is the linear contribution + KL term. See ent_reg_cost, primal_cost and dual_cost for other objective values.

  • ot_prob (LinearProblem | None) – stores the definition of the OT problem, including geometry, marginals, unbalanced regularizers, etc.

  • threshold (Array | None) – convergence threshold used to control the termination of the algorithm.

  • converged (bool | None) – whether the output corresponds to a solution whose error is below the convergence threshold.

  • inner_iterations (int | None) – number of iterations that were run between two computations of errors.

  • potentials (Tuple[Array, ...]) –

Methods

apply(inputs[, axis, lse_mode])

Apply the transport to a ndarray; axis=1 for its transpose.

cost_at_geom(other_geom)

Return reg-OT cost for matrix, evaluated at other cost matrix.

count(value, /)

Return number of occurrences of value.

index(value[, start, stop])

Return first index of value.

marginal(axis)

param axis:

set(**kwargs)

Return a copy of self, with potential overwrites.

set_cost(ot_prob, lse_mode, use_danskin)

param ot_prob:

to_dual_potentials()

Return the entropic map estimator.

transport_cost_at_geom(other_geom)

Return bare transport cost of current solution at any geometry.

Attributes

a

b

converged

Alias for field number 5

dual_cost

Return dual transport cost, without considering regularizer.

ent_reg_cost

Entropy regularized cost.

errors

Alias for field number 1

f

The first dual potential.

g

The second dual potential.

geom

inner_iterations

Alias for field number 6

kl_reg_cost

KL regularized OT transport cost.

matrix

Transport matrix if it can be instantiated.

n_iters

Returns the total number of iterations that were needed to terminate.

ot_prob

Alias for field number 3

potentials

Alias for field number 0

primal_cost

Return transport cost of current transport solution at geometry.

reg_ot_cost

Alias for field number 2

scalings

threshold

Alias for field number 4

transport_mass

Sum of transport matrix.