From d969a4d1ade5b6fd23fbdeb02b738d05cdc7dc3e Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Mon, 1 Jun 2026 15:24:37 +0000 Subject: [PATCH 1/4] fix: avoid empty checkpoints for custom jax exports --- pysr/sr.py | 226 ++++++++++++++++++++++++++---------------- pysr/test/test_jax.py | 73 ++++++++++++++ 2 files changed, 212 insertions(+), 87 deletions(-) diff --git a/pysr/sr.py b/pysr/sr.py index 050ff6cb3..0c494e564 100644 --- a/pysr/sr.py +++ b/pysr/sr.py @@ -1288,78 +1288,77 @@ def from_file( pkl_filename = Path(run_directory) / "checkpoint.pkl" if pkl_filename.exists(): pysr_logger.info(f"Attempting to load model from {pkl_filename}...") - assert binary_operators is None - assert unary_operators is None - assert operators is None - assert n_features_in is None - with open(pkl_filename, "rb") as f: - model = cast("PySRRegressor", pkl.load(f)) + model = cls._load_checkpoint(pkl_filename) + if model is not None: + assert binary_operators is None + assert unary_operators is None + assert operators is None + assert n_features_in is None - # Update any parameters if necessary, such as - # extra_sympy_mappings: - model.set_params(**pysr_kwargs) + # 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() + if ( + "equations_" not in model.__dict__ + or model.equations_ is None + or cls._equations_missing_export_formats(model.equations_) + ): + model.refresh() - if model.expression_spec is not None: - warnings.warn( - "Loading model from checkpoint file with a non-default expression spec " - "is not fully supported as it relies on dynamic objects. This may result in unexpected behavior.", - ) + if model.expression_spec is not None: + warnings.warn( + "Loading model from checkpoint file with a non-default expression spec " + "is not fully supported as it relies on dynamic objects. This may result in unexpected behavior.", + ) - return model - else: - pysr_logger.info( - f"Checkpoint file {pkl_filename} does not exist. " - "Attempting to recreate model from scratch..." + return model + pysr_logger.info( + f"Checkpoint file {pkl_filename} does not exist or could not be loaded. " + "Attempting to recreate model from CSV backups..." + ) + csv_filename = Path(run_directory) / "hall_of_fame.csv" + csv_filename_bak = Path(run_directory) / "hall_of_fame.csv.bak" + if not csv_filename.exists() and not csv_filename_bak.exists(): + raise FileNotFoundError( + f"Hall of fame file `{csv_filename}` or `{csv_filename_bak}` does not exist. " + "Please pass a `run_directory` containing a valid checkpoint file." ) - csv_filename = Path(run_directory) / "hall_of_fame.csv" - csv_filename_bak = Path(run_directory) / "hall_of_fame.csv.bak" - if not csv_filename.exists() and not csv_filename_bak.exists(): - raise FileNotFoundError( - f"Hall of fame file `{csv_filename}` or `{csv_filename_bak}` does not exist. " - "Please pass a `run_directory` containing a valid checkpoint file." - ) - if ( - operators is None - and binary_operators is None - and unary_operators is None - ): - raise ValueError( - "When recreating a model from CSV backups you must provide either " - "`operators` or legacy `binary_operators`/`unary_operators`." - ) - assert n_features_in is not None - model = cls( - binary_operators=binary_operators, - unary_operators=unary_operators, - operators=operators, - **pysr_kwargs, + if operators is None and binary_operators is None and unary_operators is None: + raise ValueError( + "When recreating a model from CSV backups you must provide either " + "`operators` or legacy `binary_operators`/`unary_operators`." ) - model.nout_ = nout - model.n_features_in_ = n_features_in + assert n_features_in is not None + model = cls( + binary_operators=binary_operators, + unary_operators=unary_operators, + operators=operators, + **pysr_kwargs, + ) + model.nout_ = nout + model.n_features_in_ = n_features_in - if feature_names_in is None: - model.feature_names_in_ = np.array( - [f"x{i}" for i in range(n_features_in)] - ) - model.display_feature_names_in_ = np.array( - [f"x{_subscriptify(i)}" for i in range(n_features_in)] - ) - else: - assert len(feature_names_in) == n_features_in - model.feature_names_in_ = feature_names_in - model.display_feature_names_in_ = feature_names_in + if feature_names_in is None: + model.feature_names_in_ = np.array( + [f"x{i}" for i in range(n_features_in)] + ) + model.display_feature_names_in_ = np.array( + [f"x{_subscriptify(i)}" for i in range(n_features_in)] + ) + else: + assert len(feature_names_in) == n_features_in + model.feature_names_in_ = feature_names_in + model.display_feature_names_in_ = feature_names_in - if selection_mask is None: - model.selection_mask_ = np.ones(n_features_in, dtype=np.bool_) - else: - model.selection_mask_ = selection_mask + if selection_mask is None: + model.selection_mask_ = np.ones(n_features_in, dtype=np.bool_) + else: + model.selection_mask_ = selection_mask - model.refresh(run_directory=run_directory) + model.refresh(run_directory=run_directory) - return model + return model def __repr__(self) -> str: """ @@ -1424,8 +1423,12 @@ def __getstate__(self) -> dict[str, Any]: show_pickle_warning = not ( "show_pickle_warnings_" in state and not state["show_pickle_warnings_"] ) - state_keys_containing_lambdas = ["extra_sympy_mappings", "extra_torch_mappings"] - for state_key in state_keys_containing_lambdas: + state_keys_to_clear = ( + "extra_sympy_mappings", + "extra_jax_mappings", + "extra_torch_mappings", + ) + for state_key in state_keys_to_clear: warn_msg = ( f"`{state_key}` cannot be pickled and will be removed from the " "serialized instance. When loading the model, please redefine " @@ -1436,8 +1439,7 @@ def __getstate__(self) -> dict[str, Any]: warnings.warn(warn_msg) else: pysr_logger.debug(warn_msg) - state_keys_to_clear = state_keys_containing_lambdas - state_keys_to_clear.append("logger_") + state_keys_to_clear = (*state_keys_to_clear, "logger_") pickled_state = { key: (None if key in state_keys_to_clear else value) for key, value in state.items() @@ -1448,38 +1450,88 @@ def __getstate__(self) -> dict[str, Any]: pickled_state["output_torch_format"] = False pickled_state["output_jax_format"] = False if self.nout_ == 1: - pickled_columns = ~pickled_state["equations_"].columns.isin( - ["jax_format", "torch_format"] - ) - pickled_state["equations_"] = ( - pickled_state["equations_"].loc[:, pickled_columns].copy() + pickled_state["equations_"] = self._drop_equation_columns( + pickled_state["equations_"], ["jax_format", "torch_format"] ) else: - pickled_columns = [ - ~dataframe.columns.isin(["jax_format", "torch_format"]) - for dataframe in pickled_state["equations_"] - ] - pickled_state["equations_"] = [ - dataframe.loc[:, signle_pickled_columns] - for dataframe, signle_pickled_columns in zip( - pickled_state["equations_"], pickled_columns - ) - ] + pickled_state["equations_"] = self._drop_equation_columns( + pickled_state["equations_"], ["jax_format", "torch_format"] + ) + try: + pkl.dumps(pickled_state["equations_"]) + except Exception as e: + warn_msg = ( + "`equations_` export formats cannot be pickled and will be " + "removed from the serialized instance. When loading the model, " + "please redefine custom mappings at runtime." + ) + if show_pickle_warning: + warnings.warn(warn_msg) + else: + pysr_logger.debug(f"{warn_msg} Error: {e}") + pickled_state["equations_"] = self._drop_equation_columns( + pickled_state["equations_"], ["sympy_format", "lambda_format"] + ) return pickled_state + @staticmethod + def _drop_equation_columns(equations, columns: list[str]): + if isinstance(equations, list): + return [ + dataframe.loc[:, ~dataframe.columns.isin(columns)].copy() + for dataframe in equations + ] + return equations.loc[:, ~equations.columns.isin(columns)].copy() + + @staticmethod + def _equations_missing_export_formats(equations) -> bool: + required_columns = {"sympy_format", "lambda_format"} + if isinstance(equations, list): + return any( + not required_columns.issubset(dataframe.columns) + for dataframe in equations + ) + return not required_columns.issubset(equations.columns) + def _checkpoint(self): """Save 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(self.get_pkl_filename(), "wb") as f: - try: + pkl_filename = self.get_pkl_filename() + tmp_filename = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", dir=pkl_filename.parent, delete=False + ) as f: + tmp_filename = Path(f.name) pkl.dump(self, f) - except Exception as e: - pysr_logger.debug(f"Error checkpointing model: {e}") - self.show_pickle_warnings_ = True + os.replace(tmp_filename, pkl_filename) + except Exception as e: + pysr_logger.debug(f"Error checkpointing model: {e}") + if tmp_filename is not None: + tmp_filename.unlink(missing_ok=True) + finally: + self.show_pickle_warnings_ = True + + @staticmethod + def _load_checkpoint(pkl_filename: Path) -> "PySRRegressor" | None: + if pkl_filename.stat().st_size == 0: + pysr_logger.warning( + f"Checkpoint file {pkl_filename} is empty. " + "Attempting to recreate model from CSV backups..." + ) + return None + try: + with open(pkl_filename, "rb") as f: + return cast("PySRRegressor", pkl.load(f)) + except (EOFError, pkl.UnpicklingError) as e: + pysr_logger.warning( + f"Could not load checkpoint file {pkl_filename}: {e}. " + "Attempting to recreate model from CSV backups..." + ) + return None def get_pkl_filename(self) -> Path: path = Path(self.output_directory_) / self.run_id_ / "checkpoint.pkl" diff --git a/pysr/test/test_jax.py b/pysr/test/test_jax.py index d261afcf5..42318cbe7 100644 --- a/pysr/test/test_jax.py +++ b/pysr/test/test_jax.py @@ -1,3 +1,4 @@ +import tempfile import unittest from functools import partial from pathlib import Path @@ -149,6 +150,78 @@ def cos_approx(x): jax_output = jax_prediction(X.values) np.testing.assert_almost_equal(y.values, jax_output, decimal=3) + def test_checkpoint_custom_jax_mapping(self): + sp_cos_approx = sympy.Function("cos_approx") + jax_mapping = { + sp_cos_approx: "(lambda x: 1 - x**2 / 2 + x**4 / 24 + x**6 / 720)" + } + + with tempfile.TemporaryDirectory() as tmpdir: + run_dir = Path(tmpdir) / "custom_jax" + run_dir.mkdir() + pd.DataFrame( + { + "Complexity": [1], + "Loss": [0.0], + "Equation": ["cos_approx(x0)"], + } + ).to_csv(run_dir / "hall_of_fame.csv") + + model = PySRRegressor( + progress=False, + unary_operators=[ + "cos_approx(x) = 1 - x^2 / 2 + x^4 / 24 + x^6 / 720" + ], + extra_sympy_mappings={"cos_approx": sp_cos_approx}, + extra_jax_mappings=jax_mapping, + output_jax_format=True, + ) + model.output_directory_ = tmpdir + model.run_id_ = "custom_jax" + model.nout_ = 1 + model.n_features_in_ = 1 + model.feature_names_in_ = np.array(["x0"]) + model.display_feature_names_in_ = np.array(["x0"]) + model.selection_mask_ = np.ones(1, dtype=np.bool_) + model.refresh() + + model._checkpoint() + assert (run_dir / "checkpoint.pkl").stat().st_size > 0 + + model2 = PySRRegressor.from_file( + run_directory=run_dir, + extra_sympy_mappings={"cos_approx": sp_cos_approx}, + extra_jax_mappings=jax_mapping, + ) + jax_format = model2.jax(index=0) + X = self.jnp.array([[0.0], [1.0], [2.0]]) + expected = 1 - X[:, 0] ** 2 / 2 + X[:, 0] ** 4 / 24 + X[:, 0] ** 6 / 720 + np.testing.assert_allclose( + np.array(jax_format["callable"](X, jax_format["parameters"])), + np.array(expected), + ) + + def test_from_file_empty_checkpoint_falls_back_to_csv(self): + with tempfile.TemporaryDirectory() as tmpdir: + run_dir = Path(tmpdir) / "empty_checkpoint" + run_dir.mkdir() + (run_dir / "checkpoint.pkl").touch() + pd.DataFrame( + { + "Complexity": [1], + "Loss": [0.0], + "Equation": ["x0"], + } + ).to_csv(run_dir / "hall_of_fame.csv") + + model = PySRRegressor.from_file( + run_directory=run_dir, + binary_operators=["+"], + unary_operators=[], + n_features_in=1, + ) + assert str(model.sympy(index=0)) == "x0" + def runtests(just_tests=False): """Run all tests in test_jax.py.""" From 31210d8cb248af91e86836a9446327ee5912c61e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 1 Jun 2026 15:25:27 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pysr/sr.py | 4 +--- pysr/test/test_jax.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/pysr/sr.py b/pysr/sr.py index 0c494e564..27b28eb89 100644 --- a/pysr/sr.py +++ b/pysr/sr.py @@ -1340,9 +1340,7 @@ def from_file( model.n_features_in_ = n_features_in if feature_names_in is None: - model.feature_names_in_ = np.array( - [f"x{i}" for i in range(n_features_in)] - ) + model.feature_names_in_ = np.array([f"x{i}" for i in range(n_features_in)]) model.display_feature_names_in_ = np.array( [f"x{_subscriptify(i)}" for i in range(n_features_in)] ) diff --git a/pysr/test/test_jax.py b/pysr/test/test_jax.py index 42318cbe7..d3d44b4b2 100644 --- a/pysr/test/test_jax.py +++ b/pysr/test/test_jax.py @@ -169,9 +169,7 @@ def test_checkpoint_custom_jax_mapping(self): model = PySRRegressor( progress=False, - unary_operators=[ - "cos_approx(x) = 1 - x^2 / 2 + x^4 / 24 + x^6 / 720" - ], + unary_operators=["cos_approx(x) = 1 - x^2 / 2 + x^4 / 24 + x^6 / 720"], extra_sympy_mappings={"cos_approx": sp_cos_approx}, extra_jax_mappings=jax_mapping, output_jax_format=True, From 2e59bfd0726c412eb649825720c9f937d7480cc9 Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Mon, 1 Jun 2026 20:59:06 +0000 Subject: [PATCH 3/4] refactor: move checkpoint helpers out of regressor Co-authored-by: Miles Cranmer --- pysr/_checkpoint.py | 125 ++++++++++++++++++++++++++++++++++++++++++++ pysr/sr.py | 116 ++++------------------------------------ 2 files changed, 135 insertions(+), 106 deletions(-) create mode 100644 pysr/_checkpoint.py diff --git a/pysr/_checkpoint.py b/pysr/_checkpoint.py new file mode 100644 index 000000000..78e536c75 --- /dev/null +++ b/pysr/_checkpoint.py @@ -0,0 +1,125 @@ +"""Checkpoint and pickle helpers for PySRRegressor.""" + +from __future__ import annotations + +import logging +import os +import pickle as pkl +import tempfile +import warnings +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +if TYPE_CHECKING: + import pandas as pd + + from .sr import PySRRegressor + + +logger = logging.getLogger(__name__) + + +def get_regressor_pickle_state(state: dict[str, Any]) -> dict[str, Any]: + """Return a pickle-safe version of a PySRRegressor state dictionary.""" + show_pickle_warning = not ( + "show_pickle_warnings_" in state and not state["show_pickle_warnings_"] + ) + state_keys_to_clear = ( + "extra_sympy_mappings", + "extra_jax_mappings", + "extra_torch_mappings", + ) + for state_key in state_keys_to_clear: + warn_msg = ( + 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." + ) + if state[state_key] is not None: + if show_pickle_warning: + warnings.warn(warn_msg) + else: + logger.debug(warn_msg) + state_keys_to_clear = (*state_keys_to_clear, "logger_") + pickled_state = { + key: (None if key in state_keys_to_clear else value) + for key, value in state.items() + } + if ("equations_" in pickled_state) and (pickled_state["equations_"] is not None): + pickled_state["output_torch_format"] = False + pickled_state["output_jax_format"] = False + pickled_state["equations_"] = drop_equation_columns( + pickled_state["equations_"], ["jax_format", "torch_format"] + ) + try: + pkl.dumps(pickled_state["equations_"]) + except Exception as e: + warn_msg = ( + "`equations_` export formats cannot be pickled and will be " + "removed from the serialized instance. When loading the model, " + "please redefine custom mappings at runtime." + ) + if show_pickle_warning: + warnings.warn(warn_msg) + else: + logger.debug(f"{warn_msg} Error: {e}") + pickled_state["equations_"] = drop_equation_columns( + pickled_state["equations_"], ["sympy_format", "lambda_format"] + ) + return pickled_state + + +def drop_equation_columns( + equations: pd.DataFrame | list[pd.DataFrame], + columns: list[str], +) -> pd.DataFrame | list[pd.DataFrame]: + if isinstance(equations, list): + return [ + dataframe.loc[:, ~dataframe.columns.isin(columns)].copy() + for dataframe in equations + ] + return equations.loc[:, ~equations.columns.isin(columns)].copy() + + +def equations_missing_export_formats( + equations: pd.DataFrame | list[pd.DataFrame], +) -> bool: + required_columns = {"sympy_format", "lambda_format"} + if isinstance(equations, list): + return any( + not required_columns.issubset(dataframe.columns) for dataframe in equations + ) + return not required_columns.issubset(equations.columns) + + +def save_checkpoint(model: PySRRegressor, pkl_filename: Path) -> None: + tmp_filename = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", dir=pkl_filename.parent, delete=False + ) as f: + tmp_filename = Path(f.name) + pkl.dump(model, f) + os.replace(tmp_filename, pkl_filename) + except Exception as e: + logger.debug(f"Error checkpointing model: {e}") + if tmp_filename is not None: + tmp_filename.unlink(missing_ok=True) + + +def load_checkpoint(pkl_filename: Path) -> PySRRegressor | None: + if pkl_filename.stat().st_size == 0: + logger.warning( + f"Checkpoint file {pkl_filename} is empty. " + "Attempting to recreate model from CSV backups..." + ) + return None + try: + with open(pkl_filename, "rb") as f: + return cast("PySRRegressor", pkl.load(f)) + except (EOFError, pkl.UnpicklingError) as e: + logger.warning( + f"Could not load checkpoint file {pkl_filename}: {e}. " + "Attempting to recreate model from CSV backups..." + ) + return None diff --git a/pysr/sr.py b/pysr/sr.py index 27b28eb89..d28e737a7 100644 --- a/pysr/sr.py +++ b/pysr/sr.py @@ -5,7 +5,6 @@ import copy import logging import os -import pickle as pkl import re import sys import tempfile @@ -26,6 +25,12 @@ from sklearn.utils.validation import _check_feature_names_in # type: ignore from sklearn.utils.validation import check_is_fitted +from ._checkpoint import ( + equations_missing_export_formats, + get_regressor_pickle_state, + load_checkpoint, + save_checkpoint, +) from .denoising import denoise, multi_denoise from .deprecated import DEPRECATED_KWARGS from .export_latex import ( @@ -1288,7 +1293,7 @@ def from_file( pkl_filename = Path(run_directory) / "checkpoint.pkl" if pkl_filename.exists(): pysr_logger.info(f"Attempting to load model from {pkl_filename}...") - model = cls._load_checkpoint(pkl_filename) + model = load_checkpoint(pkl_filename) if model is not None: assert binary_operators is None assert unary_operators is None @@ -1302,7 +1307,7 @@ def from_file( if ( "equations_" not in model.__dict__ or model.equations_ is None - or cls._equations_missing_export_formats(model.equations_) + or equations_missing_export_formats(model.equations_) ): model.refresh() @@ -1417,79 +1422,7 @@ def __getstate__(self) -> dict[str, Any]: `pickle.dumps()`. However, some attributes do not support pickling and need to be hidden, such as the JAX and Torch representations. """ - state = self.__dict__ - show_pickle_warning = not ( - "show_pickle_warnings_" in state and not state["show_pickle_warnings_"] - ) - state_keys_to_clear = ( - "extra_sympy_mappings", - "extra_jax_mappings", - "extra_torch_mappings", - ) - for state_key in state_keys_to_clear: - warn_msg = ( - 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." - ) - if state[state_key] is not None: - if show_pickle_warning: - warnings.warn(warn_msg) - else: - pysr_logger.debug(warn_msg) - state_keys_to_clear = (*state_keys_to_clear, "logger_") - pickled_state = { - key: (None if key in state_keys_to_clear else value) - for key, value in state.items() - } - 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: - pickled_state["equations_"] = self._drop_equation_columns( - pickled_state["equations_"], ["jax_format", "torch_format"] - ) - else: - pickled_state["equations_"] = self._drop_equation_columns( - pickled_state["equations_"], ["jax_format", "torch_format"] - ) - try: - pkl.dumps(pickled_state["equations_"]) - except Exception as e: - warn_msg = ( - "`equations_` export formats cannot be pickled and will be " - "removed from the serialized instance. When loading the model, " - "please redefine custom mappings at runtime." - ) - if show_pickle_warning: - warnings.warn(warn_msg) - else: - pysr_logger.debug(f"{warn_msg} Error: {e}") - pickled_state["equations_"] = self._drop_equation_columns( - pickled_state["equations_"], ["sympy_format", "lambda_format"] - ) - return pickled_state - - @staticmethod - def _drop_equation_columns(equations, columns: list[str]): - if isinstance(equations, list): - return [ - dataframe.loc[:, ~dataframe.columns.isin(columns)].copy() - for dataframe in equations - ] - return equations.loc[:, ~equations.columns.isin(columns)].copy() - - @staticmethod - def _equations_missing_export_formats(equations) -> bool: - required_columns = {"sympy_format", "lambda_format"} - if isinstance(equations, list): - return any( - not required_columns.issubset(dataframe.columns) - for dataframe in equations - ) - return not required_columns.issubset(equations.columns) + return get_regressor_pickle_state(self.__dict__) def _checkpoint(self): """Save the model's current state to a checkpoint file. @@ -1497,40 +1430,11 @@ def _checkpoint(self): This should only be used internally by PySRRegressor. """ self.show_pickle_warnings_ = False - pkl_filename = self.get_pkl_filename() - tmp_filename = None try: - with tempfile.NamedTemporaryFile( - mode="wb", dir=pkl_filename.parent, delete=False - ) as f: - tmp_filename = Path(f.name) - pkl.dump(self, f) - os.replace(tmp_filename, pkl_filename) - except Exception as e: - pysr_logger.debug(f"Error checkpointing model: {e}") - if tmp_filename is not None: - tmp_filename.unlink(missing_ok=True) + save_checkpoint(self, self.get_pkl_filename()) finally: self.show_pickle_warnings_ = True - @staticmethod - def _load_checkpoint(pkl_filename: Path) -> "PySRRegressor" | None: - if pkl_filename.stat().st_size == 0: - pysr_logger.warning( - f"Checkpoint file {pkl_filename} is empty. " - "Attempting to recreate model from CSV backups..." - ) - return None - try: - with open(pkl_filename, "rb") as f: - return cast("PySRRegressor", pkl.load(f)) - except (EOFError, pkl.UnpicklingError) as e: - pysr_logger.warning( - f"Could not load checkpoint file {pkl_filename}: {e}. " - "Attempting to recreate model from CSV backups..." - ) - return None - def get_pkl_filename(self) -> Path: path = Path(self.output_directory_) / self.run_id_ / "checkpoint.pkl" path.parent.mkdir(parents=True, exist_ok=True) From aea7cde9e6f885b611e88ed86630340eb8b059e0 Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Mon, 1 Jun 2026 21:12:25 +0000 Subject: [PATCH 4/4] refactor: keep checkpoint helpers on regressor Co-authored-by: Miles Cranmer --- pysr/_checkpoint.py | 125 ------------------------------------------ pysr/sr.py | 118 +++++++++++++++++++++++++++++++++++---- pysr/test/test_jax.py | 27 ++++++++- 3 files changed, 132 insertions(+), 138 deletions(-) delete mode 100644 pysr/_checkpoint.py diff --git a/pysr/_checkpoint.py b/pysr/_checkpoint.py deleted file mode 100644 index 78e536c75..000000000 --- a/pysr/_checkpoint.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Checkpoint and pickle helpers for PySRRegressor.""" - -from __future__ import annotations - -import logging -import os -import pickle as pkl -import tempfile -import warnings -from pathlib import Path -from typing import TYPE_CHECKING, Any, cast - -if TYPE_CHECKING: - import pandas as pd - - from .sr import PySRRegressor - - -logger = logging.getLogger(__name__) - - -def get_regressor_pickle_state(state: dict[str, Any]) -> dict[str, Any]: - """Return a pickle-safe version of a PySRRegressor state dictionary.""" - show_pickle_warning = not ( - "show_pickle_warnings_" in state and not state["show_pickle_warnings_"] - ) - state_keys_to_clear = ( - "extra_sympy_mappings", - "extra_jax_mappings", - "extra_torch_mappings", - ) - for state_key in state_keys_to_clear: - warn_msg = ( - 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." - ) - if state[state_key] is not None: - if show_pickle_warning: - warnings.warn(warn_msg) - else: - logger.debug(warn_msg) - state_keys_to_clear = (*state_keys_to_clear, "logger_") - pickled_state = { - key: (None if key in state_keys_to_clear else value) - for key, value in state.items() - } - if ("equations_" in pickled_state) and (pickled_state["equations_"] is not None): - pickled_state["output_torch_format"] = False - pickled_state["output_jax_format"] = False - pickled_state["equations_"] = drop_equation_columns( - pickled_state["equations_"], ["jax_format", "torch_format"] - ) - try: - pkl.dumps(pickled_state["equations_"]) - except Exception as e: - warn_msg = ( - "`equations_` export formats cannot be pickled and will be " - "removed from the serialized instance. When loading the model, " - "please redefine custom mappings at runtime." - ) - if show_pickle_warning: - warnings.warn(warn_msg) - else: - logger.debug(f"{warn_msg} Error: {e}") - pickled_state["equations_"] = drop_equation_columns( - pickled_state["equations_"], ["sympy_format", "lambda_format"] - ) - return pickled_state - - -def drop_equation_columns( - equations: pd.DataFrame | list[pd.DataFrame], - columns: list[str], -) -> pd.DataFrame | list[pd.DataFrame]: - if isinstance(equations, list): - return [ - dataframe.loc[:, ~dataframe.columns.isin(columns)].copy() - for dataframe in equations - ] - return equations.loc[:, ~equations.columns.isin(columns)].copy() - - -def equations_missing_export_formats( - equations: pd.DataFrame | list[pd.DataFrame], -) -> bool: - required_columns = {"sympy_format", "lambda_format"} - if isinstance(equations, list): - return any( - not required_columns.issubset(dataframe.columns) for dataframe in equations - ) - return not required_columns.issubset(equations.columns) - - -def save_checkpoint(model: PySRRegressor, pkl_filename: Path) -> None: - tmp_filename = None - try: - with tempfile.NamedTemporaryFile( - mode="wb", dir=pkl_filename.parent, delete=False - ) as f: - tmp_filename = Path(f.name) - pkl.dump(model, f) - os.replace(tmp_filename, pkl_filename) - except Exception as e: - logger.debug(f"Error checkpointing model: {e}") - if tmp_filename is not None: - tmp_filename.unlink(missing_ok=True) - - -def load_checkpoint(pkl_filename: Path) -> PySRRegressor | None: - if pkl_filename.stat().st_size == 0: - logger.warning( - f"Checkpoint file {pkl_filename} is empty. " - "Attempting to recreate model from CSV backups..." - ) - return None - try: - with open(pkl_filename, "rb") as f: - return cast("PySRRegressor", pkl.load(f)) - except (EOFError, pkl.UnpicklingError) as e: - logger.warning( - f"Could not load checkpoint file {pkl_filename}: {e}. " - "Attempting to recreate model from CSV backups..." - ) - return None diff --git a/pysr/sr.py b/pysr/sr.py index d28e737a7..343800888 100644 --- a/pysr/sr.py +++ b/pysr/sr.py @@ -5,6 +5,7 @@ import copy import logging import os +import pickle as pkl import re import sys import tempfile @@ -25,12 +26,6 @@ from sklearn.utils.validation import _check_feature_names_in # type: ignore from sklearn.utils.validation import check_is_fitted -from ._checkpoint import ( - equations_missing_export_formats, - get_regressor_pickle_state, - load_checkpoint, - save_checkpoint, -) from .denoising import denoise, multi_denoise from .deprecated import DEPRECATED_KWARGS from .export_latex import ( @@ -1293,7 +1288,7 @@ def from_file( pkl_filename = Path(run_directory) / "checkpoint.pkl" if pkl_filename.exists(): pysr_logger.info(f"Attempting to load model from {pkl_filename}...") - model = load_checkpoint(pkl_filename) + model = cls._load_checkpoint(pkl_filename) if model is not None: assert binary_operators is None assert unary_operators is None @@ -1307,7 +1302,7 @@ def from_file( if ( "equations_" not in model.__dict__ or model.equations_ is None - or equations_missing_export_formats(model.equations_) + or cls._equations_missing_export_formats(model.equations_) ): model.refresh() @@ -1422,7 +1417,97 @@ def __getstate__(self) -> dict[str, Any]: `pickle.dumps()`. However, some attributes do not support pickling and need to be hidden, such as the JAX and Torch representations. """ - return get_regressor_pickle_state(self.__dict__) + return self._get_pickle_state() + + def _get_pickle_state(self) -> dict[str, Any]: + show_pickle_warning = not ( + "show_pickle_warnings_" in self.__dict__ and not self.show_pickle_warnings_ + ) + warn_or_log = warnings.warn if show_pickle_warning else pysr_logger.debug + state_keys_to_clear = ( + "extra_sympy_mappings", + "extra_jax_mappings", + "extra_torch_mappings", + ) + for state_key in state_keys_to_clear: + warn_msg = ( + 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." + ) + if getattr(self, state_key) is not None: + warn_or_log(warn_msg) + pickled_state_keys_to_clear = (*state_keys_to_clear, "logger_") + pickled_state = { + key: (None if key in pickled_state_keys_to_clear else value) + for key, value in self.__dict__.items() + } + if ("equations_" in pickled_state) and ( + pickled_state["equations_"] is not None + ): + pickled_state["output_torch_format"] = False + pickled_state["output_jax_format"] = False + pickled_state["equations_"] = self._drop_equation_columns( + pickled_state["equations_"], ["jax_format", "torch_format"] + ) + try: + pkl.dumps(pickled_state["equations_"]) + except Exception as e: + warn_msg = ( + "`equations_` export formats cannot be pickled and will be " + "removed from the serialized instance. When loading the model, " + "please redefine custom mappings at runtime." + ) + warn_or_log( + warn_msg if show_pickle_warning else f"{warn_msg} Error: {e}" + ) + pickled_state["equations_"] = self._drop_equation_columns( + pickled_state["equations_"], ["sympy_format", "lambda_format"] + ) + return pickled_state + + @staticmethod + def _drop_equation_columns( + equations: pd.DataFrame | list[pd.DataFrame], + columns: list[str], + ) -> pd.DataFrame | list[pd.DataFrame]: + all_equations = equations if isinstance(equations, list) else [equations] + dropped_equations = [ + dataframe.loc[:, ~dataframe.columns.isin(columns)].copy() + for dataframe in all_equations + ] + return ( + dropped_equations if isinstance(equations, list) else dropped_equations[0] + ) + + @staticmethod + def _equations_missing_export_formats( + equations: pd.DataFrame | list[pd.DataFrame], + ) -> bool: + required_columns = {"sympy_format", "lambda_format"} + all_equations = equations if isinstance(equations, list) else [equations] + return any( + not required_columns.issubset(dataframe.columns) + for dataframe in all_equations + ) + + @classmethod + def _load_checkpoint(cls, pkl_filename: Path) -> "PySRRegressor | None": + if pkl_filename.stat().st_size == 0: + pysr_logger.warning( + f"Checkpoint file {pkl_filename} is empty. " + "Attempting to recreate model from CSV backups..." + ) + return None + try: + with open(pkl_filename, "rb") as f: + return cast("PySRRegressor", pkl.load(f)) + except (EOFError, pkl.UnpicklingError) as e: + pysr_logger.warning( + f"Could not load checkpoint file {pkl_filename}: {e}. " + "Attempting to recreate model from CSV backups..." + ) + return None def _checkpoint(self): """Save the model's current state to a checkpoint file. @@ -1430,8 +1515,19 @@ def _checkpoint(self): This should only be used internally by PySRRegressor. """ self.show_pickle_warnings_ = False + tmp_filename = None try: - save_checkpoint(self, self.get_pkl_filename()) + pkl_filename = self.get_pkl_filename() + with tempfile.NamedTemporaryFile( + mode="wb", dir=pkl_filename.parent, delete=False + ) as f: + tmp_filename = Path(f.name) + pkl.dump(self, f) + os.replace(tmp_filename, pkl_filename) + except Exception as e: + pysr_logger.debug(f"Error checkpointing model: {e}") + if tmp_filename is not None: + tmp_filename.unlink(missing_ok=True) finally: self.show_pickle_warnings_ = True @@ -2899,7 +2995,7 @@ def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]: if self.nout_ > 1: return ret_outputs - return ret_outputs[0] + return cast(pd.DataFrame, ret_outputs[0]) def latex_table( self, diff --git a/pysr/test/test_jax.py b/pysr/test/test_jax.py index d3d44b4b2..dacd5a456 100644 --- a/pysr/test/test_jax.py +++ b/pysr/test/test_jax.py @@ -1,3 +1,4 @@ +import pickle as pkl import tempfile import unittest from functools import partial @@ -182,9 +183,10 @@ def test_checkpoint_custom_jax_mapping(self): model.display_feature_names_in_ = np.array(["x0"]) model.selection_mask_ = np.ones(1, dtype=np.bool_) model.refresh() + model.equations_.loc[:, "lambda_format"] = [lambda x: x] - model._checkpoint() - assert (run_dir / "checkpoint.pkl").stat().st_size > 0 + with open(run_dir / "checkpoint.pkl", "wb") as f: + pkl.dump(model, f) model2 = PySRRegressor.from_file( run_directory=run_dir, @@ -220,6 +222,27 @@ def test_from_file_empty_checkpoint_falls_back_to_csv(self): ) assert str(model.sympy(index=0)) == "x0" + def test_from_file_corrupt_checkpoint_falls_back_to_csv(self): + with tempfile.TemporaryDirectory() as tmpdir: + run_dir = Path(tmpdir) / "corrupt_checkpoint" + run_dir.mkdir() + (run_dir / "checkpoint.pkl").write_bytes(b"not a pickle") + pd.DataFrame( + { + "Complexity": [1], + "Loss": [0.0], + "Equation": ["x0"], + } + ).to_csv(run_dir / "hall_of_fame.csv") + + model = PySRRegressor.from_file( + run_directory=run_dir, + binary_operators=["+"], + unary_operators=[], + n_features_in=1, + ) + assert str(model.sympy(index=0)) == "x0" + def runtests(just_tests=False): """Run all tests in test_jax.py."""