adding essential parts management in wasserstein distance

1 files changed, 138 insertions, 8 deletions

diff --git a/src/python/gudhi/wasserstein/wasserstein.py b/src/python/gudhi/wasserstein/wasserstein.py
index b37d30bb..283ecd9d 100644
--- a/src/python/gudhi/wasserstein/wasserstein.py
+++ b/src/python/gudhi/wasserstein/wasserstein.py
@@ -79,6 +79,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 +infty if the diagram has an essential part (points with infinite coordinates).
     '''
     if enable_autodiff:
         import eagerpy as ep
@@ -88,32 +91,136 @@ 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 a: (n x 2) numpy.array (point of a diagram)
+    :retuns: 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.
+    '''
+    if len(a):
+        ess_first_type  = np.where(np.isfinite(a[:,1]) & (a[:,0] == -np.inf))[0]  # coord (-inf, x)
+        ess_second_type = np.where(np.isfinite(a[:,0]) & (a[:,1] == np.inf))[0]  # coord (x, +inf)
+        ess_third_type  = np.where((a[:,0] == -np.inf) & (a[:,1] == np.inf))[0]      # coord (-inf, +inf)
+        ess_fourth_type = np.where((a[:,0] == -np.inf) & (a[:,1] == -np.inf))[0]    # coord (-inf, -inf)
+        ess_fifth_type  = np.where((a[:,0] == np.inf)  & (a[:,1] == np.inf))[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 distanc ecomputation
+    :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.
     '''
-    :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).
+    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, and then [i, -1] (or [-1, j]) to account for essential points matched to the diagonal
+    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.
+    '''
+    c = 0
+    m = []
+
+    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 zip(ess_parts_X, ess_parts_Y):
+        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 type >= 2 (both coordinates are infinite, so we essentially just align points)
+    for u, v in zip(ess_parts_X[2:], ess_parts_Y[2:]):
+        m += list(zip(u, v))  # cost is 0
+
+    return c, np.array(m)
+
+
+def _offdiag(X):
+    '''
+    :param X: (n x 2) numpy array encoding a persistence diagram.
+    :returns: The off-diagonal part of a diagram `X` (points with finite coordinates).
+    '''
+    return X[np.where(np.isfinite(X[:,0]) & np.isfinite(X[:,1]))]
+
+
+def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enable_autodiff=False,
+                         keep_essential_parts=True):
+    '''
+    :param X: (n x 2) numpy.array encoding the first diagram. Can now contain essential parts (points with infinite
+                        coordinates).
     :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.
+                     Note that if the cost is +inf (essential parts have different number of points,
+                     then the optimal matching will be set to `None`.
     :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`.
+                       default value is `np.inf`.
     :param enable_autodiff: If X and Y are torch.tensor, tensorflow.Tensor or jax.numpy.ndarray, 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
             and lets the various frameworks compute its gradient. It never pulls new points from the diagonal.
     :type enable_autodiff: bool
+    :param keep_essential_parts: If False, only considers the off-diagonal points in the diagrams.
+                                 Otherwise, computes the distance between the essential parts separately.
+    :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.
     '''
+    # Zeroth step: check compatibility of arguments
+    if keep_essential_parts and enable_autodiff:
+        import warnings
+        warnings.warn("enable_autodiff does not handle essential parts yet. These will be ignored in the following computations")
+        keep_essential_parts = False
+
+    # First step: handle empty diagrams
     n = len(X)
     m = len(Y)
 
-    # handle empty diagrams
     if n == 0:
         if m == 0:
             if not matching:
@@ -132,6 +239,25 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab
         else:
             return _perstot(X, order, internal_p, enable_autodiff), np.array([[i, -1] for i in range(n)])
 
+
+    # Second step: handle essential parts
+    if keep_essential_parts:
+        essential_cost, essential_matching = _handle_essential_parts(X, Y, order=order)
+        if (essential_cost == np.inf):
+            if matching:
+                return np.inf, None
+            else:
+                return np.inf  # avoid computing off-diagonal transport cost if essential parts do not match (saves time)
+
+    else:
+        essential_cost = 0
+        essential_matching = None
+
+    X, Y = _offdiag(X), _offdiag(Y)
+    n = len(X)
+    m = len(Y)
+
+    # Now the standard pipeline for off-diagonal parts
     if enable_autodiff:
         import eagerpy as ep
 
@@ -139,6 +265,7 @@ 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()
+
     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 +281,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)
@@ -178,4 +308,4 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab
     # 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)