fix: check for NaNs in emd loss matrix

ot/lp/__init__.py

@@ -237,6 +237,8 @@ def emd(a, b, M, numItermax=100000, log=False, center_dual=True, numThreads=1, c
 
     .. note:: An error will be raised if the vectors :math:`\mathbf{a}` and :math:`\mathbf{b}` do not sum to the same value.
 
+    .. note:: An error will be raided if the loss matrix :math:`\mathbf{M}` contains NaNs.
+
     Uses the algorithm proposed in :ref:`[1] <references-emd>`.
 
     Parameters
@@ -302,6 +304,9 @@ def emd(a, b, M, numItermax=100000, log=False, center_dual=True, numThreads=1, c
     ot.optim.cg : General regularized OT
     """
 
+    if np.isnan(M).any():
+        raise ValueError('The loss matrix should not contain NaN values.')
+
     a, b, M = list_to_array(a, b, M)
     nx = get_backend(M, a, b)
 

test/gromov/test_gw.py

@@ -832,3 +832,33 @@ def test_fgw_barycenter(nx):
     # test correspondance with utils function
     recovered_C = ot.gromov.update_kl_loss(p, lambdas, log['T'], [C1, C2])
     np.testing.assert_allclose(C, recovered_C)
+
+
+# Related to issue 469
+def test_gromov2_nan_in_source_cost():
+    # GIVEN a source cost matrix with a NaN value
+    source_cost = np.zeros((2, 2))
+    target_cost = np.ones((2, 2))
+    source_distribution = np.array([0.5, 0.5])
+    target_distribution = np.array([0.5, 0.5])
+
+    source_cost[0, 0] = np.nan
+
+    # WHEN we call gromov_wasserstein2 - THEN we expect a ValueError
+    with pytest.raises(ValueError, match='The loss matrix should not contain NaN values.'):
+        ot.gromov_wasserstein2(source_cost, target_cost, source_distribution, target_distribution)
+
+
+# Related to issue 469
+def test_gromov2_nan_in_target_cost():
+    # GIVEN - a target cost matrix with a NaN value
+    source_cost = np.zeros((2, 2))
+    target_cost = np.ones((2, 2))
+    source_distribution = np.array([0.5, 0.5])
+    target_distribution = np.array([0.5, 0.5])
+
+    target_cost[0, 0] = np.nan
+
+    # WHEN - we call
+    with pytest.raises(ValueError, match='The loss matrix should not contain NaN values.'):
+        ot.gromov_wasserstein2(source_cost, target_cost, source_distribution, target_distribution)

test/test_ot.py

@@ -245,6 +245,11 @@ def test_emd_empty():
     np.testing.assert_allclose(w, 0)
 
 
+def test_emd_nan_in_loss_matrix():
+    with pytest.raises(ValueError, match='The loss matrix should not contain NaN values.'):
+        ot.emd([], [], [np.nan])
+
+
 def test_emd2_multi():
     n = 500  # nb bins