Merge pull request #365 from tlacombe/essential_part_in_wasserstein

Essential part in wasserstein

3 files changed, 326 insertions, 34 deletions

diff --git a/src/python/doc/wasserstein_distance_user.rst b/src/python/doc/wasserstein_distance_user.rst
index 9ffc2759..76eb1469 100644
--- a/src/python/doc/wasserstein_distance_user.rst
+++ b/src/python/doc/wasserstein_distance_user.rst
@@ -44,7 +44,7 @@ Basic example
 *************
 
 This example computes the 1-Wasserstein distance from 2 persistence diagrams with Euclidean ground metric.
-Note that persistence diagrams must be submitted as (n x 2) numpy arrays and must not contain inf values.
+Note that persistence diagrams must be submitted as (n x 2) numpy arrays.
 
 .. testcode::
 
@@ -67,14 +67,16 @@ We can also have access to the optimal matching by letting `matching=True`.
 It is encoded as a list of indices (i,j), meaning that the i-th point in X
 is mapped to the j-th point in Y.
 An index of -1 represents the diagonal.
+It handles essential parts (points with infinite coordinates). However if the cardinalities of the essential parts differ, 
+any matching has a cost +inf and thus can be considered to be optimal. In such a case, the function returns `(np.inf, None)`.
 
 .. testcode::
 
     import gudhi.wasserstein
     import numpy as np
 
-    dgm1 = np.array([[2.7, 3.7],[9.6, 14.],[34.2, 34.974]])
-    dgm2 = np.array([[2.8, 4.45], [5, 6], [9.5, 14.1]])
+    dgm1 = np.array([[2.7, 3.7],[9.6, 14.],[34.2, 34.974], [3, np.inf]])
+    dgm2 = np.array([[2.8, 4.45], [5, 6], [9.5, 14.1], [4, np.inf]])
     cost, matchings = gudhi.wasserstein.wasserstein_distance(dgm1, dgm2, matching=True, order=1, internal_p=2)
 
     message_cost = "Wasserstein distance value = %.2f" %cost
@@ -90,16 +92,31 @@ An index of -1 represents the diagonal.
     for j in dgm2_to_diagonal:
         print("point %s in dgm2 is matched to the diagonal" %j)
 
-The output is:
+    # An example where essential part cardinalities differ
+    dgm3 = np.array([[1, 2], [0, np.inf]])
+    dgm4 = np.array([[1, 2], [0, np.inf], [1, np.inf]])
+    cost, matchings = gudhi.wasserstein.wasserstein_distance(dgm3, dgm4, matching=True, order=1, internal_p=2)
+    print("\nSecond example:")
+    print("cost:", cost)
+    print("matchings:", matchings)
+
+
+The output is: 
 
 .. testoutput::
 
-    Wasserstein distance value = 2.15
+    Wasserstein distance value = 3.15
     point 0 in dgm1 is matched to point 0 in dgm2
     point 1 in dgm1 is matched to point 2 in dgm2
+    point 3 in dgm1 is matched to point 3 in dgm2
     point 2 in dgm1 is matched to the diagonal
     point 1 in dgm2 is matched to the diagonal
 
+    Second example:
+    cost: inf
+    matchings: None
+
+
 Barycenters
 -----------
 
@@ -181,4 +198,4 @@ Tutorial
 
 This
 `notebook <https://github.com/GUDHI/TDA-tutorial/blob/master/Tuto-GUDHI-Barycenters-of-persistence-diagrams.ipynb>`_
-presents the concept of barycenter, or Fréchet mean, of a family of persistence diagrams.
\ No newline at end of file
+presents the concept of barycenter, or Fréchet mean, of a family of persistence diagrams.
diff --git a/src/python/gudhi/wasserstein/wasserstein.py b/src/python/gudhi/wasserstein/wasserstein.py
index a9d1cdff..dc18806e 100644
--- a/src/python/gudhi/wasserstein/wasserstein.py
+++ b/src/python/gudhi/wasserstein/wasserstein.py
@@ -9,6 +9,7 @@
 
 import numpy as np
 import scipy.spatial.distance as sc
+import warnings
 
 try:
     import ot
