Function nearest_brenier_potential_fit

ot/mapping.py:23–199

r""" Computes optimal values and gradients at X for a strongly convex potential :math:`\varphi` with Lipschitz gradients on the partitions defined by `X_classes`, where :math:`\varphi` is optimal such that :math:`\nabla \varphi \#\mu \approx \nu`, given samples :math:`X = x_1, \cdots, x_

(
    X,
    V,
    X_classes=None,
    a=None,
    b=None,
    strongly_convex_constant=0.6,
    gradient_lipschitz_constant=1.4,
    its=100,
    log=False,
    init_method="barycentric",
    solver=None,
)

def nearest_brenier_potential_fit(
    X,
    V,
    X_classes=None,
    a=None,
    b=None,
    strongly_convex_constant=0.6,
    gradient_lipschitz_constant=1.4,
    its=100,
    log=False,
    init_method="barycentric",
    solver=None,
):
    r"""
    Computes optimal values and gradients at X for a strongly convex potential :math:`\varphi` with Lipschitz gradients
    on the partitions defined by `X_classes`, where :math:`\varphi` is optimal such that
    :math:`\nabla \varphi \#\mu \approx \nu`, given samples :math:`X = x_1, \cdots, x_n \sim \mu` and
    :math:`V = v_1, \cdots, v_n \sim \nu`. Finding such a potential that has the desired regularity on the
    partition :math:`(E_k)_{k \in [K]}` (given by the classes `X_classes`) is equivalent to finding optimal values
    `phi` for the :math:`\varphi(x_i)` and its gradients :math:`\nabla \varphi(x_i)` (variable`G`).
    In practice, these optimal values are found by solving the following problem

    .. math::
        :nowrap:

        \begin{gather*}
        \text{min} \sum_{i,j}\pi_{i,j}\|g_i - v_j\|_2^2 \\
        g_1,\cdots, g_n \in \mathbb{R}^d,\; \varphi_1, \cdots, \varphi_n \in \mathbb{R},\; \pi \in \Pi(a, b) \\
        \text{s.t.}\ \forall k \in [K],\; \forall i,j \in I_k: \\
        \varphi_i-\varphi_j-\langle g_j, x_i-x_j\rangle \geq c_1\|g_i - g_j\|_2^2
        + c_2\|x_i-x_j\|_2^2 - c_3\langle g_j-g_i, x_j -x_i \rangle.
        \end{gather*}

    The constants :math:`c_1, c_2, c_3` only depend on `strongly_convex_constant` and `gradient_lipschitz_constant`.
    The constraint :math:`\pi \in \Pi(a, b)` denotes the fact that the matrix :math:`\pi` belong to the OT polytope
    of marginals a and b. :math:`I_k` is the subset of :math:`[n]` of the i such that :math:`x_i` is in the
    partition (or class) :math:`E_k`, i.e. `X_classes[i] == k`.

    This problem is solved by alternating over the variable :math:`\pi` and the variables :math:`\varphi_i, g_i`.
    For :math:`\pi`, the problem is the standard discrete OT problem, and for :math:`\varphi_i, g_i`, the
    problem is a convex QCQP solved using :code:`cvxpy` (ECOS solver).

    Accepts any compatible backend, but will perform the QCQP optimisation on Numpy arrays, and convert back at the end.

    .. warning:: This function requires the CVXPY library
    .. warning:: Accepts any backend but will convert to Numpy then back to the backend.

    Parameters
    ----------
    X : array-like (n, d)
        reference points used to compute the optimal values phi and G
    V : array-like (n, d)
        values of the gradients at the reference points X
    X_classes : array-like (n,), optional
        classes of the reference points, defaults to a single class
    a : array-like (n,), optional
        weights for the reference points X, defaults to uniform
    b : array-like (n,), optional