fix: avoid empty checkpoints for custom jax exports

Author: MilesCranmerBot

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"

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."""