enable_autodiff for POT wasserstein_distance

1 files changed, 53 insertions, 11 deletions

diff --git a/src/python/gudhi/wasserstein/wasserstein.py b/src/python/gudhi/wasserstein/wasserstein.py
index 5df66cf9..9660b99b 100644
--- a/src/python/gudhi/wasserstein/wasserstein.py
+++ b/src/python/gudhi/wasserstein/wasserstein.py
@@ -53,17 +53,30 @@ def _build_dist_matrix(X, Y, order, internal_p):
     return Cf
 
 
-def _perstot(X, order, internal_p):
+def _perstot_autodiff(X, order, internal_p):
+    '''
+    Version of _perstot that works on eagerpy tensors.
+    '''
+    return _dist_to_diag(X, internal_p).norms.lp(order)
+
+def _perstot(X, order, internal_p, enable_autodiff):
     '''
     :param X: (n x 2) numpy.array (points of a given diagram).
     :param order: exponent for Wasserstein. Default value is 2.
     :param internal_p: Ground metric on the (upper-half) plane (i.e. norm L^p in R^2); Default value is 2 (Euclidean norm).
+    :param enable_autodiff: If X is torch.tensor, tensorflow.Tensor or jax.numpy.ndarray, make the computation
+        transparent to automatic differentiation.
+    :type enable_autodiff: bool
     :returns: float, the total persistence of the diagram (that is, its distance to the empty diagram).
     '''
-    return np.linalg.norm(_dist_to_diag(X, internal_p), ord=order)
+    if enable_autodiff:
+        import eagerpy as ep
+        return _perstot_autodiff(ep.astensor(X), order, internal_p).raw
+    else:
+        return np.linalg.norm(_dist_to_diag(X, internal_p), ord=order)
 
 
-def wasserstein_distance(X, Y, matching=False, order=2., internal_p=2.):
+def wasserstein_distance(X, Y, matching=False, order=2., internal_p=2., enable_autodiff=False):
     '''
     :param X: (n x 2) numpy.array encoding the (finite points of the) first diagram. Must not contain essential points
                 (i.e. with infinite coordinate).
@@ -74,6 +87,9 @@ def wasserstein_distance(X, Y, matching=False, order=2., internal_p=2.):
     :param order: exponent for Wasserstein; Default value is 2.
     :param internal_p: Ground metric on the (upper-half) plane (i.e. norm L^p in R^2);
                        Default value is 2 (Euclidean norm).
+    :param enable_autodiff: If X and Y are torch.tensor, tensorflow.Tensor or jax.numpy.ndarray, make the computation
+        transparent to automatic differentiation.
+    :type enable_autodiff: bool
     :returns: the Wasserstein distance of order q (1 <= q < infinity) between persistence diagrams with
               respect to the internal_p-norm as ground metric.
               If matching is set to True, also returns the optimal matching between X and Y.
@@ -82,23 +98,30 @@ def wasserstein_distance(X, Y, matching=False, order=2., internal_p=2.):
     m = len(Y)
 
     # handle empty diagrams
-    if X.size == 0:
-        if Y.size == 0:
+    if n == 0:
+        if m == 0:
             if not matching:
+                # What if enable_autodiff?
                 return 0.
             else:
                 return 0., np.array([])
         else:
             if not matching:
-                return _perstot(Y, order, internal_p)
+                return _perstot(Y, order, internal_p, enable_autodiff)
             else:
-                return _perstot(Y, order, internal_p), np.array([[-1, j] for j in range(m)])
-    elif Y.size == 0:
+                return _perstot(Y, order, internal_p, enable_autodiff), np.array([[-1, j] for j in range(m)])
+    elif m == 0:
         if not matching:
-            return _perstot(X, order, internal_p)
+            return _perstot(X, order, internal_p, enable_autodiff)
         else:
-            return _perstot(X, order, internal_p), np.array([[i, -1] for i in range(n)])
-
+            return _perstot(X, order, internal_p, enable_autodiff), np.array([[i, -1] for i in range(n)])
+
+    if enable_autodiff:
+        import eagerpy as ep
+        X_orig = ep.astensor(X)
+        Y_orig = ep.astensor(Y)
+        X = X_orig.numpy()
+        Y = Y_orig.numpy()
     M = _build_dist_matrix(X, Y, order=order, internal_p=internal_p)
     a = np.ones(n+1) # weight vector of the input diagram. Uniform here.
     a[-1] = m
@@ -106,6 +129,7 @@ def wasserstein_distance(X, Y, matching=False, order=2., internal_p=2.):
     b[-1] = n
 
     if matching:
+        assert not enable_autodiff, "matching and enable_autodiff are currently incompatible"
         P = ot.emd(a=a,b=b,M=M, numItermax=2000000)
         ot_cost = np.sum(np.multiply(P,M))
         P[-1, -1] = 0  # Remove matching corresponding to the diagonal
@@ -115,6 +139,24 @@ def wasserstein_distance(X, Y, matching=False, order=2., internal_p=2.):
         match[:,1][match[:,1] >= m] = -1
         return ot_cost ** (1./order) , match
 
+    if enable_autodiff:
+        P = ot.emd(a=a,b=b,M=M, numItermax=2000000)
+        pairs = np.argwhere(P[:-1, :-1])
+        diag2 = np.nonzero(P[-1, :-1])
+        diag1 = np.nonzero(P[:-1, -1])
+        dists = []
+        # empty arrays are not handled properly by the helpers, so we avoid calling them
+        if len(pairs):
+            dists.append((Y_orig[pairs[:, 1]] - X_orig[pairs[:, 0]]).norms.lp(internal_p, axis=-1).norms.lp(order))
+        if len(diag1):
+            dists.append(_perstot_autodiff(X_orig[diag1], order, internal_p))
+        if len(diag2):
+            dists.append(_perstot_autodiff(Y_orig[diag2], order, internal_p))
+        dists = [ dist.reshape(1) for dist in dists ]
+        return ep.concatenate(dists).norms.lp(order)
+        # Should just compute the L^order norm manually?
+        # We can also concatenate the 3 vectors to compute just one norm.
+
     # Comptuation of the otcost using the ot.emd2 library.
     # Note: it is the Wasserstein distance to the power q.
     # The default numItermax=100000 is not sufficient for some examples with 5000 points, what is a good value?