@@ -70,6 +71,7 @@ def _perstot_autodiff(X, order, internal_p):
     '''
     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).
@@ -79,6 +81,9 @@ def _perstot(X, order, internal_p, enable_autodiff):
         transparent to automatic differentiation.
     :type enable_autodiff: bool
     :returns: float, the total persistence of the diagram (that is, its distance to the empty diagram).
+
+    .. note::
+        Can be +inf if the diagram has an essential part (points with infinite coordinates).
     '''
     if enable_autodiff:
         import eagerpy as ep
@@ -88,32 +93,163 @@ def _perstot(X, order, internal_p, enable_autodiff):
         return np.linalg.norm(_dist_to_diag(X, internal_p), ord=order)
 
 
-def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enable_autodiff=False):
+def _get_essential_parts(a):
     '''
-    :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 matching: if True, computes and returns the optimal matching between X and Y, encoded as
-                     a (n x 2) np.array  [...[i,j]...], meaning the i-th point in X is matched to
-                     the j-th point in Y, with the convention (-1) represents the diagonal.
-    :param order: exponent for Wasserstein; Default value is 1.
-    :param internal_p: Ground metric on the (upper-half) plane (i.e. norm L^p in R^2);
-                       Default value is `np.inf`.
-    :param enable_autodiff: If X and Y are torch.tensor or tensorflow.Tensor, make the computation
+    :param a: (n x 2) numpy.array (point of a diagram)
+    :returns: five lists of indices (between 0 and len(a)) accounting for the five types of points with infinite
+    coordinates that can occur in a diagram, namely:
+        type0 : (-inf, finite)
+        type1 : (finite, +inf)
+        type2 : (-inf, +inf)
+        type3 : (-inf, -inf)
+        type4 : (+inf, +inf)
+    .. note::
+        For instance, a[_get_essential_parts(a)[0]] returns the points in a of coordinates (-inf, x) for some finite x.
+        Note also that points with (+inf, -inf) are not handled (points (x,y) in dgm satisfy by assumption (y >= x)).
+
+        Finally, we consider that points with coordinates (-inf,-inf) and (+inf, +inf) belong to the diagonal.
+    '''
+    if len(a):
+        first_coord_finite = np.isfinite(a[:,0])
+        second_coord_finite = np.isfinite(a[:,1])
+        first_coord_infinite_positive = (a[:,0] == np.inf)
+        second_coord_infinite_positive = (a[:,1] == np.inf)
+        first_coord_infinite_negative = (a[:,0] == -np.inf)
+        second_coord_infinite_negative = (a[:,1] == -np.inf)
+
+        ess_first_type  = np.where(second_coord_finite & first_coord_infinite_negative)[0] # coord (-inf, x)
+        ess_second_type = np.where(first_coord_finite & second_coord_infinite_positive)[0]  # coord (x, +inf)
+        ess_third_type  = np.where(first_coord_infinite_negative & second_coord_infinite_positive)[0]  # coord (-inf, +inf)
+
+        ess_fourth_type = np.where(first_coord_infinite_negative & second_coord_infinite_negative)[0] # coord (-inf, -inf)
+        ess_fifth_type  = np.where(first_coord_infinite_positive  & second_coord_infinite_positive)[0]  # coord (+inf, +inf)
+        return ess_first_type, ess_second_type, ess_third_type, ess_fourth_type, ess_fifth_type
+    else:
+        return [], [], [], [], []
+
+
+def _cost_and_match_essential_parts(X, Y, idX, idY, order, axis):
+    '''
+    :param X: (n x 2) numpy.array (dgm points)
+    :param Y: (n x 2) numpy.array (dgm points)
+    :param idX: indices to consider for this one dimensional OT problem (in X)
+    :param idY: indices to consider for this one dimensional OT problem (in Y)
+    :param order: exponent for Wasserstein distance computation
+    :param axis: must be 0 or 1, correspond to the coordinate which is finite.
+    :returns: cost (float) and match for points with *one* infinite coordinate.
+
+    .. note::
+        Assume idX, idY come when calling _handle_essential_parts, thus have same length.
+    '''
+    u = X[idX, axis]
+    v = Y[idY, axis]
+
+    cost = np.sum(np.abs(np.sort(u) - np.sort(v))**(order))  # OT cost in 1D
+
+    sortidX = idX[np.argsort(u)]
+    sortidY = idY[np.argsort(v)]
+    # We return [i,j] sorted per value
+    match = list(zip(sortidX, sortidY))
+
+    return cost, match
+
+
+def _handle_essential_parts(X, Y, order):
+    '''
+    :param X: (n x 2) numpy array, first diagram.
+    :param Y: (n x 2) numpy array, second diagram.
+    :order: Wasserstein order for cost computation.
+    :returns: cost and matching due to essential parts. If cost is +inf, matching will be set to None.
+    '''
+    ess_parts_X = _get_essential_parts(X)
+    ess_parts_Y = _get_essential_parts(Y)
+
+    # Treats the case of infinite cost (cardinalities of essential parts differ).
+    for u, v in list(zip(ess_parts_X, ess_parts_Y))[:3]: # ignore types 4 and 5 as they belong to the diagonal
+        if len(u) != len(v):
+            return np.inf, None
+
+    # Now we know each essential part has the same number of points in both diagrams.
+    # Handle type 0 and type 1 essential parts (those with one finite coordinates)
+    c1, m1 = _cost_and_match_essential_parts(X, Y, ess_parts_X[0], ess_parts_Y[0], axis=1, order=order)
+    c2, m2 = _cost_and_match_essential_parts(X, Y, ess_parts_X[1], ess_parts_Y[1], axis=0, order=order)
+
+    c = c1 + c2
+    m = m1 + m2
+
+    # Handle type3 (coordinates (-inf,+inf), so we just align points)
+    m += list(zip(ess_parts_X[2], ess_parts_Y[2]))
+
+    # Handle type 4 and 5, considered as belonging to the diagonal so matched to (-1) with cost 0.
+    for z in ess_parts_X[3:]:
+        m += [(u, -1) for u in z] # points in X are matched to -1
+    for z in ess_parts_Y[3:]:
+        m += [(-1, v) for v in z] # -1 is match to points in Y
+
+    return c, np.array(m)
+
+
+def _finite_part(X):
+    '''
+    :param X: (n x 2) numpy array encoding a persistence diagram.
+    :returns: The finite part of a diagram `X` (points with finite coordinates).
+    '''
+    return X[np.where(np.isfinite(X[:,0]) & np.isfinite(X[:,1]))]
+
+
+def _warn_infty(matching):
+    '''
+    Handle essential parts with different cardinalities. Warn the user about cost being infinite and (if
+    `matching=True`) about the returned matching being `None`.
+    '''
+    if matching:
+        warnings.warn('Cardinality of essential parts differs. Distance (cost) is +inf, and the returned matching is None.')
+        return np.inf, None
+    else:
+        warnings.warn('Cardinality of essential parts differs. Distance (cost) is +inf.')
+        return np.inf
+
+
+def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enable_autodiff=False,
+                         keep_essential_parts=True):
+    '''
+    Compute the Wasserstein distance between persistence diagram using Python Optimal Transport backend.
+    Diagrams can contain points with infinity coordinates (essential parts).
+    Points with (-inf,-inf) and (+inf,+inf) coordinates are considered as belonging to the diagonal.
+    If the distance between two diagrams is +inf (which happens if the cardinalities of essential
+    parts differ) and optimal matching is required, it will be set to ``None``.
+
+    :param X: The first diagram.
+    :type X: n x 2 numpy.array
+    :param Y:  The second diagram.
+    :type Y: m x 2 numpy.array
+    :param matching: if ``True``, computes and returns the optimal matching between X and Y, encoded as
+        a (n x 2) np.array  [...[i,j]...], meaning the i-th point in X is matched to
+        the j-th point in Y, with the convention that (-1) represents the diagonal.
+    :param order: Wasserstein exponent q (1 <= q < infinity).
+    :type order: float
+    :param internal_p: Ground metric on the (upper-half) plane (i.e. norm L^p in R^2).
+    :type internal_p: float
+    :param enable_autodiff: If X and Y are ``torch.tensor`` or ``tensorflow.Tensor``, make the computation
         transparent to automatic differentiation. This requires the package EagerPy and is currently incompatible
-        with `matching=True`.
+        with ``matching=True`` and with ``keep_essential_parts=True``.
 
-        .. note:: This considers the function defined on the coordinates of the off-diagonal points of X and Y
+        .. note:: This considers the function defined on the coordinates of the off-diagonal finite points of X and Y
             and lets the various frameworks compute its gradient. It never pulls new points from the diagonal.
     :type enable_autodiff: bool
-    :returns: the Wasserstein distance of order q (1 <= q < infinity) between persistence diagrams with
+    :param keep_essential_parts: If ``False``, only considers the finite points in the diagrams.
+                                 Otherwise, include essential parts in cost and matching computation.
+    :type keep_essential_parts: 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.
+              If cost is +inf, any matching is optimal and thus it returns `None` instead.
     '''
