Fix underflow in nmf._beta_divergence (#10142)

TomDLT · agramfort · commit 2160ea0a74bf · 2017-11-16T18:31:15.000+01:00
diff --git a/sklearn/decomposition/nmf.py b/sklearn/decomposition/nmf.py
@@ -114,8 +114,9 @@ def _beta_divergence(X, W, H, beta, square_root=False):
         X_data = X.ravel()
 
     # do not affect the zeros: here 0 ** (-1) = 0 and not infinity
-    WH_data = WH_data[X_data != 0]
-    X_data = X_data[X_data != 0]
+    indices = X_data > EPSILON
+    WH_data = WH_data[indices]
+    X_data = X_data[indices]
 
     # used to avoid division by zero
     WH_data[WH_data == 0] = EPSILON
diff --git a/sklearn/decomposition/tests/test_nmf.py b/sklearn/decomposition/tests/test_nmf.py
@@ -479,3 +479,18 @@ def test_nmf_decreasing():
                 if previous_loss is not None:
                     assert_greater(previous_loss, loss)
                 previous_loss = loss
+
+
+def test_nmf_underflow():
+    # Regression test for an underflow issue in _beta_divergence
+    rng = np.random.RandomState(0)
+    n_samples, n_features, n_components = 10, 2, 2
+    X = np.abs(rng.randn(n_samples, n_features)) * 10
+    W = np.abs(rng.randn(n_samples, n_components)) * 10
+    H = np.abs(rng.randn(n_components, n_features))
+
+    X[0, 0] = 0
+    ref = nmf._beta_divergence(X, W, H, beta=1.0)
+    X[0, 0] = 1e-323
+    res = nmf._beta_divergence(X, W, H, beta=1.0)
+    assert_almost_equal(res, ref)
