feat: allow 1D input

MilesCranmer · Aug 26, 2024 · 66542ff · 66542ff
1 parent 0c9515e
commit 66542ff
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/pysr/regressor_sequence.py b/pysr/regressor_sequence.py
@@ -131,7 +131,10 @@ def fit(
         self : object
             Fitted estimator.
         """
-        X = self._validate_data(X)
+        X = self._validate_data(X, ensure_2d=False)
+        if X.ndim == 1:
+            X = X.reshape(-1, 1)
+        assert X.ndim == 2
         _check_assertions(
             X,
             self.recursive_history_length,
@@ -195,7 +198,10 @@ def predict(self, X, index=None, num_predictions=1):
         ValueError
             Raises if the `best_equation` cannot be evaluated.
         """
-        X = self._validate_data(X)
+        X = self._validate_data(X, ensure_2d=False)
+        if X.ndim == 1:
+            X = X.reshape(-1, 1)
+        assert X.ndim == 2
         _check_assertions(X, recursive_history_length=self.recursive_history_length)
         historical_X = self._sliding_window(X)[:: X.shape[1], :]
         if num_predictions < 1: