fix: clarify TemplateExpressionSpec reload errors

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: MilesCranmerBot

pysr/sr.py

@@ -39,6 +39,7 @@
     AbstractExpressionSpec,
     ExpressionSpec,
     ParametricExpressionSpec,
+    TemplateExpressionSpec,
     parametric_expression_deprecation_warning,
 )
 from .feature_selection import run_feature_selection
@@ -1465,8 +1466,23 @@ def __getstate__(self) -> dict[str, Any]:
                         pickled_state["equations_"], pickled_columns
                     )
                 ]
+        pickled_state["_loaded_from_pickle_"] = False
         return pickled_state
 
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._loaded_from_pickle_ = True
+
+    def _raise_for_reloaded_template_expression_spec(self) -> None:
+        if getattr(self, "_loaded_from_pickle_", False) and isinstance(
+            self.expression_spec, TemplateExpressionSpec
+        ):
+            raise NotImplementedError(
+                "Loading models fitted with TemplateExpressionSpec is not yet supported. "
+                "Please refit the model in the current Python/Julia session before "
+                "calling predict or export methods."
+            )
+
     def _checkpoint(self):
         """Save the model's current state to a checkpoint file.
 
@@ -2640,6 +2656,7 @@ def predict(
         check_is_fitted(
             self, attributes=["selection_mask_", "feature_names_in_", "nout_"]
         )
+        self._raise_for_reloaded_template_expression_spec()
         best_equation = self.get_best(index=index)
 
         # When X is an numpy array or a pandas dataframe with a RangeIndex,
@@ -2730,6 +2747,7 @@ def sympy(self, index: int | list[int] | None = None):
         best_equation : str, list[str] of length nout_
             SymPy representation of the best equation.
         """
+        self._raise_for_reloaded_template_expression_spec()
         if not self.expression_spec_.supports_sympy:
             raise ValueError(
                 f"`expression_spec={self.expression_spec_}` does not support sympy export."
@@ -2766,6 +2784,7 @@ def latex(
         best_equation : str or list[str] of length nout_
             LaTeX expression of the best equation.
         """
+        self._raise_for_reloaded_template_expression_spec()
         if not self.expression_spec_.supports_latex:
             raise ValueError(
                 f"`expression_spec={self.expression_spec_}` does not support latex export."
@@ -2803,6 +2822,7 @@ def jax(self, index=None):
             Dictionary of callable jax function in "callable" key,
             and jax array of parameters as "parameters" key.
         """
+        self._raise_for_reloaded_template_expression_spec()
         if not self.expression_spec_.supports_jax:
             raise ValueError(
                 f"`expression_spec={self.expression_spec_}` does not support jax export."
@@ -2839,6 +2859,7 @@ def pytorch(self, index=None):
         best_equation : torch.nn.Module
             PyTorch module representing the expression.
         """
+        self._raise_for_reloaded_template_expression_spec()
         if not self.expression_spec_.supports_torch:
             raise ValueError(
                 f"`expression_spec={self.expression_spec_}` does not support torch export."
@@ -2983,6 +3004,7 @@ def latex_table(
         latex_table_str : str
             A string that will render a table in LaTeX of the equations.
         """
+        self._raise_for_reloaded_template_expression_spec()
         if not self.expression_spec_.supports_latex:
             raise ValueError(
                 f"`expression_spec={self.expression_spec_}` does not support latex export."

pysr/test/test_main.py

@@ -2246,6 +2246,26 @@ def test_process_constraints_swaps_multiplication_constraints(self):
 
 
 class TestTemplateExpressionSpec(unittest.TestCase):
+    def test_reloaded_template_expression_spec_raises_targeted_error(self):
+        model = PySRRegressor(
+            expression_spec=TemplateExpressionSpec(
+                combine="f(x)", expressions=["f"], variable_names=["x"]
+            )
+        )
+        model.selection_mask_ = np.array([True])
+        model.feature_names_in_ = np.array(["x"])
+        model.nout_ = 1
+
+        loaded = pkl.loads(pkl.dumps(model))
+
+        message = (
+            "Loading models fitted with TemplateExpressionSpec is not yet supported"
+        )
+        with self.assertRaisesRegex(NotImplementedError, message):
+            loaded.predict(np.ones((1, 1)))
+        with self.assertRaisesRegex(NotImplementedError, message):
+            loaded.sympy()
+
     def _check_macro_str(self, spec, expected_str):
         self.assertEqual(
             spec._template_macro_str().strip(), dedent(expected_str).strip()