implementation comments

Author: eloitanguy

ot/lp/_barycenter_solvers.py

@@ -199,14 +199,12 @@ def free_support_barycenter(
     measures_weights : list of N (k_i,) array-like
         Numpy arrays where each numpy array has :math:`k_i` non-negatives values summing to one
         representing the weights of each discrete input measure
-
     X_init : (k,d) array-like
         Initialization of the support locations (on `k` atoms) of the barycenter
     b : (k,) array-like
         Initialization of the weights of the barycenter (non-negatives, sum to 1)
     weights : (N,) array-like
         Initialization of the coefficients of the barycenter (non-negatives, sum to 1)
-
     numItermax : int, optional
         Max number of iterations
     stopThr : float, optional
@@ -219,13 +217,11 @@ def free_support_barycenter(
         If compiled with OpenMP, chooses the number of threads to parallelize.
         "max" selects the highest number possible.
 
-
     Returns
     -------
     X : (k,d) array-like
         Support locations (on k atoms) of the barycenter
 
-
     .. _references-free-support-barycenter:
     References
     ----------
@@ -428,20 +424,20 @@ def generalized_free_support_barycenter(
         return Y
 
 
-class StoppingCriterionReached(Exception):
-    pass
-
-
 def free_support_barycenter_generic_costs(
     measure_locations,
     measure_weights,
     X_init,
     cost_list,
-    B,
+    ground_bary=None,
     a=None,
     numItermax=100,
     stopThr=1e-5,
     log=False,
+    ground_bary_lr=1e-2,
+    ground_bary_numItermax=100,
+    ground_bary_stopThr=1e-5,
+    ground_bary_solver="SGD",
 ):
     r"""
     Solves the OT barycenter problem for generic costs using the fixed point
@@ -507,14 +503,15 @@ def free_support_barycenter_generic_costs(
         List of K arrays of measure weights, each of shape (m_k).
     X_init : array-like
         Array of shape (n, d) representing initial barycenter points.
-    cost_list : list of callable
+    cost_list : list of callable or callable
         List of K cost functions :math:`c_k: \mathbb{R}^{n\times
         d}\times\mathbb{R}^{m_k\times d_k} \rightarrow \mathbb{R}_+^{n\times
-        m_k}`.
-    B : callable
+        m_k}`. If cost_list is a single callable, the same cost is used K times.
+    ground_bary : callable or None, optional
         Function List(array(n, d_k)) -> array(n, d) accepting a list of K arrays
         of shape (n\times d_K), computing the ground barycenters (broadcasted
-        over n).
+        over n). If not provided, done with Adam on PyTorch (requires PyTorch
+        backend)
     a : array-like, optional
         Array of shape (n,) representing weights of the barycenter
         measure.Defaults to uniform.
@@ -524,6 +521,16 @@ def free_support_barycenter_generic_costs(
         If the iterations move less than this, terminate (default is 1e-5).
     log : bool, optional
         Whether to return the log dictionary (default is False).
+    ground_bary_lr : float, optional
+        Learning rate for the ground barycenter solver (if auto is used).
+    ground_bary_numItermax : int, optional
+        Maximum number of iterations for the ground barycenter solver (if auto
+        is used).
+    ground_bary_stopThr : float, optional
+        Stop threshold for the ground barycenter solver (if auto is used).
+    ground_bary_solver : str, optional
+        Solver for auto ground bary solver (torch SGD or Adam). Default is
+        "SGD".
 
     Returns
     -------
@@ -549,49 +556,85 @@ def free_support_barycenter_generic_costs(
     See Also
     --------
     ot.lp.free_support_barycenter : Free support solver for the case where
-    :math:`c_k(x,y) = \|x-y\|_2^2`. ot.lp.generalized_free_support_barycenter :
-    Free support solver for the case where :math:`c_k(x,y) = \|P_kx-y\|_2^2`
-    with :math:`P_k` linear.
+    :math:`c_k(x,y) = \lambda_k\|x-y\|_2^2`.
+    ot.lp.generalized_free_support_barycenter : Free support solver for the case
+    where :math:`c_k(x,y) = \|P_kx-y\|_2^2` with :math:`P_k` linear.
     """
     nx = get_backend(X_init, measure_locations[0])
     K = len(measure_locations)
     n = X_init.shape[0]
     if a is None:
         a = nx.ones(n, type_as=X_init) / n
+    if callable(cost_list):  # use the given cost for all K pairs
+        cost_list = [cost_list] * K
+    auto_ground_bary = False
+
+    if ground_bary is None:
+        auto_ground_bary = True
+        assert str(nx) == "torch", (
+            f"Backend {str(nx)} is not compatible with ground_bary=None, it"
+            "must be provided if not using PyTorch backend"
+        )
+        try:
+            import torch
+            from torch.optim import Adam, SGD
+
+            def ground_bary(y, x_init):
+                x = x_init.clone().detach().requires_grad_(True)
+                solver = Adam if ground_bary_solver == "Adam" else SGD
+                opt = solver([x], lr=ground_bary_lr)
+                for _ in range(ground_bary_numItermax):
+                    x_prev = x.data.clone()
+                    opt.zero_grad()
+                    # inefficient cost computation but compatible
+                    # with the choice of cost_list[k] giving the cost matrix
+                    loss = torch.sum(
+                        torch.stack(
+                            [torch.diag(cost_list[k](x, y[k])) for k in range(K)]
+                        )
+                    )
+                    loss.backward()
+                    opt.step()
+                    diff = torch.sum((x.data - x_prev) ** 2)
+                    if diff < ground_bary_stopThr:
+                        break
+                return x.detach()
+
+        except ImportError:
+            raise ImportError("PyTorch is required to use ground_bary=None")
+
     X_list = [X_init] if log else []  # store the iterations
     X = X_init
     dX_list = []  # store the displacement squared norms
-    exit_status = "Unknown"
-
-    try:
-        for _ in range(numItermax):
-            pi_list = [  # compute the pairwise transport plans
-                emd(a, measure_weights[k], cost_list[k](X, measure_locations[k]))
-                for k in range(K)
-            ]
-            Y_perm = []
-            for k in range(K):  # compute barycentric projections
-                Y_perm.append(n * pi_list[k] @ measure_locations[k])
-            X_next = B(Y_perm)
-
-            if log:
-                X_list.append(X_next)
+    exit_status = "Max iterations reached"
+
+    for _ in range(numItermax):
+        pi_list = [  # compute the pairwise transport plans
+            emd(a, measure_weights[k], cost_list[k](X, measure_locations[k]))
+            for k in range(K)
+        ]
+        Y_perm = []
+        for k in range(K):  # compute barycentric projections
+            Y_perm.append(n * pi_list[k] @ measure_locations[k])
+        if auto_ground_bary:  # use previous position as initialization
+            X_next = ground_bary(Y_perm, X)
+        else:
+            X_next = ground_bary(Y_perm)
 
-            # stationary criterion: move less than the threshold
-            dX = nx.sum((X - X_next) ** 2)
-            X = X_next
+        if log:
+            X_list.append(X_next)
 
-            if log:
-                dX_list.append(dX)
+        # stationary criterion: move less than the threshold
+        dX = nx.sum((X - X_next) ** 2)
+        X = X_next
 
-            if dX < stopThr:
-                exit_status = "Stationary Point"
-                raise StoppingCriterionReached
+        if log:
+            dX_list.append(dX)
 
-        exit_status = "Max iterations reached"
-        raise StoppingCriterionReached
+        if dX < stopThr:
+            exit_status = "Stationary Point"
+            break
 
-    except StoppingCriterionReached:
-        if log:
-            return X, {"X_list": X_list, "exit_status": exit_status, "dX_list": dX_list}
-        return X
+    if log:
+        return X, {"X_list": X_list, "exit_status": exit_status, "dX_list": dX_list}
+    return X

test/test_ot.py

@@ -13,8 +13,6 @@
 from ot.datasets import make_1D_gauss as gauss
 from ot.backend import torch, tf
 
-# import ot.lp._barycenter_solvers  # TODO: remove this import
-
 
 def test_emd_dimension_and_mass_mismatch():
     # test emd and emd2 for dimension mismatch
@@ -414,14 +412,14 @@ def cost(x, y):
 
     cost_list = [cost, cost]
 
-    def B(y):
+    def ground_bary(y):
         out = 0
         for yk in y:
             out += yk / len(y)
         return out
 
     X = ot.lp.free_support_barycenter_generic_costs(
-        measures_locations, measures_weights, X_init, cost_list, B
+        measures_locations, measures_weights, X_init, cost_list, ground_bary
     )
 
     np.testing.assert_allclose(X, bar_locations, rtol=1e-5, atol=1e-7)
@@ -432,7 +430,7 @@ def B(y):
         measures_weights,
         X_init,
         cost_list,
-        B,
+        ground_bary,
         a=ot.unif(1),
         log=True,
     )
@@ -449,12 +447,95 @@ def B(y):
         measures_weights,
         X_init,
         cost_list,
-        B,
+        ground_bary,
         numItermax=1,
         log=True,
     )
     assert log2["exit_status"] == "Max iterations reached"
 
+    # test with a single callable cost
+    X3, log3 = ot.lp.free_support_barycenter_generic_costs(
+        measures_locations,
+        measures_weights,
+        X_init,
+        cost,
+        ground_bary,
+        numItermax=1,
+        log=True,
+    )
+
+    # test with no ground_bary but in numpy: requires pytorch backend
+    with pytest.raises(AssertionError):
+        ot.lp.free_support_barycenter_generic_costs(
+            measures_locations,
+            measures_weights,
+            X_init,
+            cost_list,
+            ground_bary=None,
+            numItermax=1,
+        )
+
+
+@pytest.mark.skipif(not torch, reason="No torch available")
+def test_free_support_barycenter_generic_costs_auto_ground_bary():
+    measures_locations = [
+        torch.tensor([1.0]).reshape((1, 1)),
+        torch.tensor([2.0]).reshape((1, 1)),
+    ]
+    measures_weights = [torch.tensor([1.0]), torch.tensor([1.0])]
+
+    X_init = torch.tensor([1.2]).reshape((1, 1))
+
+    def cost(x, y):
+        return ot.dist(x, y)
+
+    cost_list = [cost, cost]
+
+    def ground_bary(y):
+        out = 0
+        for yk in y:
+            out += yk / len(y)
+        return out
+
+    X = ot.lp.free_support_barycenter_generic_costs(
+        measures_locations,
+        measures_weights,
+        X_init,
+        cost_list,
+        ground_bary,
+        numItermax=1,
+    )
+
+    X2, log2 = ot.lp.free_support_barycenter_generic_costs(
+        measures_locations,
+        measures_weights,
+        X_init,
+        cost_list,
+        ground_bary=None,
+        ground_bary_lr=1e-2,
+        ground_bary_stopThr=1e-20,
+        ground_bary_numItermax=50,
+        numItermax=10,
+        log=True,
+    )
+
+    np.testing.assert_allclose(X2.numpy(), X.numpy(), rtol=1e-4, atol=1e-4)
+
+    X3 = ot.lp.free_support_barycenter_generic_costs(
+        measures_locations,
+        measures_weights,
+        X_init,
+        cost_list,
+        ground_bary=None,
+        ground_bary_lr=1e-2,
+        ground_bary_stopThr=1e-20,
+        ground_bary_numItermax=50,
+        numItermax=10,
+        ground_bary_solver="Adam",
+    )
+
+    np.testing.assert_allclose(X2.numpy(), X3.numpy(), rtol=1e-3, atol=1e-3)
+
 
 def test_free_support_barycenter_generic_costs_backends(nx):
     measures_locations = [
@@ -469,22 +550,22 @@ def cost(x, y):
 
     cost_list = [cost, cost]
 
-    def B(y):
+    def ground_bary(y):
         out = 0
         for yk in y:
             out += yk / len(y)
         return out
 
     X = ot.lp.free_support_barycenter_generic_costs(
-        measures_locations, measures_weights, X_init, cost_list, B
+        measures_locations, measures_weights, X_init, cost_list, ground_bary
     )
 
     measures_locations2 = nx.from_numpy(*measures_locations)
     measures_weights2 = nx.from_numpy(*measures_weights)
     X_init2 = nx.from_numpy(X_init)
 
     X2 = ot.lp.free_support_barycenter_generic_costs(
-        measures_locations2, measures_weights2, X_init2, cost_list, B
+        measures_locations2, measures_weights2, X_init2, cost_list, ground_bary
     )
 
     np.testing.assert_allclose(X, nx.to_numpy(X2))