1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
|
# This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
# See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
# Author(s): Theo Lacombe
#
# Copyright (C) 2019 Inria
#
# Modification(s):
# - YYYY/MM Author: Description of the modification
import numpy as np
import scipy.spatial.distance as sc
try:
import ot
except ImportError:
print("POT (Python Optimal Transport) package is not installed. Try to run $ conda install -c conda-forge pot ; or $ pip install POT")
# Currently unused, but Théo says it is likely to be used again.
def _proj_on_diag(X):
'''
:param X: (n x 2) array encoding the points of a persistent diagram.
:returns: (n x 2) array encoding the (respective orthogonal) projections of the points onto the diagonal
'''
Z = (X[:,0] + X[:,1]) / 2.
return np.array([Z , Z]).T
def _dist_to_diag(X, internal_p):
'''
:param X: (n x 2) array encoding the points of a persistent diagram.
:param internal_p: Ground metric (i.e. norm L^p).
:returns: (n) array encoding the (respective orthogonal) distances of the points to the diagonal
.. note::
Assumes that the points are above the diagonal.
'''
return (X[:, 1] - X[:, 0]) * 2 ** (1.0 / internal_p - 1)
def _build_dist_matrix(X, Y, order, internal_p):
'''
: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 order: exponent for the Wasserstein metric.
:param internal_p: Ground metric (i.e. norm L^p).
:returns: (n+1) x (m+1) np.array encoding the cost matrix C.
For 0 <= i < n, 0 <= j < m, C[i,j] encodes the distance between X[i] and Y[j],
while C[i, m] (resp. C[n, j]) encodes the distance (to the p) between X[i] (resp Y[j])
and its orthogonal projection onto the diagonal.
note also that C[n, m] = 0 (it costs nothing to move from the diagonal to the diagonal).
'''
Cxd = _dist_to_diag(X, internal_p)**order
Cdy = _dist_to_diag(Y, internal_p)**order
if np.isinf(internal_p):
C = sc.cdist(X,Y, metric='chebyshev')**order
else:
C = sc.cdist(X,Y, metric='minkowski', p=internal_p)**order
Cf = np.hstack((C, Cxd[:,None]))
Cdy = np.append(Cdy, 0)
Cf = np.vstack((Cf, Cdy[None,:]))
return Cf
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.
:param internal_p: Ground metric on the (upper-half) plane (i.e. norm L^p in R^2).
: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).
'''
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=1., internal_p=np.inf, 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).
: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, 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`.
.. 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
: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.
'''
n = len(X)
m = len(Y)
# handle empty diagrams
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, enable_autodiff)
else:
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, enable_autodiff)
else:
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
b = np.ones(m+1) # weight vector of the input diagram. Uniform here.
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
match = np.argwhere(P)
# 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
if enable_autodiff:
P = ot.emd(a=a, b=b, M=M, numItermax=2000000)
pairs_X_Y = np.argwhere(P[:-1, :-1])
pairs_X_diag = np.nonzero(P[:-1, -1])
pairs_Y_diag = np.nonzero(P[-1, :-1])
dists = []
# empty arrays are not handled properly by the helpers, so we avoid calling them
if len(pairs_X_Y):
dists.append((Y_orig[pairs_X_Y[:, 1]] - X_orig[pairs_X_Y[:, 0]]).norms.lp(internal_p, axis=-1).norms.lp(order))
if len(pairs_X_diag):
dists.append(_perstot_autodiff(X_orig[pairs_X_diag], order, internal_p))
if len(pairs_Y_diag):
dists.append(_perstot_autodiff(Y_orig[pairs_Y_diag], order, internal_p))
dists = [dist.reshape(1) for dist in dists]
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.
# 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)
|
