From ccf71e9b9eb54f0bdb5cca88414cfb34639ec7f8 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Wed, 20 Jul 2022 14:28:04 -0400
Subject: [PATCH 01/21] `load` function to init model from saved equations

---
 pysr/__init__.py |  1 +
 pysr/sr.py       | 74 ++++++++++++++++++++++++++++++++++++++++++++++++
 test/test.py     | 28 +++++++++++++++++-
 3 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/pysr/__init__.py b/pysr/__init__.py
index e303becb2..210e85cb7 100644
--- a/pysr/__init__.py
+++ b/pysr/__init__.py
@@ -6,6 +6,7 @@
     best_tex,
     best_callable,
     best_row,
+    load,
 )
 from .julia_helpers import install
 from .feynman_problems import Problem, FeynmanProblem
diff --git a/pysr/sr.py b/pysr/sr.py
index 3e2112975..ec12fe877 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -2034,3 +2034,77 @@ def run_feature_selection(X, y, select_k_features, random_state=None):
         clf, threshold=-np.inf, max_features=select_k_features, prefit=True
     )
     return selector.get_support(indices=True)
+
+
+def load(
+    equation_file,
+    *,
+    binary_operators,
+    unary_operators,
+    n_features_in,
+    feature_names_in=None,
+    selection_mask=None,
+    nout=1,
+    **pysr_kwargs,
+):
+    """
+    Create a model from equations stored as a csv file
+
+    Parameters
+    ----------
+    equation_file : str
+        Path to a csv file containing equations.
+
+    binary_operators : list[str], default=["+", "-", "*", "/"]
+        The same binary operators used when creating the model.
+
+    unary_operators : list[str], default=None
+        The same unary operators used when creating the model.
+
+    n_features_in : int
+        Number of features passed to the model.
+
+    feature_names_in : list[str], default=None
+        Names of the features passed to the model.
+
+    selection_mask : list[bool], default=None
+        If using select_k_features, you must pass `model.selection_mask_` here.
+
+    nout : int, default=1
+        Number of outputs of the model.
+
+    pysr_kwargs : dict
+        Any other keyword arguments to initialize the PySRRegressor object.
+
+    Returns
+    -------
+    model : PySRRegressor
+        The model with fitted equations.
+    """
+
+    # TODO: copy .bkup file if exists.
+    model = PySRRegressor(
+        equation_file=equation_file,
+        binary_operators=binary_operators,
+        unary_operators=unary_operators,
+        **pysr_kwargs,
+    )
+
+    model.equation_file_ = equation_file
+    model.nout_ = nout
+    model.n_features_in_ = n_features_in
+
+    if feature_names_in is None:
+        model.feature_names_in_ = [f"x{i}" for i in range(n_features_in)]
+    else:
+        assert len(feature_names_in) == n_features_in
+        model.feature_names_in_ = feature_names_in
+
+    if selection_mask is None:
+        model.selection_mask_ = np.ones(n_features_in, dtype=bool)
+    else:
+        model.selection_mask_ = selection_mask
+
+    model.refresh()
+
+    return model
diff --git a/test/test.py b/test/test.py
index 4c82a17e1..1c581b7fe 100644
--- a/test/test.py
+++ b/test/test.py
@@ -4,7 +4,7 @@
 import unittest
 import numpy as np
 from sklearn import model_selection
-from pysr import PySRRegressor
+from pysr import PySRRegressor, load
 from pysr.sr import run_feature_selection, _handle_feature_selection
 from sklearn.utils.estimator_checks import check_estimator
 import sympy
@@ -280,6 +280,32 @@ def test_high_dim_selection_early_stop(self):
         model.fit(X.values, y.values, Xresampled=Xresampled.values)
         self.assertLess(np.average((model.predict(X.values) - y.values) ** 2), 1e-4)
 
+    def test_load_model(self):
+        """See if we can load a ran model from the equation file."""
+        csv_file_data = """
+        Complexity|MSE|Equation
+        1|0.19951081|1.9762075
+        3|0.12717344|(f0 + 1.4724599)
+        4|0.104823045|pow_abs(2.2683423, cos(f3))"""
+        # Strip the indents:
+        csv_file_data = "\n".join([l.strip() for l in csv_file_data.split("\n")])
+        with open("equation_file.csv", "w") as f:
+            f.write(csv_file_data)
+        with open("equation_file.csv.bkup", "w") as f:
+            f.write(csv_file_data)
+        model = load(
+            "equation_file.csv",
+            n_features_in=5,
+            feature_names_in=["f0", "f1", "f2", "f3", "f4"],
+            binary_operators=["+", "*", "/", "-", "^"],
+            unary_operators=["cos"],
+        )
+        X = self.rstate.rand(100, 5)
+        y_truth = 2.2683423 ** np.cos(X[:, 3])
+        y_test = model.predict(X, 2)
+
+        np.testing.assert_allclose(y_truth, y_test)
+
 
 class TestBest(unittest.TestCase):
     def setUp(self):

From e5b4869851bd5a5f7b6fa783cf6d0f8ff10ca8a5 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Wed, 20 Jul 2022 14:32:27 -0400
Subject: [PATCH 02/21] Call `refresh` in load function

---
 pysr/sr.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index ec12fe877..07b34cd31 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -2090,7 +2090,6 @@ def load(
         **pysr_kwargs,
     )
 
-    model.equation_file_ = equation_file
     model.nout_ = nout
     model.n_features_in_ = n_features_in
 
@@ -2105,6 +2104,6 @@ def load(
     else:
         model.selection_mask_ = selection_mask
 
-    model.refresh()
+    model.refresh(checkpoint_file=equation_file)
 
     return model

From 179fef6351ee7d5356ec2a1ce9efcda8241dd935 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:42:59 -0400
Subject: [PATCH 03/21] Correctly set path names

---
 test/test.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/test/test.py b/test/test.py
index 1c581b7fe..59c7d76bc 100644
--- a/test/test.py
+++ b/test/test.py
@@ -12,6 +12,7 @@
 import warnings
 import pickle as pkl
 import tempfile
+from pathlib import Path
 
 DEFAULT_PARAMS = inspect.signature(PySRRegressor.__init__).parameters
 DEFAULT_NITERATIONS = DEFAULT_PARAMS["niterations"].default
@@ -289,12 +290,14 @@ def test_load_model(self):
         4|0.104823045|pow_abs(2.2683423, cos(f3))"""
         # Strip the indents:
         csv_file_data = "\n".join([l.strip() for l in csv_file_data.split("\n")])
-        with open("equation_file.csv", "w") as f:
+        rand_dir = Path(tempfile.mkdtemp())
+        equation_filename = rand_dir / "equation.csv"
+        with open(equation_filename, "w") as f:
             f.write(csv_file_data)
-        with open("equation_file.csv.bkup", "w") as f:
+        with open(equation_filename + ".bkup", "w") as f:
             f.write(csv_file_data)
         model = load(
-            "equation_file.csv",
+            equation_filename,
             n_features_in=5,
             feature_names_in=["f0", "f1", "f2", "f3", "f4"],
             binary_operators=["+", "*", "/", "-", "^"],

From 85371bb899ddc546f448adec34e5b2977f080f9d Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:44:41 -0400
Subject: [PATCH 04/21] Allow pickling without equations_ stored

---
 pysr/sr.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index 07b34cd31..67c11f58c 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -883,7 +883,9 @@ def __getstate__(self):
             key: None if key == "raw_julia_state_" else value
             for key, value in state.items()
         }
-        if "equations_" in pickled_state:
+        if ("equations_" in pickled_state) and (
+            pickled_state["equations_"] is not None
+        ):
             pickled_state["output_torch_format"] = False
             pickled_state["output_jax_format"] = False
             if self.nout_ == 1:

From dde0ef7e3c2606415b9d7c03c56370402c398e3e Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:45:13 -0400
Subject: [PATCH 05/21] Remove extra_sympy_mappings from pickle file

---
 pysr/sr.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index 67c11f58c..e147d08a7 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -562,6 +562,9 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
     equation_file_contents_ : list[pandas.DataFrame]
         Contents of the equation file output by the Julia backend.
 
+    show_pickle_warnings_ : bool
+        Whether to show warnings about what attributes can be pickled.
+
     Notes
     -----
     Most default parameters have been tuned over several example equations,
@@ -873,14 +876,26 @@ def __getstate__(self):
         from the pickled instance.
         """
         state = self.__dict__
-        if "raw_julia_state_" in state:
+        show_pickle_warning = not (
+            "show_pickle_warnings_" in state and not state["show_pickle_warnings_"]
+        )
+        if "raw_julia_state_" in state and show_pickle_warning:
             warnings.warn(
                 "raw_julia_state_ cannot be pickled and will be removed from the "
                 "serialized instance. This will prevent a `warm_start` fit of any "
                 "model that is deserialized via `pickle.load()`."
             )
+        state_keys_containing_lambdas = ["extra_sympy_mappings", "extra_torch_mappings"]
+        for state_key in state_keys_containing_lambdas:
+            if state[state_key] is not None and show_pickle_warning:
+                warnings.warn(
+                    f"`{state_key}` cannot be pickled and will be removed from the "
+                    "serialized instance. When loading the model, please redefine "
+                    f"`{state_key}` at runtime."
+                )
+        state_keys_to_clear = ["raw_julia_state_"] + state_keys_containing_lambdas
         pickled_state = {
-            key: None if key == "raw_julia_state_" else value
+            key: (None if key in state_keys_to_clear else value)
             for key, value in state.items()
         }
         if ("equations_" in pickled_state) and (

From b16d9efb3c83ced7870fbb7641fa97cfd9452a2d Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:45:42 -0400
Subject: [PATCH 06/21] Automatically pickle file at initialization

---
 pysr/sr.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pysr/sr.py b/pysr/sr.py
index e147d08a7..63d74fe61 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -1623,6 +1623,11 @@ def fit(
             y,
         )
 
+        # Save model state:
+        self.show_pickle_warnings_ = False
+        with open(str(self.equation_file_) + ".pkl", "wb") as f:
+            pkl.dump(self, f)
+        self.show_pickle_warnings_ = True
         # Fitting procedure
         return self._run(X, y, mutated_params, weights=weights, seed=seed)
 

From 5c0ad5569248da926a646aeb6194ad9d03fc9844 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:46:11 -0400
Subject: [PATCH 07/21] Allow loading from pickle file

---
 pysr/sr.py | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index 63d74fe61..cd851c16b 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -2061,9 +2061,9 @@ def run_feature_selection(X, y, select_k_features, random_state=None):
 def load(
     equation_file,
     *,
-    binary_operators,
-    unary_operators,
-    n_features_in,
+    binary_operators=None,
+    unary_operators=None,
+    n_features_in=None,
     feature_names_in=None,
     selection_mask=None,
     nout=1,
@@ -2097,12 +2097,33 @@ def load(
 
     pysr_kwargs : dict
         Any other keyword arguments to initialize the PySRRegressor object.
+        These will overwrite those stored in the pickle file.
 
     Returns
     -------
     model : PySRRegressor
         The model with fitted equations.
     """
+    # Try to load model from <equation_file>.pkl
+    print(f"Checking if {equation_file}.pkl exists...")
+    if os.path.exists(str(equation_file) + ".pkl"):
+        assert binary_operators is None
+        assert unary_operators is None
+        assert n_features_in is None
+        with open(str(equation_file) + ".pkl", "rb") as f:
+            model = pkl.load(f)
+        model.set_params(**pysr_kwargs)
+        model.refresh()
+        return model
+
+    # Else, we re-create it.
+    print(
+        f"{equation_file}.pkl does not exist, "
+        "so we must create the model from scratch."
+    )
+    assert binary_operators is not None
+    assert unary_operators is not None
+    assert n_features_in is not None
 
     # TODO: copy .bkup file if exists.
     model = PySRRegressor(

From dc1d66378e25d7a8ef5d45811e0a67d0b00c449e Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:46:43 -0400
Subject: [PATCH 08/21] Add pickle files to gitignore

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index 23004701d..f0daf5e88 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
 *.csv
 *.csv.out*
 *.bkup
+*.pkl
 performance*txt
 *.out
 trials*

From 4ae8a5c2380b0fa6dd34f2f56207d2dc3970a362 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:47:00 -0400
Subject: [PATCH 09/21] Add missing pickle import

---
 pysr/sr.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pysr/sr.py b/pysr/sr.py
index cd851c16b..c093d3619 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -8,6 +8,7 @@
 import tempfile
 import shutil
 from pathlib import Path
+import pickle as pkl
 from datetime import datetime
 import warnings
 from multiprocessing import cpu_count

From 78cdb0e736efe0575ecd7797e7ad7e07b6ecd447 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:47:15 -0400
Subject: [PATCH 10/21] Add test for loading from pickle file

---
 test/test.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/test/test.py b/test/test.py
index 59c7d76bc..33999e087 100644
--- a/test/test.py
+++ b/test/test.py
@@ -309,6 +309,33 @@ def test_load_model(self):
 
         np.testing.assert_allclose(y_truth, y_test)
 
+    def test_load_model_simple(self):
+        # Test that we can simply load a model from its equation file.
+        y = self.X[:, [0, 1]] ** 2
+        model = PySRRegressor(
+            # Test that passing a single operator works:
+            unary_operators="sq(x) = x^2",
+            binary_operators="plus",
+            extra_sympy_mappings={"sq": lambda x: x**2},
+            **self.default_test_kwargs,
+            procs=0,
+            denoise=True,
+            early_stop_condition="stop_if(loss, complexity) = loss < 0.05 && complexity == 2",
+        )
+        rand_dir = Path(tempfile.mkdtemp())
+        equation_file = rand_dir / "equations.csv"
+        model.set_params(temp_equation_file=False)
+        model.set_params(equation_file=equation_file)
+        model.fit(self.X, y)
+
+        # lambda functions are removed from the pickling, so we need
+        # to pass it during the loading:
+        model2 = load(
+            model.equation_file_, extra_sympy_mappings={"sq": lambda x: x**2}
+        )
+
+        np.testing.assert_allclose(model.predict(self.X), model2.predict(self.X))
+
 
 class TestBest(unittest.TestCase):
     def setUp(self):

From 214744b5ce5f1a375f5af936b774ca3a3b26bdd4 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 22:59:41 -0400
Subject: [PATCH 11/21] Fix filename concat in test

---
 test/test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test.py b/test/test.py
index 33999e087..0e1108400 100644
--- a/test/test.py
+++ b/test/test.py
@@ -294,7 +294,7 @@ def test_load_model(self):
         equation_filename = rand_dir / "equation.csv"
         with open(equation_filename, "w") as f:
             f.write(csv_file_data)
-        with open(equation_filename + ".bkup", "w") as f:
+        with open(str(equation_filename) + ".bkup", "w") as f:
             f.write(csv_file_data)
         model = load(
             equation_filename,

From 58e25a9d7261ceb4509125d6f8d102b68bfe5fd4 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 23:30:52 -0400
Subject: [PATCH 12/21] Test both with and without `bkup` file

---
 test/test.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/test/test.py b/test/test.py
index 0e1108400..30a50d1cc 100644
--- a/test/test.py
+++ b/test/test.py
@@ -290,24 +290,24 @@ def test_load_model(self):
         4|0.104823045|pow_abs(2.2683423, cos(f3))"""
         # Strip the indents:
         csv_file_data = "\n".join([l.strip() for l in csv_file_data.split("\n")])
-        rand_dir = Path(tempfile.mkdtemp())
-        equation_filename = rand_dir / "equation.csv"
-        with open(equation_filename, "w") as f:
-            f.write(csv_file_data)
-        with open(str(equation_filename) + ".bkup", "w") as f:
-            f.write(csv_file_data)
-        model = load(
-            equation_filename,
-            n_features_in=5,
-            feature_names_in=["f0", "f1", "f2", "f3", "f4"],
-            binary_operators=["+", "*", "/", "-", "^"],
-            unary_operators=["cos"],
-        )
-        X = self.rstate.rand(100, 5)
-        y_truth = 2.2683423 ** np.cos(X[:, 3])
-        y_test = model.predict(X, 2)
 
-        np.testing.assert_allclose(y_truth, y_test)
+        for from_backup in [False, True]:
+            rand_dir = Path(tempfile.mkdtemp())
+            equation_filename = str(rand_dir / "equation.csv")
+            with open(equation_filename + (".bkup" if from_backup else ""), "w") as f:
+                f.write(csv_file_data)
+            model = load(
+                equation_filename,
+                n_features_in=5,
+                feature_names_in=["f0", "f1", "f2", "f3", "f4"],
+                binary_operators=["+", "*", "/", "-", "^"],
+                unary_operators=["cos"],
+            )
+            X = self.rstate.rand(100, 5)
+            y_truth = 2.2683423 ** np.cos(X[:, 3])
+            y_test = model.predict(X, 2)
+
+            np.testing.assert_allclose(y_truth, y_test)
 
     def test_load_model_simple(self):
         # Test that we can simply load a model from its equation file.

From 1f019764c8af52477ed2f066b50e17e7474ec26e Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 23:31:28 -0400
Subject: [PATCH 13/21] Don't check for `equation_file_` until after
 checkpoint_file set

---
 pysr/sr.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index c093d3619..be9224893 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -1642,10 +1642,10 @@ def refresh(self, checkpoint_file=None):
         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.equation_file_contents_ = None
+        check_is_fitted(self, attributes=["equation_file_"])
         self.equations_ = self.get_hof()
 
     def predict(self, X, index=None):

From f1ac7043f981a5d2fa1202d172986ab9ee261f99 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 23:32:00 -0400
Subject: [PATCH 14/21] Allow both `bkup` and `csv` file

---
 pysr/sr.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index be9224893..b53d222a0 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -1834,10 +1834,10 @@ def _read_equation_file(self):
             if self.nout_ > 1:
                 all_outputs = []
                 for i in range(1, self.nout_ + 1):
-                    df = pd.read_csv(
-                        str(self.equation_file_) + f".out{i}" + ".bkup",
-                        sep="|",
-                    )
+                    cur_filename = str(self.equation_file_) + f".out{i}" + ".bkup"
+                    if not os.path.exists(cur_filename):
+                        cur_filename = str(self.equation_file_) + f".out{i}"
+                    df = pd.read_csv(cur_filename, sep="|")
                     # Rename Complexity column to complexity:
                     df.rename(
                         columns={
@@ -1850,7 +1850,10 @@ def _read_equation_file(self):
 
                     all_outputs.append(df)
             else:
-                all_outputs = [pd.read_csv(str(self.equation_file_) + ".bkup", sep="|")]
+                filename = str(self.equation_file_) + ".bkup"
+                if not os.path.exists(filename):
+                    filename = str(self.equation_file_)
+                all_outputs = [pd.read_csv(filename, sep="|")]
                 all_outputs[-1].rename(
                     columns={
                         "Complexity": "complexity",

From c6902b714c3a993d00d41a90100cc6de79c5f50f Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Sun, 24 Jul 2022 23:32:41 -0400
Subject: [PATCH 15/21] Additional logging messages during load

---
 pysr/sr.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pysr/sr.py b/pysr/sr.py
index b53d222a0..d669147e2 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -2111,6 +2111,7 @@ def load(
     # Try to load model from <equation_file>.pkl
     print(f"Checking if {equation_file}.pkl exists...")
     if os.path.exists(str(equation_file) + ".pkl"):
+        print(f"Loading model from {equation_file}.pkl.")
         assert binary_operators is None
         assert unary_operators is None
         assert n_features_in is None

From 6501ca074793bde8af2ef943394c4e03ef43a775 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Mon, 1 Aug 2022 14:24:05 -0400
Subject: [PATCH 16/21] Checkpoint model before and after fit

---
 pysr/sr.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index d669147e2..2d8e5463b 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -924,6 +924,16 @@ def __getstate__(self):
                 ]
         return pickled_state
 
+    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
+        with open(str(self.equation_file_) + ".pkl", "wb") as f:
+            pkl.dump(self, f)
+        self.show_pickle_warnings_ = True
+
     @property
     def equations(self):  # pragma: no cover
         warnings.warn(
@@ -1624,13 +1634,18 @@ def fit(
             y,
         )
 
-        # Save model state:
-        self.show_pickle_warnings_ = False
-        with open(str(self.equation_file_) + ".pkl", "wb") as f:
-            pkl.dump(self, f)
-        self.show_pickle_warnings_ = True
-        # Fitting procedure
-        return self._run(X, y, mutated_params, weights=weights, seed=seed)
+        # Initially, just save model parameters, so that
+        # it can be loaded from an early exit:
+        self._checkpoint()
+
+        # Perform the search:
+        self._run(X, y, mutated_params, weights=weights, seed=seed)
+
+        # Then, after fit, we save again, so the pickle file contains
+        # the equations:
+        self._checkpoint()
+
+        return self
 
     def refresh(self, checkpoint_file=None):
         """

From b53e7fafda3ee0a06d9e8ee56f98cc46bd7ddd57 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Mon, 1 Aug 2022 14:40:12 -0400
Subject: [PATCH 17/21] Add additional test for loading from pickle file

---
 pysr/sr.py   |  8 ++++++--
 test/test.py | 10 ++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index 2d8e5463b..f99180dd2 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -926,7 +926,7 @@ def __getstate__(self):
 
     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.
diff --git a/test/test.py b/test/test.py
index 30a50d1cc..f5e84570e 100644
--- a/test/test.py
+++ b/test/test.py
@@ -336,6 +336,16 @@ def test_load_model_simple(self):
 
         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):

From b8a97f177e29858c39aaeabd1d998b5207be2c95 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Thu, 4 Aug 2022 15:23:41 -0400
Subject: [PATCH 18/21] Use .pkl instead of .csv.pkl

---
 pysr/sr.py   | 38 ++++++++++++++++++++++++++++----------
 test/test.py | 21 ++++++++++++++++++++-
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index f99180dd2..d1a209cd6 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -930,7 +930,7 @@ def _checkpoint(self):
         This should only be used internally by PySRRegressor."""
         # Save model state:
         self.show_pickle_warnings_ = False
-        with open(str(self.equation_file_) + ".pkl", "wb") as f:
+        with open(_csv_filename_to_pkl_filename(self.equation_file_), "wb") as f:
             pkl.dump(self, f)
         self.show_pickle_warnings_ = True
 
@@ -1636,14 +1636,16 @@ def fit(
 
         # Initially, just save model parameters, so that
         # it can be loaded from an early exit:
-        self._checkpoint()
+        if not self.temp_equation_file:
+            self._checkpoint()
 
         # Perform the search:
         self._run(X, y, mutated_params, weights=weights, seed=seed)
 
         # Then, after fit, we save again, so the pickle file contains
         # the equations:
-        self._checkpoint()
+        if not self.temp_equation_file:
+            self._checkpoint()
 
         return self
 
@@ -2077,6 +2079,17 @@ def run_feature_selection(X, y, select_k_features, random_state=None):
     return selector.get_support(indices=True)
 
 
+def _csv_filename_to_pkl_filename(csv_filename) -> str:
+    # Assume that the csv filename is of the form "foo.csv"
+    dirname = str(os.path.dirname(csv_filename))
+    basename = str(os.path.basename(csv_filename))
+    base = str(os.path.splitext(basename)[0])
+
+    pkl_basename = base + ".pkl"
+
+    return os.path.join(dirname, pkl_basename)
+
+
 def load(
     equation_file,
     *,
@@ -2094,7 +2107,8 @@ def load(
     Parameters
     ----------
     equation_file : str
-        Path to a csv file containing equations.
+        Path to a csv file containing equations, or a pickle file
+        containing the model.
 
     binary_operators : list[str], default=["+", "-", "*", "/"]
         The same binary operators used when creating the model.
@@ -2123,14 +2137,19 @@ def load(
     model : PySRRegressor
         The model with fitted equations.
     """
+    if os.path.splitext(equation_file)[1] != ".pkl":
+        pkl_filename = _csv_filename_to_pkl_filename(equation_file)
+    else:
+        pkl_filename = equation_file
+
     # Try to load model from <equation_file>.pkl
-    print(f"Checking if {equation_file}.pkl exists...")
-    if os.path.exists(str(equation_file) + ".pkl"):
-        print(f"Loading model from {equation_file}.pkl.")
+    print(f"Checking if {pkl_filename} exists...")
+    if os.path.exists(pkl_filename):
+        print(f"Loading model from {pkl_filename}")
         assert binary_operators is None
         assert unary_operators is None
         assert n_features_in is None
-        with open(str(equation_file) + ".pkl", "rb") as f:
+        with open(pkl_filename, "rb") as f:
             model = pkl.load(f)
         # Update any parameters if necessary, such as
         # extra_sympy_mappings:
@@ -2142,8 +2161,7 @@ def load(
 
     # Else, we re-create it.
     print(
-        f"{equation_file}.pkl does not exist, "
-        "so we must create the model from scratch."
+        f"{equation_file} does not exist, " "so we must create the model from scratch."
     )
     assert binary_operators is not None
     assert unary_operators is not None
diff --git a/test/test.py b/test/test.py
index f5e84570e..dd07c601f 100644
--- a/test/test.py
+++ b/test/test.py
@@ -5,7 +5,11 @@
 import numpy as np
 from sklearn import model_selection
 from pysr import PySRRegressor, load
-from pysr.sr import run_feature_selection, _handle_feature_selection
+from pysr.sr import (
+    run_feature_selection,
+    _handle_feature_selection,
+    _csv_filename_to_pkl_filename,
+)
 from sklearn.utils.estimator_checks import check_estimator
 import sympy
 import pandas as pd
@@ -341,6 +345,7 @@ def test_load_model_simple(self):
             if os.path.exists(file_to_delete):
                 os.remove(file_to_delete)
 
+        pickle_file = rand_dir / "equations.pkl"
         model3 = load(
             model.equation_file_, extra_sympy_mappings={"sq": lambda x: x**2}
         )
@@ -430,6 +435,20 @@ def test_feature_selection_handler(self):
 class TestMiscellaneous(unittest.TestCase):
     """Test miscellaneous functions."""
 
+    def test_csv_to_pkl_conversion(self):
+        """Test that csv filename to pkl filename works as expected."""
+        tmpdir = Path(tempfile.mkdtemp())
+        equation_file = tmpdir / "equations.389479384.28378374.csv"
+        expected_pkl_file = tmpdir / "equations.389479384.28378374.pkl"
+
+        # First, test inputting the paths:
+        test_pkl_file = _csv_filename_to_pkl_filename(equation_file)
+        self.assertEqual(test_pkl_file, str(expected_pkl_file))
+
+        # Next, test inputting the strings.
+        test_pkl_file = _csv_filename_to_pkl_filename(str(equation_file))
+        self.assertEqual(test_pkl_file, str(expected_pkl_file))
+
     def test_deprecation(self):
         """Ensure that deprecation works as expected.
 

From a6bed2c01177ba435c141e1cd540409c3d3e34ec Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Thu, 4 Aug 2022 18:25:42 -0400
Subject: [PATCH 19/21] Fix bug with inplace editing of equation_file_contents_

---
 pysr/sr.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/pysr/sr.py b/pysr/sr.py
index d1a209cd6..3a5415554 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -1,3 +1,4 @@
+import copy
 import os
 import sys
 import numpy as np
@@ -1928,7 +1929,9 @@ def get_hof(self):
 
         ret_outputs = []
 
-        for output in self.equation_file_contents_:
+        equation_file_contents = copy.deepcopy(self.equation_file_contents_)
+
+        for output in equation_file_contents:
 
             scores = []
             lastMSE = None

From f5577eac29e49fe913ce12a214b88cba787f2e6d Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Thu, 4 Aug 2022 18:57:30 -0400
Subject: [PATCH 20/21] Reduce precision of tests

---
 test/test.py     | 6 +++---
 test/test_jax.py | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/test.py b/test/test.py
index dd07c601f..dd1ece29f 100644
--- a/test/test.py
+++ b/test/test.py
@@ -140,7 +140,7 @@ def test_multioutput_weighted_with_callable_temp_equation(self):
         # These tests are flaky, so don't fail test:
         try:
             np.testing.assert_almost_equal(
-                model.predict(X.copy())[:, 0], X[:, 0] ** 2, decimal=4
+                model.predict(X.copy())[:, 0], X[:, 0] ** 2, decimal=3
             )
         except AssertionError:
             print("Error in test_multioutput_weighted_with_callable_temp_equation")
@@ -149,7 +149,7 @@ def test_multioutput_weighted_with_callable_temp_equation(self):
 
         try:
             np.testing.assert_almost_equal(
-                model.predict(X.copy())[:, 1], X[:, 1] ** 2, decimal=4
+                model.predict(X.copy())[:, 1], X[:, 1] ** 2, decimal=3
             )
         except AssertionError:
             print("Error in test_multioutput_weighted_with_callable_temp_equation")
@@ -401,7 +401,7 @@ def test_best_lambda(self):
         X = self.X
         y = self.y
         for f in [self.model.predict, self.equations_.iloc[-1]["lambda_format"]]:
-            np.testing.assert_almost_equal(f(X), y, decimal=4)
+            np.testing.assert_almost_equal(f(X), y, decimal=3)
 
 
 class TestFeatureSelection(unittest.TestCase):
diff --git a/test/test_jax.py b/test/test_jax.py
index 58d1a6067..e885a8d3b 100644
--- a/test/test_jax.py
+++ b/test/test_jax.py
@@ -76,7 +76,7 @@ def test_pipeline(self):
         np.testing.assert_almost_equal(
             np.array(jformat["callable"](jnp.array(X), jformat["parameters"])),
             np.square(np.cos(X[:, 1])),  # Select feature 1
-            decimal=4,
+            decimal=3,
         )
 
     def test_feature_selection_custom_operators(self):

From 34f4e3f83fb1f1dd691ad5b57572bf1e7673125e Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 5 Aug 2022 00:22:07 -0400
Subject: [PATCH 21/21] Change model load to classmethod

---
 pysr/__init__.py |   1 -
 pysr/sr.py       | 217 ++++++++++++++++++++++++-----------------------
 test/test.py     |   8 +-
 3 files changed, 117 insertions(+), 109 deletions(-)

diff --git a/pysr/__init__.py b/pysr/__init__.py
index 210e85cb7..e303becb2 100644
--- a/pysr/__init__.py
+++ b/pysr/__init__.py
@@ -6,7 +6,6 @@
     best_tex,
     best_callable,
     best_row,
-    load,
 )
 from .julia_helpers import install
 from .feynman_problems import Problem, FeynmanProblem
diff --git a/pysr/sr.py b/pysr/sr.py
index 3a5415554..e98d36b67 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -810,6 +810,119 @@ def __init__(
                         f"{k} is not a valid keyword argument for PySRRegressor."
                     )
 
+    @classmethod
+    def from_file(
+        cls,
+        equation_file,
+        *,
+        binary_operators=None,
+        unary_operators=None,
+        n_features_in=None,
+        feature_names_in=None,
+        selection_mask=None,
+        nout=1,
+        **pysr_kwargs,
+    ):
+        """
+        Create a model from a saved model checkpoint or equation file.
+
+        Parameters
+        ----------
+        equation_file : str
+            Path to a pickle file containing a saved model, or a csv file
+            containing equations.
+
+        binary_operators : list[str]
+            The same binary operators used when creating the model.
+            Not needed if loading from a pickle file.
+
+        unary_operators : list[str]
+            The same unary operators used when creating the model.
+            Not needed if loading from a pickle file.
+
+        n_features_in : int
+            Number of features passed to the model.
+            Not needed if loading from a pickle file.
+
+        feature_names_in : list[str]
+            Names of the features passed to the model.
+            Not needed if loading from a pickle file.
+
+        selection_mask : list[bool]
+            If using select_k_features, you must pass `model.selection_mask_` here.
+            Not needed if loading from a pickle file.
+
+        nout : int, default=1
+            Number of outputs of the model.
+            Not needed if loading from a pickle file.
+
+        pysr_kwargs : dict
+            Any other keyword arguments to initialize the PySRRegressor object.
+            These will overwrite those stored in the pickle file.
+            Not needed if loading from a pickle file.
+
+        Returns
+        -------
+        model : PySRRegressor
+            The model with fitted equations.
+        """
+        if os.path.splitext(equation_file)[1] != ".pkl":
+            pkl_filename = _csv_filename_to_pkl_filename(equation_file)
+        else:
+            pkl_filename = equation_file
+
+        # Try to load model from <equation_file>.pkl
+        print(f"Checking if {pkl_filename} exists...")
+        if os.path.exists(pkl_filename):
+            print(f"Loading model from {pkl_filename}")
+            assert binary_operators is None
+            assert unary_operators is None
+            assert n_features_in is None
+            with open(pkl_filename, "rb") as f:
+                model = pkl.load(f)
+            # Update any parameters if necessary, such as
+            # extra_sympy_mappings:
+            model.set_params(**pysr_kwargs)
+            if "equations_" not in model.__dict__ or model.equations_ is None:
+                model.refresh()
+
+            return model
+
+        # Else, we re-create it.
+        print(
+            f"{equation_file} does not exist, "
+            "so we must create the model from scratch."
+        )
+        assert binary_operators is not None
+        assert unary_operators is not None
+        assert n_features_in is not None
+
+        # TODO: copy .bkup file if exists.
+        model = cls(
+            equation_file=equation_file,
+            binary_operators=binary_operators,
+            unary_operators=unary_operators,
+            **pysr_kwargs,
+        )
+
+        model.nout_ = nout
+        model.n_features_in_ = n_features_in
+
+        if feature_names_in is None:
+            model.feature_names_in_ = [f"x{i}" for i in range(n_features_in)]
+        else:
+            assert len(feature_names_in) == n_features_in
+            model.feature_names_in_ = feature_names_in
+
+        if selection_mask is None:
+            model.selection_mask_ = np.ones(n_features_in, dtype=bool)
+        else:
+            model.selection_mask_ = selection_mask
+
+        model.refresh(checkpoint_file=equation_file)
+
+        return model
+
     def __repr__(self):
         """
         Prints all current equations fitted by the model.
@@ -2091,107 +2204,3 @@ def _csv_filename_to_pkl_filename(csv_filename) -> str:
     pkl_basename = base + ".pkl"
 
     return os.path.join(dirname, pkl_basename)
-
-
-def load(
-    equation_file,
-    *,
-    binary_operators=None,
-    unary_operators=None,
-    n_features_in=None,
-    feature_names_in=None,
-    selection_mask=None,
-    nout=1,
-    **pysr_kwargs,
-):
-    """
-    Create a model from equations stored as a csv file
-
-    Parameters
-    ----------
-    equation_file : str
-        Path to a csv file containing equations, or a pickle file
-        containing the model.
-
-    binary_operators : list[str], default=["+", "-", "*", "/"]
-        The same binary operators used when creating the model.
-
-    unary_operators : list[str], default=None
-        The same unary operators used when creating the model.
-
-    n_features_in : int
-        Number of features passed to the model.
-
-    feature_names_in : list[str], default=None
-        Names of the features passed to the model.
-
-    selection_mask : list[bool], default=None
-        If using select_k_features, you must pass `model.selection_mask_` here.
-
-    nout : int, default=1
-        Number of outputs of the model.
-
-    pysr_kwargs : dict
-        Any other keyword arguments to initialize the PySRRegressor object.
-        These will overwrite those stored in the pickle file.
-
-    Returns
-    -------
-    model : PySRRegressor
-        The model with fitted equations.
-    """
-    if os.path.splitext(equation_file)[1] != ".pkl":
-        pkl_filename = _csv_filename_to_pkl_filename(equation_file)
-    else:
-        pkl_filename = equation_file
-
-    # Try to load model from <equation_file>.pkl
-    print(f"Checking if {pkl_filename} exists...")
-    if os.path.exists(pkl_filename):
-        print(f"Loading model from {pkl_filename}")
-        assert binary_operators is None
-        assert unary_operators is None
-        assert n_features_in is None
-        with open(pkl_filename, "rb") as f:
-            model = pkl.load(f)
-        # Update any parameters if necessary, such as
-        # extra_sympy_mappings:
-        model.set_params(**pysr_kwargs)
-        if "equations_" not in model.__dict__ or model.equations_ is None:
-            model.refresh()
-
-        return model
-
-    # Else, we re-create it.
-    print(
-        f"{equation_file} does not exist, " "so we must create the model from scratch."
-    )
-    assert binary_operators is not None
-    assert unary_operators is not None
-    assert n_features_in is not None
-
-    # TODO: copy .bkup file if exists.
-    model = PySRRegressor(
-        equation_file=equation_file,
-        binary_operators=binary_operators,
-        unary_operators=unary_operators,
-        **pysr_kwargs,
-    )
-
-    model.nout_ = nout
-    model.n_features_in_ = n_features_in
-
-    if feature_names_in is None:
-        model.feature_names_in_ = [f"x{i}" for i in range(n_features_in)]
-    else:
-        assert len(feature_names_in) == n_features_in
-        model.feature_names_in_ = feature_names_in
-
-    if selection_mask is None:
-        model.selection_mask_ = np.ones(n_features_in, dtype=bool)
-    else:
-        model.selection_mask_ = selection_mask
-
-    model.refresh(checkpoint_file=equation_file)
-
-    return model
diff --git a/test/test.py b/test/test.py
index dd1ece29f..fcde9ff9e 100644
--- a/test/test.py
+++ b/test/test.py
@@ -4,7 +4,7 @@
 import unittest
 import numpy as np
 from sklearn import model_selection
-from pysr import PySRRegressor, load
+from pysr import PySRRegressor
 from pysr.sr import (
     run_feature_selection,
     _handle_feature_selection,
@@ -300,7 +300,7 @@ def test_load_model(self):
             equation_filename = str(rand_dir / "equation.csv")
             with open(equation_filename + (".bkup" if from_backup else ""), "w") as f:
                 f.write(csv_file_data)
-            model = load(
+            model = PySRRegressor.from_file(
                 equation_filename,
                 n_features_in=5,
                 feature_names_in=["f0", "f1", "f2", "f3", "f4"],
@@ -334,7 +334,7 @@ def test_load_model_simple(self):
 
         # lambda functions are removed from the pickling, so we need
         # to pass it during the loading:
-        model2 = load(
+        model2 = PySRRegressor.from_file(
             model.equation_file_, extra_sympy_mappings={"sq": lambda x: x**2}
         )
 
@@ -346,7 +346,7 @@ def test_load_model_simple(self):
                 os.remove(file_to_delete)
 
         pickle_file = rand_dir / "equations.pkl"
-        model3 = load(
+        model3 = PySRRegressor.from_file(
             model.equation_file_, extra_sympy_mappings={"sq": lambda x: x**2}
         )
         np.testing.assert_allclose(model.predict(self.X), model3.predict(self.X))