Fix unit-test

Signed-off-by: Adam Li <[email protected]>
neurodata · Sep 28, 2023 · dbf079e · dbf079e
1 parent dc6dd29
commit dbf079e
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 2 deletions.
diff --git a/benchmarks_nonasv/notebooks/compare_coleman_and_permutation_forest.ipynb b/benchmarks_nonasv/notebooks/compare_coleman_and_permutation_forest.ipynb
@@ -1345,7 +1345,41 @@
       "(125, 600, 1) (125, 600, 1)\n",
       "X500:  1.0\n",
       "(125, 600, 1) (125, 600, 1)\n",
-      "X2:  0.004975124378109453\n"
+      "X2:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X1:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X500:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X2:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X1:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X500:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X2:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X1:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X500:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X2:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X1:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X500:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X2:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X1:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X500:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X2:  0.004975124378109453\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X1:  1.0\n",
+      "(125, 600, 1) (125, 600, 1)\n",
+      "X500:  1.0\n"
      ]
     }
    ],

diff --git a/sktree/stats/forestht.py b/sktree/stats/forestht.py
@@ -785,8 +785,10 @@ def _statistic(
                 )
 
         # determine if there are any nans in the final posterior array
+        temp_posterior_forest = np.nanmean(posterior_arr, axis=0)
+        nonnan_indices = np.where(~np.isnan(temp_posterior_forest).any(axis=1))[0]
+
         # Find the row indices with NaN values in any column
-        nonnan_indices = np.all(~np.isnan(posterior_arr), axis=(0, 2))
         samples = nonnan_indices
 
         # Ignore all NaN values (samples not tested)