FIX refine stop criteria

Author: MatthewSZhang

fastcan/_refine.py

@@ -120,7 +120,7 @@ def refine(selector, drop=1, max_iter=None, verbose=1):
     indices_temp = best_indices
     for drop_n in drop:
         i = 0
-        while i < n_features:
+        while i < n_features_to_select:
             rolled_indices = np.r_[
                 indices_include, np.roll(indices_temp[n_inclusions:], -1)
             ]

tests/test_refine.py

@@ -6,7 +6,7 @@
 from fastcan import FastCan, refine
 
 
-def test_select_refine_cls():
+def test_select_refine_cls(capsys):
     # Test whether refine work correctly with random samples.
     n_samples = 200
     n_features = 20
@@ -30,6 +30,8 @@ def test_select_refine_cls():
 
     selector = FastCan(n_to_select).fit(X, y)
     _, scores_1 = refine(selector, drop=1)
+    captured = capsys.readouterr()
+    assert "No. of iterations: 19," in captured.out
     _, scores_23 = refine(selector, drop=[2, 3], verbose=0)
     _, scores_all = refine(selector, drop="all", max_iter=20, verbose=1)