Merge pull request #116 from kilianFatras/emd_dimension

[MRG] Add assert for emd dimension mismatch

Author: rflamary

ot/lp/__init__.py

@@ -114,6 +114,9 @@ def emd(a, b, M, numItermax=100000, log=False, dense=True):
     if len(b) == 0:
         b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1]
 
+    assert (a.shape[0] == M.shape[0] and b.shape[0] == M.shape[1]), \
+        "Dimension mismatch, check dimensions of M with a and b"
+
     if dense:
         G, cost, u, v, result_code = emd_c(a, b, M, numItermax,dense)
     else:
@@ -226,6 +229,9 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(),
     if len(b) == 0:
         b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1]
 
+    assert (a.shape[0] == M.shape[0] and b.shape[0] == M.shape[1]), \
+        "Dimension mismatch, check dimensions of M with a and b"
+
     if log or return_matrix:
         def f(b):
             if dense:

test/test_ot.py

@@ -14,6 +14,22 @@
 import pytest
 
 
+def test_emd_dimension_mismatch():
+    # test emd and emd2 for dimension mismatch
+    n_samples = 100
+    n_features = 2
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n_samples, n_features)
+    a = ot.utils.unif(n_samples + 1)
+
+    M = ot.dist(x, x)
+
+    np.testing.assert_raises(AssertionError, ot.emd, a, a, M)
+
+    np.testing.assert_raises(AssertionError, ot.emd2, a, a, M)
+
+
 def test_emd_emd2():
     # test emd and emd2 for simple identity
     n = 100