Add __repr__ method that lists selected equation

TODO.md

@@ -65,6 +65,9 @@
 
 - [ ] Automatically convert log, log10, log2, pow to the correct operators.
 - [ ] I think the simplification isn't working correctly (post-merging SymbolicUtils.)
+- [ ] Show demo of PySRRegressor. Fit equations, then show how to view equations.
+- [ ] Add "selected" column string to regular equations dict.
+- [ ] List "Loss" instead of "MSE"
 
 ## Feature ideas
 

example.py

@@ -1,25 +1,23 @@
 import numpy as np
-from pysr import pysr, best
+from pysr import PySRRegressor
 
 # Dataset
-X = 2 * np.random.randn(100, 5)
-y = 2 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2
+X = 3 * np.random.randn(100, 5)
+y = 3 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2
 
 # Learn equations
-equations = pysr(
-    X,
-    y,
-    niterations=5,
+model = PySRRegressor(
+    niterations=6,
     binary_operators=["plus", "mult"],
     unary_operators=[
         "cos",
         "exp",
         "sin",  # Pre-defined library of operators (see https://pysr.readthedocs.io/en/latest/docs/operators/)
-        "inv(x) = 1/x",
+        "inv(x) = 2/x",
     ],
     loss="loss(x, y) = abs(x - y)",  # Custom loss function
 )  # Define your own operator! (Julia syntax)
 
-...  # (you can use ctl-c to exit early)
+model.fit(X, y)
 
-print(best(equations))
+print(model)

pysr/sklearn.py

@@ -1,9 +1,10 @@
 from pysr import pysr, best_row
-from sklearn.base import BaseEstimator
+from sklearn.base import BaseEstimator, RegressorMixin
 import inspect
+import pandas as pd
 
 
-class PySRRegressor(BaseEstimator):
+class PySRRegressor(BaseEstimator, RegressorMixin):
     def __init__(self, model_selection="accuracy", **params):
         """Initialize settings for pysr.pysr call.
 
@@ -18,7 +19,30 @@ def __init__(self, model_selection="accuracy", **params):
         self.equations = None
 
     def __repr__(self):
-        return f"PySRRegressor(equations={self.get_best()['sympy_format']})"
+        if self.equations is None:
+            return "PySRRegressor.equations=None"
+
+        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()
+        else:
+            raise NotImplementedError
+        selected[chosen_row] = ">"
+        output = "PySRRegressor.equations=[\n"
+        repr_equations = pd.DataFrame(
+            dict(
+                selected=selected,
+                score=equations["score"],
+                MSE=equations["MSE"],
+                Complexity=equations["Complexity"],
+            )
+        )
+        output += repr_equations.__repr__()
+        output += "\n]"
+        return output
 
     def set_params(self, **params):
         """Set parameters for pysr.pysr call or model_selection strategy."""