Add additional test for loading from pickle file

pysr/sr.py

@@ -926,7 +926,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
 
     def _checkpoint(self):
         """Saves the model's current state to a checkpoint file.
-        
+
         This should only be used internally by PySRRegressor."""
         # Save model state:
         self.show_pickle_warnings_ = False
@@ -2132,8 +2132,12 @@ def load(
         assert n_features_in is None
         with open(str(equation_file) + ".pkl", "rb") as f:
             model = pkl.load(f)
+        # Update any parameters if necessary, such as
+        # extra_sympy_mappings:
         model.set_params(**pysr_kwargs)
-        model.refresh()
+        if "equations_" not in model.__dict__ or model.equations_ is None:
+            model.refresh()
+
         return model
 
     # Else, we re-create it.

test/test.py

@@ -336,6 +336,16 @@ class TestPipeline(unittest.TestCase):
 
         np.testing.assert_allclose(model.predict(self.X), model2.predict(self.X))
 
+        # Try again, but using only the pickle file:
+        for file_to_delete in [str(equation_file), str(equation_file) + ".bkup"]:
+            if os.path.exists(file_to_delete):
+                os.remove(file_to_delete)
+
+        model3 = load(
+            model.equation_file_, extra_sympy_mappings={"sq": lambda x: x**2}
+        )
+        np.testing.assert_allclose(model.predict(self.X), model3.predict(self.X))
+
 
 class TestBest(unittest.TestCase):
     def setUp(self):