+
+    # First step: handle empty diagrams
     n = len(X)
     m = len(Y)
 
-    # handle empty diagrams
     if n == 0:
         if m == 0:
             if not matching:
@@ -122,16 +258,45 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab
             else:
                 return 0., np.array([])
         else:
-            if not matching:
-                return _perstot(Y, order, internal_p, enable_autodiff)
+            cost = _perstot(Y, order, internal_p, enable_autodiff)
+            if cost == np.inf:
+                return _warn_infty(matching)
             else:
-                return _perstot(Y, order, internal_p, enable_autodiff), np.array([[-1, j] for j in range(m)])
+                if not matching:
+                    return cost
+                else:
+                    return cost, np.array([[-1, j] for j in range(m)])
     elif m == 0:
-        if not matching:
-            return _perstot(X, order, internal_p, enable_autodiff)
+        cost = _perstot(X, order, internal_p, enable_autodiff)
+        if cost == np.inf:
+            return _warn_infty(matching)
         else:
-            return _perstot(X, order, internal_p, enable_autodiff), np.array([[i, -1] for i in range(n)])
+            if not matching:
+                return cost
+            else:
+                return cost, np.array([[i, -1] for i in range(n)])
+
 
+    # Check essential part and enable autodiff together
+    if enable_autodiff and keep_essential_parts:
+        warnings.warn('''enable_autodiff=True and keep_essential_parts=True are incompatible together.
+                      keep_essential_parts is set to False: only points with finite coordinates are considered
+                      in the following.
+                      ''')
+        keep_essential_parts = False
+
+    # Second step: handle essential parts if needed.
+    if keep_essential_parts:
+        essential_cost, essential_matching = _handle_essential_parts(X, Y, order=order)
+        if (essential_cost == np.inf):
+            return _warn_infty(matching)  # Tells the user that cost is infty and matching (if True) is None.
+            # avoid computing transport cost between the finite parts if essential parts
+            # cardinalities do not match (saves time)
+    else:
+        essential_cost = 0
+        essential_matching = None
+
+    # Now the standard pipeline for finite parts
     if enable_autodiff:
         import eagerpy as ep
 
