message: Simplify distance-to-diagonal in Wasserstein
diff: @@ -15,16 +15,19 @@ try:
  except ImportError:
  print("POT (Python Optimal Transport) package is not installed. Try to run $ conda install -c conda-forge pot ; or $ pip install POT")
  -def _proj_on_diag(X):
  +def _dist_to_diag(X, internal_p):
  '''
  :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
  + :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.
  '''
  - Z = (X[:,0] + X[:,1]) / 2.
  - return np.array([Z , Z]).T
  + return (X[:, 1] - X[:, 0]) * 2 ** (1.0 / internal_p - 1)
  -def _build_dist_matrix(X, Y, order=2., internal_p=2.):
  +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.
  @@ -36,16 +39,12 @@ def _build_dist_matrix(X, Y, order=2., internal_p=2.):
  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).
  '''
  - Xdiag = _proj_on_diag(X)
  - Ydiag = _proj_on_diag(Y)
  + 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
  - Cxd = np.linalg.norm(X - Xdiag, ord=internal_p, axis=1)**order
  - Cdy = np.linalg.norm(Y - Ydiag, ord=internal_p, axis=1)**order
  else:
  C = sc.cdist(X,Y, metric='minkowski', p=internal_p)**order
  - Cxd = np.linalg.norm(X - Xdiag, ord=internal_p, axis=1)**order
  - Cdy = np.linalg.norm(Y - Ydiag, ord=internal_p, axis=1)**order
  Cf = np.hstack((C, Cxd[:,None]))
  Cdy = np.append(Cdy, 0)
  @@ -61,8 +60,7 @@ def _perstot(X, order, internal_p):
  :param internal_p: Ground metric on the (upper-half) plane (i.e. norm L^p in R^2); Default value is 2 (Euclidean norm).
  :returns: float, the total persistence of the diagram (that is, its distance to the empty diagram).
  '''
  - Xdiag = _proj_on_diag(X)
  - return (np.sum(np.linalg.norm(X - Xdiag, ord=internal_p, axis=1)**order))**(1./order)
  + return np.linalg.norm(_dist_to_diag(X, internal_p), ord=order)
  def wasserstein_distance(X, Y, matching=False, order=2., internal_p=2.):
  