ott.geometry.costs.SqPNorm.transport_map

ott.geometry.costs.SqPNorm.transport_map#

SqPNorm.transport_map(g)#

Get an optimal transport map for a concave function \(g\).

Uses Proposition 1 from [Klein et al., 2024] to define an OT map \(x - (\nabla h^*) \circ \nabla \bar g^h(x)\), where \(h^*\) is the Legendre transform of \(h\) and \(\bar g^h\) is the h_transform() of a concave function \(g\).

Parameters:

g (Callable[[Array], float]) – Concave function.

Return type:

Callable[[Array, Any], Array]

Returns:

The transport map with a signature (x, **kwargs).