Merge pull request #192 from tlacombe/wdist

Wdist

4 files changed, 44 insertions, 45 deletions

diff --git a/src/python/doc/wasserstein_distance_sum.inc b/src/python/doc/wasserstein_distance_sum.inc
index ffd4d312..a97f428d 100644
--- a/src/python/doc/wasserstein_distance_sum.inc
+++ b/src/python/doc/wasserstein_distance_sum.inc
@@ -2,12 +2,12 @@
    :widths: 30 50 20
 
    +-----------------------------------------------------------------+----------------------------------------------------------------------+------------------------------------------------------------------+
-   | .. figure::                                                     | The p-Wasserstein distance measures the similarity between two       | :Author: Theo Lacombe                                            |
+   | .. figure::                                                     | The q-Wasserstein distance measures the similarity between two       | :Author: Theo Lacombe                                            |
    |      ../../doc/Bottleneck_distance/perturb_pd.png               | persistence diagrams. It's the minimum value c that can be achieved  |                                                                  |
    |      :figclass: align-center                                    | by a perfect matching between the points of the two diagrams (+ all  | :Introduced in: GUDHI 3.1.0                                      |
    |                                                                 | diagonal points), where the value of a matching is defined as the    |                                                                  |
-   |      Wasserstein distance is the p-th root of the sum of the    | p-th root of the sum of all edge lengths to the power p. Edge lengths| :Copyright: MIT                                                  |
-   |      edge lengths to the power p.                               | are measured in norm q, for :math:`1 \leq q \leq \infty`.            |                                                                  |
+   |      Wasserstein distance is the q-th root of the sum of the    | q-th root of the sum of all edge lengths to the power q. Edge lengths| :Copyright: MIT                                                  |
+   |      edge lengths to the power q.                               | are measured in norm p, for :math:`1 \leq p \leq \infty`.            |                                                                  |
    |                                                                 |                                                                      | :Requires: Python Optimal Transport (POT) :math:`\geq` 0.5.1     |
    +-----------------------------------------------------------------+----------------------------------------------------------------------+------------------------------------------------------------------+
    | * :doc:`wasserstein_distance_user`                              |                                                                                                                                         |
diff --git a/src/python/doc/wasserstein_distance_user.rst b/src/python/doc/wasserstein_distance_user.rst
index a049cfb5..32999a0c 100644
--- a/src/python/doc/wasserstein_distance_user.rst
+++ b/src/python/doc/wasserstein_distance_user.rst
@@ -30,7 +30,7 @@ Note that persistence diagrams must be submitted as (n x 2) numpy arrays and mus
     diag1 = np.array([[2.7, 3.7],[9.6, 14.],[34.2, 34.974]])
     diag2 = np.array([[2.8, 4.45],[9.5, 14.1]])
 
-    message = "Wasserstein distance value = " + '%.2f' % gudhi.wasserstein.wasserstein_distance(diag1, diag2, q=2., p=1.)
+    message = "Wasserstein distance value = " + '%.2f' % gudhi.wasserstein.wasserstein_distance(diag1, diag2, order=1., internal_p=2.)
     print(message)
 
 The output is:
diff --git a/src/python/gudhi/wasserstein.py b/src/python/gudhi/wasserstein.py
index 1906eeb1..db5ddff2 100644
--- a/src/python/gudhi/wasserstein.py
+++ b/src/python/gudhi/wasserstein.py
@@ -23,26 +23,26 @@ def _proj_on_diag(X):
     return np.array([Z , Z]).T
 
 
-def _build_dist_matrix(X, Y, p=2., q=2.):
+def _build_dist_matrix(X, Y, order=2., internal_p=2.):
     '''
     :param X: (n x 2) numpy.array encoding the (points of the) first diagram.
     :param Y: (m x 2) numpy.array encoding the second diagram.
-    :param q: Ground metric (i.e. norm l_q).
-    :param p: exponent for the Wasserstein metric.
+    :param internal_p: Ground metric (i.e. norm l_p).
+    :param order: exponent for the Wasserstein metric.
     :returns: (n+1) x (m+1) np.array encoding the cost matrix C. 
                 For 1 <= i <= n, 1 <= j <= m, C[i,j] encodes the distance between X[i] and Y[j], while C[i, m+1] (resp. C[n+1, j]) encodes the distance (to the p) between X[i] (resp Y[j]) and its orthogonal proj onto the diagonal.
                 note also that C[n+1, m+1] = 0  (it costs nothing to move from the diagonal to the diagonal).
     '''
     Xdiag = _proj_on_diag(X)
     Ydiag = _proj_on_diag(Y)
