Share tests for wasserstein_distance

1 files changed, 42 insertions, 17 deletions

diff --git a/src/python/test/test_wasserstein_distance.py b/src/python/test/test_wasserstein_distance.py
index 43dda77e..46a7079f 100755
--- a/src/python/test/test_wasserstein_distance.py
+++ b/src/python/test/test_wasserstein_distance.py
@@ -8,41 +8,66 @@
       - YYYY/MM Author: Description of the modification
 """
 
-from gudhi.wasserstein import wasserstein_distance
+from gudhi.wasserstein import wasserstein_distance as pot
+from gudhi.hera import wasserstein_distance as hera
 import numpy as np
 
 __author__ = "Theo Lacombe"
 __copyright__ = "Copyright (C) 2019 Inria"
 __license__ = "MIT"
 
-
-def test_basic_wasserstein():
+def _basic_wasserstein(wasserstein_distance, delta, test_infinity=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]])
     diag3 = np.array([[0, 2], [4, 6]])
     diag4 = np.array([[0, 3], [4, 8]])
-    emptydiag = np.array([[]])
+    emptydiag = np.array([])
+
+    # We just need to handle positive numbers here
+    def approx(a, b):
+        f = 1 + delta
+        return a <= b*f and b <= a*f
 
     assert wasserstein_distance(emptydiag, emptydiag, internal_p=2.,     order=1.) == 0.
     assert wasserstein_distance(emptydiag, emptydiag, internal_p=np.inf, order=1.) == 0.
     assert wasserstein_distance(emptydiag, emptydiag, internal_p=np.inf, order=2.) == 0.
     assert wasserstein_distance(emptydiag, emptydiag, internal_p=2.,     order=2.) == 0.
 
-    assert wasserstein_distance(diag3, emptydiag, internal_p=np.inf,     order=1.) == 2.
-    assert wasserstein_distance(diag3, emptydiag, internal_p=1.,         order=1.) == 4.
+    assert approx(wasserstein_distance(diag3, emptydiag, internal_p=np.inf,     order=1.), 2.)
+    assert approx(wasserstein_distance(diag3, emptydiag, internal_p=1.,         order=1.), 4.)
+
+    assert approx(wasserstein_distance(diag4, emptydiag, internal_p=1.,     order=2.), 5.)  # thank you Pythagorician triplets
+    assert approx(wasserstein_distance(diag4, emptydiag, internal_p=np.inf, order=2.), 2.5)
+    assert approx(wasserstein_distance(diag4, emptydiag, internal_p=2.,     order=2.), 3.5355339059327378)
+
+    assert approx(wasserstein_distance(diag1, diag2, internal_p=2.,   order=1.)  , 1.4453593023967701)
+    assert approx(wasserstein_distance(diag1, diag2, internal_p=2.35, order=1.74), 0.9772734057168739)
+
+    assert approx(wasserstein_distance(diag1, emptydiag, internal_p=2.35, order=1.7863), 3.141592214572228)
+
+    assert approx(wasserstein_distance(diag3, diag4, internal_p=1.,     order=1.), 3.)
+    assert approx(wasserstein_distance(diag3, diag4, internal_p=np.inf, order=1.), 3.)  # no diag matching here
+    assert approx(wasserstein_distance(diag3, diag4, internal_p=np.inf, order=2.), np.sqrt(5))
+    assert approx(wasserstein_distance(diag3, diag4, internal_p=1.,     order=2.), np.sqrt(5))
+    assert approx(wasserstein_distance(diag3, diag4, internal_p=4.5,    order=2.), np.sqrt(5))
+
+    if(not test_infinity):
+        return
 
-    assert wasserstein_distance(diag4, emptydiag, internal_p=1.,     order=2.) == 5.  # thank you Pythagorician triplets
-    assert wasserstein_distance(diag4, emptydiag, internal_p=np.inf, order=2.) == 2.5
-    assert wasserstein_distance(diag4, emptydiag, internal_p=2.,     order=2.) == 3.5355339059327378
+    diag5 = np.array([[0, 3], [4, np.inf]])
+    diag6 = np.array([[7, 8], [4, 6], [3, np.inf]])
 
-    assert wasserstein_distance(diag1, diag2, internal_p=2.,   order=1.)   == 1.4453593023967701
-    assert wasserstein_distance(diag1, diag2, internal_p=2.35, order=1.74) == 0.9772734057168739
+    assert wasserstein_distance(diag4, diag5) == np.inf
+    assert approx(wasserstein_distance(diag5, diag6, order=1, internal_p=np.inf), 4.)
 
-    assert wasserstein_distance(diag1, emptydiag, internal_p=2.35, order=1.7863) == 3.141592214572228
+def hera_wrap(delta):
+    def fun(*kargs,**kwargs):
+        return hera(*kargs,**kwargs,delta=delta)
+    return fun
 
-    assert wasserstein_distance(diag3, diag4, internal_p=1.,     order=1.) == 3.
-    assert wasserstein_distance(diag3, diag4, internal_p=np.inf, order=1.) == 3.  # no diag matching here
-    assert wasserstein_distance(diag3, diag4, internal_p=np.inf, order=2.) == np.sqrt(5)
-    assert wasserstein_distance(diag3, diag4, internal_p=1.,     order=2.) == np.sqrt(5)
-    assert wasserstein_distance(diag3, diag4, internal_p=4.5,    order=2.) == np.sqrt(5)
+def test_wasserstein_distance_pot():
+    _basic_wasserstein(pot, 1e-15, False)
 
+def test_wasserstein_distance_hera():
+    _basic_wasserstein(hera_wrap(1e-12), 1e-12)
+    _basic_wasserstein(hera_wrap(.1), .1)