Resolve "Adding new datasets"

Closes #17 (closed)

code/bolsonaro/models/omp_forest.py

@@ -33,6 +33,8 @@ class OmpForest(BaseEstimator, metaclass=ABCMeta):
 
     # sklearn baseestimator api methods
     def fit(self, X_forest, y_forest, X_omp, y_omp):
+        # print(y_forest.shape)
+        # print(set([type(y) for y in y_forest]))
         self._base_forest_estimator.fit(X_forest, y_forest)
         self._extract_subforest(X_omp, y_omp) # type: OrthogonalMatchingPursuit
         return self
@@ -140,8 +142,8 @@ class SingleOmpForest(OmpForest):
             forest_predictions /= self._forest_norms
 
         weights = self._omp.coef_
-        omp_trees_indices = np.nonzero(weights)
+        omp_trees_indices = np.nonzero(weights)[0]
 
         select_trees = np.mean(forest_predictions[omp_trees_indices], axis=0)
-
+        print(len(omp_trees_indices))
         return select_trees