Adding test for issue scikit-learn#4297, isotonic infinite loop

mjbommar · mjbommar · commit da04cd998d85 · 2015-02-27T18:02:14.000-05:00
diff --git a/sklearn/tests/test_isotonic.py b/sklearn/tests/test_isotonic.py
@@ -308,6 +308,25 @@ def test_isotonic_duplicate_min_entry():
     all_predictions_finite = np.all(np.isfinite(ir.predict(x)))
     assert_true(all_predictions_finite)
 
+
+def test_isotonic_zero_weight_loop():
+    # Test from @ogrisel's issue:
+    # https://github.com/scikit-learn/scikit-learn/issues/4297
+
+    regression = IsotonicRegression()
+    n_samples = 50
+    x = np.linspace(-3, 3, n_samples)
+    y = x + np.random.uniform(size=n_samples)
+
+    # Get some random weights and zero out
+    w = np.random.uniform(size=n_samples)
+    w[5:8] = 0
+    regression.fit(x, y, sample_weight=w)
+
+    # This will hang in failure case.
+    regression.fit(x, y, sample_weight=w)
+
+
 if __name__ == "__main__":
     import nose
     nose.run(argv=['', __file__])