@@ -139,6 +304,12 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab
         Y_orig = ep.astensor(Y)
         X = X_orig.numpy()
         Y = Y_orig.numpy()
+
+    # Extract finite points of the diagrams. 
+    X, Y = _finite_part(X), _finite_part(Y)
+    n = len(X)
+    m = len(Y)
+
     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
@@ -154,7 +325,10 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab
         # Now we turn to -1 points encoding the diagonal
         match[:,0][match[:,0] >= n] = -1
         match[:,1][match[:,1] >= m] = -1
-        return ot_cost ** (1./order) , match
+        # Finally incorporate the essential part matching
+        if essential_matching is not None:
+            match = np.concatenate([match, essential_matching]) if essential_matching.size else match
+        return (ot_cost + essential_cost) ** (1./order) , match
 
     if enable_autodiff:
         P = ot.emd(a=a, b=b, M=M, numItermax=2000000)
@@ -173,9 +347,9 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab
         return ep.concatenate(dists).norms.lp(order).raw
         # We can also concatenate the 3 vectors to compute just one norm.
 
-    # Comptuation of the otcost using the ot.emd2 library.
+    # Comptuation of the ot cost 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?
     ot_cost = ot.emd2(a, b, M, numItermax=2000000)
 
-    return ot_cost ** (1./order)
+    return (ot_cost + essential_cost) ** (1./order)
diff --git a/src/python/test/test_wasserstein_distance.py b/src/python/test/test_wasserstein_distance.py
index e3b521d6..3a004d77 100755
--- a/src/python/test/test_wasserstein_distance.py
+++ b/src/python/test/test_wasserstein_distance.py
@@ -5,25 +5,97 @@
     Copyright (C) 2019 Inria
 
     Modification(s):
