diff options
Diffstat (limited to 'src/python')
-rw-r--r-- | src/python/doc/wasserstein_distance_user.rst | 24 | ||||
-rw-r--r-- | src/python/gudhi/wasserstein/wasserstein.py | 159 | ||||
-rwxr-xr-x | src/python/test/test_wasserstein_distance.py | 91 |
3 files changed, 258 insertions, 16 deletions
diff --git a/src/python/doc/wasserstein_distance_user.rst b/src/python/doc/wasserstein_distance_user.rst index 9ffc2759..b3d17495 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,30 @@ An index of -1 represents the diagonal. for j in dgm2_to_diagonal: print("point %s in dgm2 is matched to the diagonal" %j) + 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 ----------- diff --git a/src/python/gudhi/wasserstein/wasserstein.py b/src/python/gudhi/wasserstein/wasserstein.py index a9d1cdff..572d4249 100644 --- a/src/python/gudhi/wasserstein/wasserstein.py +++ b/src/python/gudhi/wasserstein/wasserstein.py @@ -70,6 +70,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 +80,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 +92,137 @@ 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. + ''' + 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, enable_autodiff): + ''' + :param X: (n x 2) numpy array encoding a persistence diagram. + :param enable_autodiff: boolean, to handle the case where X is a eagerpy tensor. + :returns: The off-diagonal part of a diagram `X` (points with finite coordinates). + ''' + if enable_autodiff: + # Assumes the diagrams only have finite coordinates. Thus, return X directly. + # TODO improve this to get rid of essential parts if there are any. + return X + else: + 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 (finite points of the) first diagram. Must not contain essential points - (i.e. with infinite coordinate). + :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`. :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 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. ''' + + # First step: handle empty diagrams n = len(X) m = len(Y) - # handle empty diagrams if n == 0: if m == 0: if not matching: @@ -125,13 +234,41 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab if not matching: return _perstot(Y, order, internal_p, enable_autodiff) else: - return _perstot(Y, order, internal_p, enable_autodiff), np.array([[-1, j] for j in range(m)]) + cost = _perstot(Y, order, internal_p, enable_autodiff) + if cost == np.inf: # We had some essential part here. + return cost, None + 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) else: - return _perstot(X, order, internal_p, enable_autodiff), np.array([[i, -1] for i in range(n)]) + cost = _perstot(X, order, internal_p, enable_autodiff) + if cost == np.inf: + return cost, None + else: + return cost, 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 + + # Extract off-diaognal points of the diagrams. Note that if enable_autodiff is True, nothing is done here (X,Y are + # assumed to be tensors with only finite coordinates). + X, Y = _offdiag(X, enable_autodiff), _offdiag(Y, enable_autodiff) + n = len(X) + m = len(Y) + + # Now the standard pipeline for off-diagonal parts if enable_autodiff: import eagerpy as ep @@ -139,6 +276,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 +292,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 +319,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) diff --git a/src/python/test/test_wasserstein_distance.py b/src/python/test/test_wasserstein_distance.py index e3b521d6..6701c7ba 100755 --- a/src/python/test/test_wasserstein_distance.py +++ b/src/python/test/test_wasserstein_distance.py @@ -5,25 +5,79 @@ 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, _offdiag, _handle_essential_parts, _get_essential_parts 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_offdiag(): + 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(_offdiag(diag, enable_autodiff=False), [[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], + [-np.inf, 8], [-np.inf, 11], + [-np.inf, -np.inf], [-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. + assert np.array_equal(m, [[4, 4], [5, 5], [2, 2], [3, 3], [8, 8], [9, 9], [6, 6], [7, 7]]) + + c, m = _handle_essential_parts(diag1, diag3, order=1) + assert c == np.inf + assert (m is None) + + +def test_get_essential_parts(): + 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]]) + + res = _get_essential_parts(diag) + 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] ) + + 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 +118,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 +132,13 @@ 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]]) + + 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) def hera_wrap(**extra): @@ -92,8 +153,32 @@ def pot_wrap(**extra): 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_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) + +def test_wasserstein_distance_grad(): + import torch + + diag1 = torch.tensor([[2.7, 3.7], [9.6, 14.0], [34.2, 34.974]], requires_grad=True) + diag2 = torch.tensor([[2.8, 4.45], [9.5, 14.1]], requires_grad=True) + diag3 = torch.tensor([[2.8, 4.45], [9.5, 14.1]], requires_grad=True) + assert diag1.grad is None and diag2.grad is None and diag3.grad is None + dist12 = pot(diag1, diag2, internal_p=2, order=2, enable_autodiff=True, keep_essential_parts=False) + dist30 = pot(diag3, torch.tensor([]), internal_p=2, order=2, enable_autodiff=True, keep_essential_parts=False) + dist12.backward() + dist30.backward() + assert not torch.isnan(diag1.grad).any() and not torch.isnan(diag2.grad).any() and not torch.isnan(diag3.grad).any() + diag4 = torch.tensor([[0., 10.]], requires_grad=True) + diag5 = torch.tensor([[1., 11.], [3., 4.]], requires_grad=True) + dist45 = pot(diag4, diag5, internal_p=1, order=1, enable_autodiff=True, keep_essential_parts=False) + assert dist45 == 3. + dist45.backward() + assert np.array_equal(diag4.grad, [[-1., -1.]]) + assert np.array_equal(diag5.grad, [[1., 1.], [-1., 1.]]) + diag6 = torch.tensor([[5., 10.]], requires_grad=True) + pot(diag6, diag6, internal_p=2, order=2, enable_autodiff=True, keep_essential_parts=False).backward() + # https://github.com/jonasrauber/eagerpy/issues/6 + # assert np.array_equal(diag6.grad, [[0., 0.]]) |