Wasserstein is back on track now

Author: ctralie

RELEASE.txt

@@ -1,3 +1,7 @@
+0.3.1
+    - Fixed bug with repeated intervals in bottleneck
+    - Tidied up API for indicating matchings for bottleneck and wasserstein, and updated notebook
+
 0.3.0
     - Add implementations of Persistence Landscapes, including plotting methods, a transformer, and additional notebooks.
 

persim/_version.py

@@ -1 +1 @@
-__version__ = "0.3.0"
+__version__ = "0.3.1"

persim/wasserstein.py

@@ -81,7 +81,7 @@ def wasserstein(dgm1, dgm2, matching=False):
     R = np.array([[cp, -sp], [sp, cp]])
     S = S[:, 0:2].dot(R)
     T = T[:, 0:2].dot(R)
-    D = np.inf*np.ones((M+N, M+N))
+    D = np.zeros((M+N, M+N))
     np.fill_diagonal(D, 0)
     D[0:M, 0:N] = DUL
     UR = np.inf*np.ones((M, M))
@@ -90,7 +90,6 @@ def wasserstein(dgm1, dgm2, matching=False):
     UL = np.inf*np.ones((N, N))
     np.fill_diagonal(UL, T[:, 1])
     D[M:N+M, 0:N] = UL
-    print(D)
 
     # Step 2: Run the hungarian algorithm
     matchi, matchj = optimize.linear_sum_assignment(D)
@@ -104,6 +103,8 @@ def wasserstein(dgm1, dgm2, matching=False):
         # Indicate diagonally matched points
         ret[ret[:, 0] >= M, 0] = -1
         ret[ret[:, 1] >= N, 1] = -1
+        # Exclude diagonal to diagonal
+        ret = ret[ret[:, 0] + ret[:, 1] != -2, :] 
         return matchdist, ret
 
     return matchdist

test/test_distances.py

@@ -183,10 +183,32 @@ def test_inf_deathtime(self):
         assert (np.allclose(dist1, np.sqrt(2)/2)) and (np.allclose(dist2, np.sqrt(2)/2))
 
     def test_repeated(self):
-        dgm1 = np.array([[0, 10], [0,10], [0,11]])
+        dgm1 = np.array([[0, 10], [0,10]])
         dgm2 = np.array([[0, 10]])
         dist = wasserstein(dgm1, dgm2)
-        assert dist == 10.5
+        assert dist == 5*np.sqrt(2)
+
+    def test_matching(self):
+        dgm1 = np.array([
+            [0.5, 1],
+            [0.6, 1.1]
+        ])
+        dgm2 = np.array([
+            [0.5, 1.1],
+            [0.6, 1.1],
+            [0.8, 1.1],
+            [1.0, 1.1],
+        ])
+
+        d, m = wasserstein(
+            dgm1, dgm2,
+            matching=True
+        )
+        u1 = np.unique(m[:, 0])
+        u1 = u1[u1 >= 0]
+        u2 = np.unique(m[:, 1])
+        u2 = u2[u2 >= 0]
+        assert u1.size == dgm1.shape[0] and u2.size == dgm2.shape[0]
 
 
 class TestSliced: