jaxopt.projection.projection_transport
- jaxopt.projection.projection_transport(sim_matrix, marginals, make_solver=None, use_semi_dual=True)[source]
Projection onto the transportation polytope.
We solve
\[\underset{P \ge 0}{\text{argmin}} ~ ||P - S||_2^2 \quad \textrm{subject to} \quad P^\top \mathbf{1} = a, P \mathbf{1} = b\]or equivalently
\[\underset{P \ge 0}{\text{argmin}} ~ \langle P, C \rangle + \frac{1}{2} \|P\|_2^2 \quad \textrm{subject to} \quad P^\top \mathbf{1} = a, P \mathbf{1} = b\]where \(S\) is a similarity matrix, \(C\) is a cost matrix and \(S = -C\).
This implementation solves the semi-dual (see equation 10 in reference below) using LBFGS but the solver can be overidden using the
make_solveroption.For an entropy-regularized version, see
kl_projection_transport.- Parameters
sim_matrix (
Array) – similarity matrix, shape=(size_a, size_b).marginals (
Tuple) – a tuple (marginals_a, marginals_b), where marginals_a has shape=(size_a,) and marginals_b has shape=(size_b,).make_solver (
Optional[Callable]) – a function of the form make_solver(fun), for creating an iterative solver to minimize fun.use_semi_dual (
bool) – if true, use the semi-dual formulation in Equation (10) of https://arxiv.org/abs/1710.06276. Otherwise, use the dual-formulation in Equation (7).
- Returns
transportation matrix, shape=(size_a, size_b).
- Return type
plan
References
Smooth and Sparse Optimal Transport. Mathieu Blondel, Vivien Seguy, Antoine Rolet. In Proceedings of Artificial Intelligence and Statistics (AISTATS), 2018. https://arxiv.org/abs/1710.06276