New updated checkpoint loading and tests

pysr/sr.py

@@ -1402,7 +1402,6 @@ class PySRRegressor(BaseEstimator, RegressorMixin, MultiOutputMixin):
         Xresampled=None,
         weights=None,
         variable_names=None,
-        from_equation_file=False,
     ):
         """
         Search for equations to fit the dataset and store them in `self.equations_`.
@@ -1500,18 +1499,21 @@ class PySRRegressor(BaseEstimator, RegressorMixin, MultiOutputMixin):
         )
 
         # Fitting procedure
-        if not from_equation_file:
-            self._run(X=X, y=y, weights=weights, seed=seed)
-        else:
-            self.equations_ = self.get_hof()
-
-        return self
+        return self._run(X=X, y=y, weights=weights, seed=seed)
 
-    def refresh(self):
+    def refresh(self, checkpoint_file=None):
         """
         Updates self.equations_ with any new options passed, such as
         :param`extra_sympy_mappings`.
+
+        Parameters
+        ----------
+        checkpoint_file : str, default=None
+            Path to checkpoint hall of fame file to be loaded.
         """
+        check_is_fitted(self, attributes=["equation_file_"])
+        if checkpoint_file:
+            self.equation_file_ = checkpoint_file
         self.equations_ = self.get_hof()
 
     def _decision_function(self, X, best_equation):
@@ -1709,6 +1711,15 @@ class PySRRegressor(BaseEstimator, RegressorMixin, MultiOutputMixin):
     def get_hof(self):
         """Get the equations from a hall of fame file. If no arguments
         entered, the ones used previously from a call to PySR will be used."""
+        check_is_fitted(
+            self,
+            attributes=[
+                "nout_",
+                "equation_file_",
+                "selection_mask_",
+                "feature_names_in_",
+            ],
+        )
         try:
             if self.nout_ > 1:
                 all_outputs = []

test/test.py

@@ -231,6 +231,17 @@ class TestPipeline(unittest.TestCase):
 
 class TestBest(unittest.TestCase):
     def setUp(self):
+        self.rstate = np.random.RandomState(0)
+        self.X = self.rstate.randn(10, 2)
+        self.y = np.cos(self.X[:, 0]) ** 2
+        self.model = PySRRegressor(
+            niterations=1,
+            extra_sympy_mappings={},
+            output_jax_format=False,
+            model_selection="accuracy",
+            equation_file="equation_file.csv",
+        )
+        self.model.fit(self.X, self.y)
         equations = pd.DataFrame(
             {
                 "equation": ["1.0", "cos(x0)", "square(cos(x0))"],
@@ -243,18 +254,6 @@ class TestBest(unittest.TestCase):
             "equation_file.csv.bkup", sep="|"
         )
 
-        self.model = PySRRegressor(
-            equation_file="equation_file.csv",
-            variable_names="x0 x1".split(" "),
-            extra_sympy_mappings={},
-            output_jax_format=False,
-            model_selection="accuracy",
-        )
-        self.rstate = np.random.RandomState(0)
-        # Placeholder values needed to fit the model from an equation file
-        self.X = self.rstate.randn(10, 2)
-        self.y = np.cos(self.X[:, 0]) ** 2
-        self.model.fit(self.X, self.y, from_equation_file=True)
         self.model.refresh()
         self.equations_ = self.model.equations_
 
@@ -308,6 +307,10 @@ class TestFeatureSelection(unittest.TestCase):
 class TestMiscellaneous(unittest.TestCase):
     """Test miscellaneous functions."""
 
+    def setUp(self):
+        # Allows all scikit-learn exception messages to be read.
+        self.maxDiff = None
+
     def test_deprecation(self):
         """Ensure that deprecation works as expected.
 

test/test_jax.py

@@ -23,6 +23,13 @@ class TestJAX(unittest.TestCase):
 
     def test_pipeline_pandas(self):
         X = pd.DataFrame(np.random.randn(100, 10))
