fix: boolean selection masks in pandas eval

pysr/export_numpy.py

@@ -1,10 +1,12 @@
 """Code for exporting discovered expressions to numpy"""
 
 import warnings
+from typing import List, Union
 
 import numpy as np
 import pandas as pd
-from sympy import lambdify
+from numpy.typing import NDArray
+from sympy import Expr, Symbol, lambdify
 
 
 def sympy2numpy(eqn, sympy_symbols, *, selection=None):
@@ -14,6 +16,10 @@ def sympy2numpy(eqn, sympy_symbols, *, selection=None):
 class CallableEquation:
     """Simple wrapper for numpy lambda functions built with sympy"""
 
+    _sympy: Expr
+    _sympy_symbols: List[Symbol]
+    _selection: Union[NDArray[np.bool_], None]
+
     def __init__(self, eqn, sympy_symbols, selection=None):
         self._sympy = eqn
         self._sympy_symbols = sympy_symbols
@@ -29,8 +35,9 @@ class CallableEquation:
             return self._lambda(
                 **{k: X[k].values for k in map(str, self._sympy_symbols)}
             ) * np.ones(expected_shape)
+
         if self._selection is not None:
-            if X.shape[1] != len(self._selection):
+            if X.shape[1] != self._selection.sum():
                 warnings.warn(
                     "`X` should be of shape (n_samples, len(self._selection)). "
                     "Automatically filtering `X` to selection. "
@@ -38,6 +45,7 @@ class CallableEquation:
                     "this may lead to incorrect predictions and other errors."
                 )
                 X = X[:, self._selection]
+
         return self._lambda(*X.T) * np.ones(expected_shape)
 
     @property

pysr/sr.py

@@ -2056,7 +2056,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             if self.selection_mask_ is not None:
                 # RangeIndex enforces column order allowing columns to
                 # be correctly filtered with self.selection_mask_
-                X = X.iloc[:, self.selection_mask_]
+                X = X[X.columns[self.selection_mask_]]
             X.columns = self.feature_names_in_
         # Without feature information, CallableEquation/lambda_format equations
         # require that the column order of X matches that of the X used during

pysr/test/test.py

@@ -526,7 +526,7 @@ class TestFeatureSelection(unittest.TestCase):
         X = self.rstate.randn(20000, 5)
         y = X[:, 2] ** 2 + X[:, 3] ** 2
         selected = run_feature_selection(X, y, select_k_features=2)
-        self.assertEqual(sorted(selected), [2, 3])
+        np.testing.assert_array_equal(selected, [False, False, True, True, False])
 
     def test_feature_selection_handler(self):
         X = self.rstate.randn(20000, 5)
@@ -538,8 +538,8 @@ class TestFeatureSelection(unittest.TestCase):
             variable_names=var_names,
             y=y,
         )
-        self.assertTrue((2 in selection) and (3 in selection))
-        selected_var_names = [var_names[i] for i in selection]
+        np.testing.assert_array_equal(selection, [False, False, True, True, False])
+        selected_var_names = [var_names[i] for i in range(5) if selection[i]]
         self.assertEqual(set(selected_var_names), set("x2 x3".split(" ")))
         np.testing.assert_array_equal(
             np.sort(selected_X, axis=1), np.sort(X[:, [2, 3]], axis=1)