-    if np.isinf(q):
-        C = sc.cdist(X,Y, metric='chebyshev')**p
-        Cxd = np.linalg.norm(X - Xdiag, ord=q, axis=1)**p
-        Cdy = np.linalg.norm(Y - Ydiag, ord=q, axis=1)**p
+    if np.isinf(internal_p):
+        C = sc.cdist(X,Y, metric='chebyshev')**order
+        Cxd = np.linalg.norm(X - Xdiag, ord=internal_p, axis=1)**order
+        Cdy = np.linalg.norm(Y - Ydiag, ord=internal_p, axis=1)**order
     else:
-        C = sc.cdist(X,Y, metric='minkowski', p=q)**p
-        Cxd = np.linalg.norm(X - Xdiag, ord=q, axis=1)**p
-        Cdy = np.linalg.norm(Y - Ydiag, ord=q, axis=1)**p
+        C = sc.cdist(X,Y, metric='minkowski', p=internal_p)**order
+        Cxd = np.linalg.norm(X - Xdiag, ord=internal_p, axis=1)**order
+        Cdy = np.linalg.norm(Y - Ydiag, ord=internal_p, axis=1)**order
     Cf = np.hstack((C, Cxd[:,None]))
     Cdy = np.append(Cdy, 0)
 
@@ -51,24 +51,24 @@ def _build_dist_matrix(X, Y, p=2., q=2.):
     return Cf
 
 
-def _perstot(X, p, q):
+def _perstot(X, order, internal_p):
     '''
     :param X: (n x 2) numpy.array (points of a given diagram).
-    :param q: Ground metric on the (upper-half) plane (i.e. norm l_q in R^2); Default value is 2 (Euclidean norm).
-    :param p: 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 order: exponent for Wasserstein. Default value is 2.
     :returns: float, the total persistence of the diagram (that is, its distance to the empty diagram).    
     '''
     Xdiag = _proj_on_diag(X)
-    return (np.sum(np.linalg.norm(X - Xdiag, ord=q, axis=1)**p))**(1./p)
+    return (np.sum(np.linalg.norm(X - Xdiag, ord=internal_p, axis=1)**order))**(1./order)
 
 
-def wasserstein_distance(X, Y, p=2., q=2.):
+def wasserstein_distance(X, Y, order=2., internal_p=2.):
     '''
     :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).
     :param Y: (m x 2) numpy.array encoding the second diagram.
-    :param q: Ground metric on the (upper-half) plane (i.e. norm l_q in R^2); Default value is 2 (euclidean norm).
-    :param p: exponent for Wasserstein; Default value is 2.
-    :returns: the p-Wasserstein distance (1 <= p < infinity) with respect to the q-norm as ground metric.
+    :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 order: exponent for Wasserstein; Default value is 2.
+    :returns: the Wasserstein distance of order q (1 <= q < infinity) between persistence diagrams with respect to the internal_p-norm as ground metric.
     :rtype: float
     '''
     n = len(X)
@@ -79,20 +79,19 @@ def wasserstein_distance(X, Y, p=2., q=2.):
         if Y.size == 0:
             return 0.
         else:
-            return _perstot(Y, p, q)
+            return _perstot(Y, order, internal_p)
     elif Y.size == 0:
-        return _perstot(X, p, q)
+        return _perstot(X, order, internal_p)
 
-    M = _build_dist_matrix(X, Y, p=p, q=q)
+    M = _build_dist_matrix(X, Y, order=order, internal_p=internal_p)
     a = np.full(n+1, 1. / (n + m) )  # weight vector of the input diagram. Uniform here.
     a[-1] = a[-1] * m                # normalized so that we have a probability measure, required by POT
     b = np.full(m+1, 1. / (n + m) )  # weight vector of the input diagram. Uniform here.
     b[-1] = b[-1] * n                # so that we have a probability measure, required by POT
 
     # Comptuation of the otcost using the ot.emd2 library.