+      - 2020/07 Théo Lacombe: Added tests about handling essential parts in diagrams.
       - YYYY/MM Author: Description of the modification
 """
 
-from gudhi.wasserstein.wasserstein import _proj_on_diag
+from gudhi.wasserstein.wasserstein import _proj_on_diag, _finite_part, _handle_essential_parts, _get_essential_parts
+from gudhi.wasserstein.wasserstein import _warn_infty
 from gudhi.wasserstein import wasserstein_distance as pot
 from gudhi.hera import wasserstein_distance as hera
 import numpy as np
 import pytest
 
+
 __author__ = "Theo Lacombe"
 __copyright__ = "Copyright (C) 2019 Inria"
 __license__ = "MIT"
 
+
 def test_proj_on_diag():
     dgm = np.array([[1., 1.], [1., 2.], [3., 5.]])
     assert np.array_equal(_proj_on_diag(dgm), [[1., 1.], [1.5, 1.5], [4., 4.]])
     empty = np.empty((0, 2))
     assert np.array_equal(_proj_on_diag(empty), empty)
 
+
+def test_finite_part():
+    diag = np.array([[0, 1], [3, 5], [2, np.inf], [3, np.inf], [-np.inf, 8], [-np.inf, 12], [-np.inf, -np.inf],
+                     [np.inf, np.inf], [-np.inf, np.inf], [-np.inf, np.inf]])
+    assert np.array_equal(_finite_part(diag), [[0, 1], [3, 5]])
+
+
+def test_handle_essential_parts():
+    diag1 = np.array([[0, 1], [3, 5],
+                      [2, np.inf], [3, np.inf],
+                      [-np.inf, 8], [-np.inf, 12],
+                      [-np.inf, -np.inf],
+                      [np.inf, np.inf],
+                      [-np.inf, np.inf], [-np.inf, np.inf]])
+
+    diag2 = np.array([[0, 2], [3, 5],
+                      [2, np.inf], [4, np.inf],
+                      [-np.inf, 8], [-np.inf, 11],
+                      [-np.inf, -np.inf],
+                      [np.inf, np.inf],
+                      [-np.inf, np.inf], [-np.inf, np.inf]])
+
+    diag3 = np.array([[0, 2], [3, 5],
+                      [2, np.inf], [4, np.inf], [6, np.inf],
+                      [-np.inf, 8], [-np.inf, 11],
+                      [-np.inf, -np.inf],
+                      [np.inf, np.inf],
+                      [-np.inf, np.inf], [-np.inf, np.inf]])
+
+    c, m = _handle_essential_parts(diag1, diag2, order=1)
+    assert c == pytest.approx(2, 0.0001)  # Note: here c is only the cost due to essential part (thus 2, not 3)
+    # Similarly, the matching only corresponds to essential parts.
+    # Note that (-inf,-inf) and (+inf,+inf) coordinates are matched to the diagonal.
+    assert np.array_equal(m, [[4, 4], [5, 5], [2, 2], [3, 3], [8, 8], [9, 9], [6, -1], [7, -1], [-1, 6], [-1, 7]])
+
+    c, m = _handle_essential_parts(diag1, diag3, order=1)
+    assert c == np.inf
+    assert (m is None)
+
+
+def test_get_essential_parts():
+    diag1 = np.array([[0, 1], [3, 5], [2, np.inf], [3, np.inf], [-np.inf, 8], [-np.inf, 12], [-np.inf, -np.inf],
+                     [np.inf, np.inf], [-np.inf, np.inf], [-np.inf, np.inf]])
+
+    diag2 = np.array([[0, 1], [3, 5], [2, np.inf], [3, np.inf]])
+
+    res  = _get_essential_parts(diag1)
+    res2 = _get_essential_parts(diag2)
+    assert np.array_equal(res[0], [4, 5])
+    assert np.array_equal(res[1], [2, 3])
+    assert np.array_equal(res[2], [8, 9])
+    assert np.array_equal(res[3], [6]   )
+    assert np.array_equal(res[4], [7]   )
+
+    assert np.array_equal(res2[0], []    )
+    assert np.array_equal(res2[1], [2, 3])
+    assert np.array_equal(res2[2], []    )
+    assert np.array_equal(res2[3], []    )
+    assert np.array_equal(res2[4], []    )
+
+
+def test_warn_infty():
+    assert _warn_infty(matching=False)==np.inf
+    c, m = _warn_infty(matching=True)
+    assert (c == np.inf)
+    assert (m is None)
+
+
 def _basic_wasserstein(wasserstein_distance, delta, test_infinity=True, test_matching=True):
     diag1 = np.array([[2.7, 3.7], [9.6, 14.0], [34.2, 34.974]])
     diag2 = np.array([[2.8, 4.45], [9.5, 14.1]])
@@ -64,7 +136,7 @@ def _basic_wasserstein(wasserstein_distance, delta, test_infinity=True, test_mat
 
         assert wasserstein_distance(diag4, diag5) == np.inf
         assert wasserstein_distance(diag5, diag6, order=1, internal_p=np.inf) == approx(4.)
-
+        assert wasserstein_distance(diag5, emptydiag) == np.inf
 
     if test_matching:
         match = wasserstein_distance(emptydiag, emptydiag, matching=True, internal_p=1., order=2)[1]
@@ -78,6 +150,31 @@ def _basic_wasserstein(wasserstein_distance, delta, test_infinity=True, test_mat
         match = wasserstein_distance(diag1, diag2, matching=True, internal_p=2., order=2.)[1]
         assert np.array_equal(match, [[0, 0], [1, 1], [2, -1]])
 
+    if test_matching and test_infinity:
+        diag7 = np.array([[0, 3], [4, np.inf], [5, np.inf]])
+        diag8 = np.array([[0,1], [0, np.inf], [-np.inf, -np.inf], [np.inf, np.inf]])
+        diag9 = np.array([[-np.inf, -np.inf], [np.inf, np.inf]])
+        diag10 = np.array([[0,1], [-np.inf, -np.inf], [np.inf, np.inf]])
+
+        match = wasserstein_distance(diag5, diag6, matching=True, internal_p=2., order=2.)[1]
+        assert np.array_equal(match, [[0, -1], [-1,0], [-1, 1], [1, 2]])
+        match = wasserstein_distance(diag5, diag7, matching=True, internal_p=2., order=2.)[1]
+        assert (match is None)
+        cost, match = wasserstein_distance(diag7, emptydiag, matching=True, internal_p=2., order=2.3)
+        assert (cost == np.inf)
+        assert (match is None)
+        cost, match = wasserstein_distance(emptydiag, diag7, matching=True, internal_p=2.42, order=2.)
+        assert (cost == np.inf)
+        assert (match is None)
+        cost, match = wasserstein_distance(diag8, diag9, matching=True, internal_p=2., order=2.)
+        assert (cost == np.inf)
+        assert (match is None)
+        cost, match = wasserstein_distance(diag9, diag10, matching=True, internal_p=1., order=1.)
+        assert (cost == 1)
+        assert (match == [[0, -1],[1, -1],[-1, 0], [-1, 1], [-1, 2]]) # type 4 and 5 are match to the diag anyway.
+        cost, match = wasserstein_distance(diag9, emptydiag, matching=True, internal_p=2., order=2.)
+        assert (cost == 0.)
+        assert (match == [[0, -1], [1, -1]])
 
 
 def hera_wrap(**extra):
@@ -85,15 +182,19 @@ def hera_wrap(**extra):
         return hera(*kargs,**kwargs,**extra)
     return fun
 
+
 def pot_wrap(**extra):
     def fun(*kargs,**kwargs):
         return pot(*kargs,**kwargs,**extra)
     return fun
 
+
 def test_wasserstein_distance_pot():
-    _basic_wasserstein(pot, 1e-15, test_infinity=False, test_matching=True)
-    _basic_wasserstein(pot_wrap(enable_autodiff=True), 1e-15, test_infinity=False, test_matching=False)
+    _basic_wasserstein(pot, 1e-15, test_infinity=False, test_matching=True)  # pot with its standard args
+    _basic_wasserstein(pot_wrap(enable_autodiff=True, keep_essential_parts=False), 1e-15, test_infinity=False, test_matching=False)
+
 
 def test_wasserstein_distance_hera():
     _basic_wasserstein(hera_wrap(delta=1e-12), 1e-12, test_matching=False)
     _basic_wasserstein(hera_wrap(delta=.1), .1, test_matching=False)
+