Get PySRRegressor working with multi-output

pysr/sr.py

@@ -665,27 +665,46 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
         if self.equations is None:
             return "PySRRegressor.equations = None"
 
+        output = "PySRRegressor.equations = [\n"
+
         equations = self.equations
-        selected = ["" for _ in range(len(equations))]
-        if self.model_selection == "accuracy":
-            chosen_row = -1
-        elif self.model_selection == "best":
-            chosen_row = equations["score"].idxmax()
+        if not isinstance(equations, list):
+            all_equations = [equations]
         else:
-            raise NotImplementedError
-        selected[chosen_row] = ">>>>"
-        output = "PySRRegressor.equations = [\n"
-        repr_equations = pd.DataFrame(
-            dict(
-                pick=selected,
-                score=equations["score"],
-                equation=equations["equation"],
-                loss=equations["loss"],
-                complexity=equations["complexity"],
+            all_equations = equations
+
+        for i, equations in enumerate(all_equations):
+            selected = ["" for _ in range(len(equations))]
+            if self.model_selection == "accuracy":
+                chosen_row = -1
+            elif self.model_selection == "best":
+                chosen_row = equations["score"].idxmax()
+            else:
+                raise NotImplementedError
+            selected[chosen_row] = ">>>>"
+            repr_equations = pd.DataFrame(
+                dict(
+                    pick=selected,
+                    score=equations["score"],
+                    equation=equations["equation"],
+                    loss=equations["loss"],
+                    complexity=equations["complexity"],
+                )
             )
-        )
-        output += repr_equations.__repr__()
-        output += "\n]"
+
+            if len(all_equations) > 1:
+                output += "[\n"
+
+            for line in repr_equations.__repr__().split("\n"):
+                output += "\t" + line + "\n"
+
+            if len(all_equations) > 1:
+                output += "]"
+
+            if i < len(all_equations) - 1:
+                output += ", "
+
+        output += "]"
         return output
 
     def set_params(self, **params):
@@ -708,13 +727,19 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
 
     def get_best(self):
         if self.equations is None:
-            return 0.0
+            raise ValueError("No equations have been generated yet.")
         if self.model_selection == "accuracy":
+            if isinstance(self.equations, list):
+                return [eq.iloc[-1] for eq in self.equations]
             return self.equations.iloc[-1]
         elif self.model_selection == "best":
-            return best_row(self.equations)
+            if isinstance(self.equations, list):
+                return [eq.iloc[eq["score"].idxmax()] for eq in self.equations]
+            return self.equations.iloc[self.equations["score"].idxmax()]
         else:
-            raise NotImplementedError
+            raise NotImplementedError(
+                f"{self.model_selection} is not a valid model selection strategy."
+            )
 
     def fit(self, X, y, weights=None, variable_names=None):
         """Search for equations to fit the dataset.
@@ -747,26 +772,40 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
 
     def predict(self, X):
         self.refresh()
-        np_format = self.get_best()["lambda_format"]
-        return np_format(X)
+        best = self.get_best()
+        if self.multioutput:
+            return np.stack([eq["lambda_format"](X) for eq in best], axis=1)
+        return best["lambda_format"](X)
 
     def sympy(self):
         self.refresh()
-        return self.get_best()["sympy_format"]
+        best = self.get_best()
+        if self.multioutput:
+            return [eq["sympy_format"] for eq in best]
+        return best["sympy_format"]
 
     def latex(self):
         self.refresh()
-        return self.sympy().simplify()
+        sympy_representation = self.sympy()
+        if self.multioutput:
+            return [sympy.latex(s) for s in sympy_representation]
+        return sympy.latex(sympy_representation)
 
     def jax(self):
         self.set_params(output_jax_format=True)
         self.refresh()
-        return self.get_best()["jax_format"]
-
+        best = self.get_best()
+        if self.multioutput:
+            return [eq["jax_format"] for eq in best]
+        return best["jax_format"]
+        
     def pytorch(self):
         self.set_params(output_torch_format=True)
         self.refresh()
-        return self.get_best()["torch_format"]
+        best = self.get_best()
+        if self.multioutput:
+            return [eq["torch_format"] for eq in best]
+        return best["torch_format"]
 
     def _run(self, X, y, weights, variable_names):
         global already_ran
@@ -846,11 +885,11 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
 
         if len(y.shape) == 1 or (len(y.shape) == 2 and y.shape[1] == 1):
             self.multioutput = False
-            nout = 1
+            self.nout = 1
             y = y.reshape(-1)
         elif len(y.shape) == 2:
             self.multioutput = True
-            nout = y.shape[1]
+            self.nout = y.shape[1]
         else:
             raise NotImplementedError("y shape not supported!")
 
@@ -1182,3 +1221,8 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
         if self.multioutput:
             return ret_outputs
         return ret_outputs[0]
+
+    def score(self, X, y):
+        del X
+        del y
+        raise NotImplementedError

test/test.py

@@ -171,13 +171,13 @@ class TestBest(unittest.TestCase):
     def setUp(self):
         equations = pd.DataFrame(
             {
-                "Equation": ["1.0", "cos(x0)", "square(cos(x0))"],
-                "MSE": [1.0, 0.1, 1e-5],
-                "Complexity": [1, 2, 3],
+                "equation": ["1.0", "cos(x0)", "square(cos(x0))"],
+                "loss": [1.0, 0.1, 1e-5],
+                "complexity": [1, 2, 3],
             }
         )
 
-        equations["Complexity MSE Equation".split(" ")].to_csv(
+        equations["complexity loss equation".split(" ")].to_csv(
             "equation_file.csv.bkup", sep="|"
         )
 
@@ -195,19 +195,15 @@ class TestBest(unittest.TestCase):
         self.model.equations = self.equations
 
     def test_best(self):
-        self.assertEqual(best(self.equations), sympy.cos(sympy.Symbol("x0")) ** 2)
-        self.assertEqual(best(), sympy.cos(sympy.Symbol("x0")) ** 2)
         self.assertEqual(self.model.sympy(), sympy.cos(sympy.Symbol("x0")) ** 2)
 
     def test_best_tex(self):
-        self.assertEqual(best_tex(self.equations), "\\cos^{2}{\\left(x_{0} \\right)}")
-        self.assertEqual(best_tex(), "\\cos^{2}{\\left(x_{0} \\right)}")
         self.assertEqual(self.model.latex(), "\\cos^{2}{\\left(x_{0} \\right)}")
 
     def test_best_lambda(self):
         X = np.random.randn(10, 2)
         y = np.cos(X[:, 0]) ** 2
-        for f in [best_callable(), best_callable(self.equations)]:
+        for f in [self.model.predict, self.equations.iloc[-1]['lambda_format']]:
             np.testing.assert_almost_equal(f(X), y, decimal=4)