-    # Note: it is the squared Wasserstein distance.
+    # 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?
     ot_cost = (n+m) * ot.emd2(a, b, M, numItermax=2000000)
 
-    return ot_cost ** (1./p)
-
+    return ot_cost ** (1./order)
diff --git a/src/python/test/test_wasserstein_distance.py b/src/python/test/test_wasserstein_distance.py
index bb59b3c2..43dda77e 100755
--- a/src/python/test/test_wasserstein_distance.py
+++ b/src/python/test/test_wasserstein_distance.py
@@ -23,26 +23,26 @@ def test_basic_wasserstein():
     diag4 = np.array([[0, 3], [4, 8]])
     emptydiag = np.array([[]])
 
-    assert wasserstein_distance(emptydiag, emptydiag, q=2., p=1.) == 0.
-    assert wasserstein_distance(emptydiag, emptydiag, q=np.inf, p=1.) == 0.
-    assert wasserstein_distance(emptydiag, emptydiag, q=np.inf, p=2.) == 0.
-    assert wasserstein_distance(emptydiag, emptydiag, q=2., p=2.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, internal_p=2.,     order=1.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, internal_p=np.inf, order=1.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, internal_p=np.inf, order=2.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, internal_p=2.,     order=2.) == 0.
 
-    assert wasserstein_distance(diag3, emptydiag, q=np.inf, p=1.) == 2.
-    assert wasserstein_distance(diag3, emptydiag, q=1., p=1.) == 4.
+    assert wasserstein_distance(diag3, emptydiag, internal_p=np.inf,     order=1.) == 2.
+    assert wasserstein_distance(diag3, emptydiag, internal_p=1.,         order=1.) == 4.
 
-    assert wasserstein_distance(diag4, emptydiag, q=1., p=2.) == 5.  # thank you Pythagorician triplets
-    assert wasserstein_distance(diag4, emptydiag, q=np.inf, p=2.) == 2.5
-    assert wasserstein_distance(diag4, emptydiag, q=2., p=2.) == 3.5355339059327378
+    assert wasserstein_distance(diag4, emptydiag, internal_p=1.,     order=2.) == 5.  # thank you Pythagorician triplets
+    assert wasserstein_distance(diag4, emptydiag, internal_p=np.inf, order=2.) == 2.5
+    assert wasserstein_distance(diag4, emptydiag, internal_p=2.,     order=2.) == 3.5355339059327378
 
-    assert wasserstein_distance(diag1, diag2, q=2., p=1.) == 1.4453593023967701
-    assert wasserstein_distance(diag1, diag2, q=2.35, p=1.74) == 0.9772734057168739
+    assert wasserstein_distance(diag1, diag2, internal_p=2.,   order=1.)   == 1.4453593023967701
+    assert wasserstein_distance(diag1, diag2, internal_p=2.35, order=1.74) == 0.9772734057168739
 
-    assert wasserstein_distance(diag1, emptydiag, q=2.35, p=1.7863) == 3.141592214572228
+    assert wasserstein_distance(diag1, emptydiag, internal_p=2.35, order=1.7863) == 3.141592214572228
 
-    assert wasserstein_distance(diag3, diag4, q=1., p=1.) == 3.
-    assert wasserstein_distance(diag3, diag4, q=np.inf, p=1.) == 3.  # no diag matching here
-    assert wasserstein_distance(diag3, diag4, q=np.inf, p=2.) == np.sqrt(5)
-    assert wasserstein_distance(diag3, diag4, q=1., p=2.) == np.sqrt(5)
-    assert wasserstein_distance(diag3, diag4, q=4.5, p=2.) == np.sqrt(5)
+    assert wasserstein_distance(diag3, diag4, internal_p=1.,     order=1.) == 3.
+    assert wasserstein_distance(diag3, diag4, internal_p=np.inf, order=1.) == 3.  # no diag matching here
+    assert wasserstein_distance(diag3, diag4, internal_p=np.inf, order=2.) == np.sqrt(5)
+    assert wasserstein_distance(diag3, diag4, internal_p=1.,     order=2.) == np.sqrt(5)
+    assert wasserstein_distance(diag3, diag4, internal_p=4.5,    order=2.) == np.sqrt(5)