Change model load to classmethod

pysr/__init__.py

@@ -6,7 +6,6 @@ from .sr import (
     best_tex,
     best_callable,
     best_row,
-    load,
 )
 from .julia_helpers import install
 from .feynman_problems import Problem, FeynmanProblem

pysr/sr.py

@@ -810,6 +810,119 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
                         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

test/test.py

@@ -4,7 +4,7 @@ import inspect
 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 @@ class TestPipeline(unittest.TestCase):
             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 @@ class TestPipeline(unittest.TestCase):
 
         # 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 @@ class TestPipeline(unittest.TestCase):
                 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))