From 10bac3908f15817d55cbd7725ef0eb33829a3973 Mon Sep 17 00:00:00 2001 From: MilesCranmer Date: Tue, 26 Apr 2022 16:28:32 -0400 Subject: [PATCH] Fix edgecase when expression is a constant --- pysr/sr.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pysr/sr.py b/pysr/sr.py index acd47253..aff4d411 100644 --- a/pysr/sr.py +++ b/pysr/sr.py @@ -271,12 +271,13 @@ def __repr__(self): return f"PySRFunction(X=>{self._sympy})" def __call__(self, X): + expected_shape = (X.shape[0],) if isinstance(X, pd.DataFrame): # Lambda function takes as argument: - return self._lambda(**{k: X[k].values for k in X.columns}) + return self._lambda(**{k: X[k].values for k in X.columns}) * np.ones(expected_shape) elif self._selection is not None: - return self._lambda(*X[:, self._selection].T) - return self._lambda(*X.T) + return self._lambda(*X[:, self._selection].T) * np.ones(expected_shape) + return self._lambda(*X.T) * np.ones(expected_shape) def _get_julia_project(julia_project):