+        y = np.ones(X.shape[0])
+        model = PySRRegressor(
+            max_evals=10000,
+            output_jax_format=True,
+        )
+        model.fit(X, y)
+
         equations = pd.DataFrame(
             {
                 "Equation": ["1.0", "cos(x1)", "square(cos(x1))"],
@@ -35,14 +42,7 @@ class TestJAX(unittest.TestCase):
             "equation_file.csv.bkup", sep="|"
         )
 
-        model = PySRRegressor(
-            equation_file="equation_file.csv",
-            output_jax_format=True,
-            variable_names="x1 x2 x3".split(" "),
-        )
-
-        model.fit(X, y=np.ones(X.shape[0]), from_equation_file=True)
-        model.refresh()
+        model.refresh(checkpoint_file="equation_file.csv")
         jformat = model.jax()
 
         np.testing.assert_almost_equal(
@@ -53,6 +53,10 @@ class TestJAX(unittest.TestCase):
 
     def test_pipeline(self):
         X = np.random.randn(100, 10)
+        y = np.ones(X.shape[0])
+        model = PySRRegressor(max_evals=10000, output_jax_format=True)
+        model.fit(X, y)
+
         equations = pd.DataFrame(
             {
                 "Equation": ["1.0", "cos(x1)", "square(cos(x1))"],
@@ -65,14 +69,7 @@ class TestJAX(unittest.TestCase):
             "equation_file.csv.bkup", sep="|"
         )
 
-        model = PySRRegressor(
-            equation_file="equation_file.csv",
-            output_jax_format=True,
-            variable_names="x1 x2 x3".split(" "),
-        )
-
-        model.fit(X, y=np.ones(X.shape[0]), from_equation_file=True)
-        model.refresh()
+        model.refresh(checkpoint_file="equation_file.csv")
         jformat = model.jax()
 
         np.testing.assert_almost_equal(

test/test_torch.py

@@ -3,7 +3,6 @@ import numpy as np
 import pandas as pd
 from pysr import sympy2torch, PySRRegressor
 import sympy
-from functools import partial
 
 
 class TestTorch(unittest.TestCase):
@@ -15,6 +14,7 @@ class TestTorch(unittest.TestCase):
         cosx = 1.0 * sympy.cos(x) + y
 
         import torch
+
         X = torch.tensor(np.random.randn(1000, 3))
         true = 1.0 * torch.cos(X[:, 0]) + X[:, 1]
         torch_module = sympy2torch(cosx, [x, y, z])
@@ -23,7 +23,6 @@ class TestTorch(unittest.TestCase):
         )
 
     def test_pipeline_pandas(self):
-        X = pd.DataFrame(np.random.randn(100, 10))
         equations = pd.DataFrame(
             {
                 "Equation": ["1.0", "cos(x1)", "square(cos(x1))"],
@@ -36,18 +35,16 @@ class TestTorch(unittest.TestCase):
             "equation_file.csv.bkup", sep="|"
         )
 
+        X = pd.DataFrame(np.random.randn(100, 10))
+        y = np.ones(X.shape[0])
         model = PySRRegressor(
+            max_evals=10000,
             model_selection="accuracy",
-            equation_file="equation_file.csv",
             extra_sympy_mappings={},
             output_torch_format=True,
         )
-        # Because a model hasn't been fit via the `fit` method, some
-        # attributes will not/cannot be set. For the purpose of
-        # testing, these attributes will be set manually here.
-        model.fit(X, y=np.ones(X.shape[0]), from_equation_file=True)
-        model.refresh()
-
+        model.fit(X, y)
+        model.refresh(checkpoint_file="equation_file.csv")
         tformat = model.pytorch()
         self.assertEqual(str(tformat), "_SingleSymPyModule(expression=cos(x1)**2)")
         import torch
@@ -60,6 +57,14 @@ class TestTorch(unittest.TestCase):
 
     def test_pipeline(self):
         X = np.random.randn(100, 10)
+        y = np.ones(X.shape[0])
+        model = PySRRegressor(
+            max_evals=10000,
+            model_selection="accuracy",
+            output_torch_format=True,
+        )
+        model.fit(X, y)
+
         equations = pd.DataFrame(
             {
                 "Equation": ["1.0", "cos(x1)", "square(cos(x1))"],
@@ -72,20 +77,13 @@ class TestTorch(unittest.TestCase):
             "equation_file.csv.bkup", sep="|"
         )
 
-        model = PySRRegressor(
-            model_selection="accuracy",
-            equation_file="equation_file.csv",
-            extra_sympy_mappings={},
-            output_torch_format=True,
-        )
-
-        model.fit(X, y=np.ones(X.shape[0]), from_equation_file=True)
-        model.refresh()
+        model.refresh(checkpoint_file="equation_file.csv")
 
         tformat = model.pytorch()
         self.assertEqual(str(tformat), "_SingleSymPyModule(expression=cos(x1)**2)")
 
         import torch
+
         np.testing.assert_almost_equal(
             tformat(torch.tensor(X)).detach().numpy(),
             np.square(np.cos(X[:, 1])),  # 2nd feature
@@ -99,6 +97,7 @@ class TestTorch(unittest.TestCase):
         module = sympy2torch(expression, [x, y, z])
 
         import torch
+
         X = torch.rand(100, 3).float() * 10
 
         true_out = (
@@ -112,6 +111,13 @@ class TestTorch(unittest.TestCase):
 
     def test_custom_operator(self):
         X = np.random.randn(100, 3)
+        y = np.ones(X.shape[0])
+        model = PySRRegressor(
+            max_evals=10000,
+            model_selection="accuracy",
+            output_torch_format=True,
+        )
+        model.fit(X, y)
 
         equations = pd.DataFrame(
             {
@@ -126,15 +132,13 @@ class TestTorch(unittest.TestCase):
         )
 
         import torch
-        model = PySRRegressor(
-            model_selection="accuracy",
+
+        model.set_params(
             equation_file="equation_file_custom_operator.csv",
             extra_sympy_mappings={"mycustomoperator": sympy.sin},
             extra_torch_mappings={"mycustomoperator": torch.sin},
-            output_torch_format=True,
         )
-        model.fit(X, y=np.ones(X.shape[0]), from_equation_file=True)
-        model.refresh()
+        model.refresh(checkpoint_file="equation_file_custom_operator.csv")
         self.assertEqual(str(model.sympy()), "sin(x1)")
         # Will automatically use the set global state from get_hof.
 
@@ -160,6 +164,7 @@ class TestTorch(unittest.TestCase):
 
         np_output = model.predict(X.values)
         import torch
+
         torch_output = torch_module(torch.tensor(X.values)).detach().numpy()
 
         np.testing.assert_almost_equal(np_output, torch_output, decimal=4)