da class and tests

PythonOT · rflamary · Oct 11, 2023 · Sep 20, 2023 · Sep 20, 2023 · Sep 20, 2023
commit ff4697552da9a75771dd4975c2b7439242460b45
diff --git a/examples/others/plot_SSNB.py b/examples/others/plot_SSNB.py
@@ -22,11 +22,13 @@
     \varphi \in \mathrm{argmin}_{\varphi \in \mathcal{F}}\ \mathrm{W}_2(\nabla \varphi \#\mu_s, \mu_t),
 
 where :math:`\mathcal{F}` is the space functions that are on every set :math:`E_k` l-strongly convex
-with an L-Lipschitz gradient, given :math:`(E_k)_{k \in [K]}` a partition of the ambiant source space.
+with an L-Lipschitz gradient, given :math:`(E_k)_{k \in [K]}` a partition of the ambient source space.
 
 We perform the optimisation on a low amount of fitting samples and with few iterations,
 since solving the SSNB problem is quite computationally expensive.
 
+THIS EXAMPLE REQUIRES CVXPY
+
 .. [58] François-Pierre Paty, Alexandre d’Aspremont, and Marco Cuturi. Regularity as regularization:
         Smooth and strongly convex brenier potentials in optimal transport. In International Conference
         on Artificial Intelligence and Statistics, pages 1222–1232. PMLR, 2020.
@@ -42,7 +44,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import ot
-import os
 
 # %%
 # Generating the fitting data

diff --git a/ot/da.py b/ot/da.py
@@ -8,6 +8,7 @@
 #         Michael Perrot <michael.perrot@univ-st-etienne.fr>
 #         Nathalie Gayraud <nat.gayraud@gmail.com>
 #         Ievgen Redko <ievgen.redko@univ-st-etienne.fr>
+#         Eloi Tanguy <eloi.tanguy@u-paris.fr>
 #
 # License: MIT License
 
@@ -22,6 +23,7 @@
 from .gaussian import empirical_bures_wasserstein_mapping, empirical_gaussian_gromov_wasserstein_mapping
 from .optim import cg
 from .optim import gcg
+from .mapping import nearest_brenier_potential_fit, nearest_brenier_potential_predict_bounds
 
 
 def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
@@ -2635,3 +2637,163 @@ def inverse_transform_labels(self, yt=None):
                 transp_ys.append(nx.dot(D1, transp.T).T)
 
             return transp_ys
+
+
+class NearestBrenierPotential(BaseTransport):
+    r"""
+    Smooth Strongly Convex Nearest Brenier Potentials (SSNB) is a method from :ref:`[58]` that computes
+    an l-strongly convex potential :math:`\varphi` with an L-Lipschitz gradient such that
+    :math:`\nabla \varphi \# \mu \approx \nu`. This regularity can be enforced only on the components of a partition
+    of the ambient space (encoded by point classes), which is a relaxation compared to imposing global regularity.
+
+    SSNBs approach the target measure by solving the optimisation problem:
+
+    .. math::
+        \\varphi \in \\text{argmin}_{\\varphi \in \\mathcal{F}}\ \\text{W}_2(\\nabla \\varphi \#\\mu_s, \\mu_t),
+
+    where :math:`\mathcal{F}` is the space functions that are on every set :math:`E_k` l-strongly convex
+    with an L-Lipschitz gradient, given :math:`(E_k)_{k \in [K]}` a partition of the ambient source space.
+
+    The problem is solved on "fitting" source and target data via a convex Quadratically Constrained Quadratic Program,
+    yielding the values :code:`phi` and the gradients :code:`G` at at the source points.
+    The images of "new" source samples are then found by solving a (simpler) Quadratically Constrained Linear Program
+    at each point, using the fitting "parameters" :code:`phi` and :code:`G`. We provide two possible images, which
+    correspond to "lower" and "upper potentials" (:ref:`[59]`, Theorem 3.14). Each of these two images are optimal
+    solutions of the SSNB problem, and can be used in practice.
+
+    Parameters
+    ----------
+    strongly_convex_constant : float, optional
+        constant for the strong convexity of the input potential phi, defaults to 0.6
+    gradient_lipschitz_constant : float, optional
+        constant for the Lipschitz property of the input gradient G, defaults to 1.4
+    its: int, optional
+        number of iterations, defaults to 100
+    log : bool, optional
+        record log if true
+    seed: int or RandomState or None, optional
+        Seed used for random number generator (for the initialisation in :code:`fit`.
+
+    References
+    ----------
+
+    .. [58] François-Pierre Paty, Alexandre d’Aspremont, and Marco Cuturi. Regularity as regularization:
+            Smooth and strongly convex brenier potentials in optimal transport. In International Conference
+            on Artificial Intelligence and Statistics, pages 1222–1232. PMLR, 2020.
+
+    .. [59] Adrien B Taylor. Convex interpolation and performance estimation of first-order methods for
+            convex optimization. PhD thesis, Catholic University of Louvain, Louvain-la-Neuve, Belgium,
+            2017.
+
+    See Also
+    --------
+    ot.mapping.nearest_brenier_potential_fit : Fitting the SSNB on source and target data
+    ot.mapping.nearest_brenier_potential_predict_bounds : Predicting SSNB images on new source data
+    """
+    def __init__(self, strongly_convex_constant=0.6, gradient_lipschitz_constant=1.4, log=False, its=100, seed=None):
+        self.strongly_convex_constant = strongly_convex_constant
+        self.gradient_lipschitz_constant = gradient_lipschitz_constant
+        self.log = log
+        self.its = its
+        self.seed = seed
+        self.fit_log, self.predict_log = None, None
+        self.phi, self.G = None, None
+        self.fit_Xs, self.fit_ys, self.fit_Xt = None, None, None
+
+    def fit(self, Xs=None, ys=None, Xt=None, yt=None):
+        r"""
+        Fits the Smooth Strongly Convex Nearest Brenier Potential [58] to the source data :code:`Xs` to the target data
+        :code:`Xt`, with the partition given by the (optional) labels :code:`ys`.
+
+        Wrapper for :code:`ot.mapping.nearest_brenier_potential_fit`.
+
+        THIS METHOD REQUIRES THE CVXPY LIBRARY
+
+        Parameters
+        ----------
+        Xs : array-like (n, d)
+            source points used to compute the optimal values phi and G
+        ys : array-like (n,), optional
+            classes of the reference points, defaults to a single class
+        Xt : array-like (n, d)
+            values of the gradients at the reference points X
+        yt : optional
+            ignored.
+
+        Returns
+        -------
+        self : object
+            Returns self.
+
+        References
+        ----------
+
+        .. [58] François-Pierre Paty, Alexandre d’Aspremont, and Marco Cuturi. Regularity as regularization:
+                Smooth and strongly convex brenier potentials in optimal transport. In International Conference
+                on Artificial Intelligence and Statistics, pages 1222–1232. PMLR, 2020.
+
+        See Also
+        --------
+        ot.mapping.nearest_brenier_potential_fit : Fitting the SSNB on source and target data
+
+        """
+        self.fit_Xs, self.fit_ys, self.fit_Xt = Xs, ys, Xt
+        returned = nearest_brenier_potential_fit(Xs, Xt, X_classes=ys,
+                                                 strongly_convex_constant=self.strongly_convex_constant,
+                                                 gradient_lipschitz_constant=self.gradient_lipschitz_constant,
+                                                 its=self.its, log=self.log)
+
+        if self.log:
+            self.phi, self.G, self.fit_log = returned
+        else:
+            self.phi, self.G = returned
+
+        return self
+
+    def transform(self, Xs, ys=None):
+        r"""
+        Computes the images of the new source samples :code:`Xs` of classes :code:`ys` by the fitted
+        Smooth Strongly Convex Nearest Brenier Potentials (SSNB) :ref:`[58]`. The output is the images of two SSNB optimal
+        maps, called 'lower' and 'upper' potentials (from :ref:`[59]`, Theorem 3.14).
+
+        Wrapper for :code:`nearest_brenier_potential_predict_bounds`.
+
+        THIS METHOD REQUIRES THE CVXPY LIBRARY
+
+        Parameters
+        ----------
+        Xs : array-like (m, d)
+            input source points
+        ys : : array_like (m,), optional
+            classes of the input source points, defaults to a single class
+
+        Returns
+        -------
+        G_lu : array-like (2, m, d)
+            gradients of the lower and upper bounding potentials at Y (images of the source inputs)
+
+        References
+        ----------
+
+        .. [58] François-Pierre Paty, Alexandre d’Aspremont, and Marco Cuturi. Regularity as regularization:
+                Smooth and strongly convex brenier potentials in optimal transport. In International Conference
+                on Artificial Intelligence and Statistics, pages 1222–1232. PMLR, 2020.
+
+        .. [59] Adrien B Taylor. Convex interpolation and performance estimation of first-order methods for
+                convex optimization. PhD thesis, Catholic University of Louvain, Louvain-la-Neuve, Belgium,
+                2017.
+
+        See Also
+        --------
+        ot.mapping.nearest_brenier_potential_predict_bounds : Predicting SSNB images on new source data
+
+        """
+        returned = nearest_brenier_potential_predict_bounds(
+            self.fit_Xs, self.phi, self.G, Xs, X_classes=self.fit_ys, Y_classes=ys,
+            strongly_convex_constant=self.strongly_convex_constant,
+            gradient_lipschitz_constant=self.gradient_lipschitz_constant, log=self.log)
+        if self.log:
+            _, G_lu, self.predict_log = returned
+        else:
+            _, G_lu = returned
+        return G_lu
diff --git a/ot/mapping.py b/ot/mapping.py
@@ -45,26 +45,26 @@ def nearest_brenier_potential_fit(X, V, X_classes=None, a=None, b=None, strongly
 
     Accepts any compatible backend, but will perform the QCQP optimisation on Numpy arrays, and convert back at the end.
 
+    THIS FUNCTION REQUIRES THE CVXPY LIBRARY
+
     Parameters
     ----------
-    X: array-like (n, d)
+    X : array-like (n, d)
         reference points used to compute the optimal values phi and G
-    V: array-like (n, d)
+    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
+    a : array-like (n,), optional
         weights for the reference points X, defaults to uniform
-    b: array-like (n,), optional
+    b : array-like (n,), optional
         weights for the target points V, defaults to uniform
     strongly_convex_constant : float, optional
         constant for the strong convexity of the input potential phi, defaults to 0.6
     gradient_lipschitz_constant : float, optional
         constant for the Lipschitz property of the input gradient G, defaults to 1.4
     its: int, optional
         number of iterations, defaults to 100
-    pbar: bool, optional
-        if True show a progress bar, defaults to False
     log : bool, optional
         record log if true
     seed: int or RandomState or None, optional
@@ -87,6 +87,11 @@ def nearest_brenier_potential_fit(X, V, X_classes=None, a=None, b=None, strongly
             Smooth and strongly convex brenier potentials in optimal transport. In International Conference
             on Artificial Intelligence and Statistics, pages 1222–1232. PMLR, 2020.
 
+    See Also
+    --------
+    ot.mapping.nearest_brenier_potential_predict_bounds : Predicting SSNB images on new source data
+    ot.da.NearestBrenierPotential : BaseTransport wrapper for SSNB
+
     """
     try:
         import cvxpy as cvx
@@ -197,7 +202,8 @@ def nearest_brenier_potential_predict_bounds(X, phi, G, Y, X_classes=None, Y_cla
     Compute the values of the lower and upper bounding potentials at the input points Y, using the potential optimal
     values phi at X and their gradients G at X. The 'lower' potential corresponds to the method from :ref:`[58]`,
     Equation 2, while the bounding property and 'upper' potential come from :ref:`[59]`, Theorem 3.14 (taking into
-    account the fact that this theorem's statement has a min instead of a max, which is a typo).
+    account the fact that this theorem's statement has a min instead of a max, which is a typo). Both potentials are
+    optimal for the SSNB problem.
 
     If :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`, for each :math:`y \in E_k`, this function solves the convex QCQP problems,
@@ -221,6 +227,8 @@ def nearest_brenier_potential_predict_bounds(X, phi, G, Y, X_classes=None, Y_cla
 
     The constants :math:`c_1, c_2, c_3` only depend on `strongly_convex_constant` and `gradient_lipschitz_constant`.
 
+    THIS FUNCTION REQUIRES THE CVXPY LIBRARY
+
     Parameters
     ----------
     X : array-like (n, d)
@@ -264,6 +272,11 @@ def nearest_brenier_potential_predict_bounds(X, phi, G, Y, X_classes=None, Y_cla
             convex optimization. PhD thesis, Catholic University of Louvain, Louvain-la-Neuve, Belgium,
             2017.
 
+    See Also
+    --------
+    ot.mapping.nearest_brenier_potential_fit : Fitting the SSNB on source and target data
+    ot.da.NearestBrenierPotential : BaseTransport wrapper for SSNB
+
     """
     try:
         import cvxpy as cvx