diff options
Diffstat (limited to 'src/python/gudhi/wasserstein')
-rw-r--r-- | src/python/gudhi/wasserstein/wasserstein.py | 16 |
1 files changed, 6 insertions, 10 deletions
diff --git a/src/python/gudhi/wasserstein/wasserstein.py b/src/python/gudhi/wasserstein/wasserstein.py index 495142c4..142385b1 100644 --- a/src/python/gudhi/wasserstein/wasserstein.py +++ b/src/python/gudhi/wasserstein/wasserstein.py @@ -178,13 +178,13 @@ def _handle_essential_parts(X, Y, order): 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: - import eagerpy as ep - - return ep.astensor(X[np.where(np.isfinite(X[:,0]) & np.isfinite(X[:,1]))]) - + # 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]))] @@ -218,11 +218,6 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab 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. keep_essential_parts set to False.") - keep_essential_parts = False # First step: handle empty diagrams n = len(X) @@ -267,7 +262,8 @@ def wasserstein_distance(X, Y, matching=False, order=1., internal_p=np.inf, enab essential_cost = 0 essential_matching = None - # Extract off-diaognal points of the diagrams. + # 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) |