perf: cache lambdified equations and avoid redundant export work

pysr/export.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import copy
 from collections.abc import Callable
 
 import numpy as np
@@ -29,23 +28,22 @@ def add_export_formats(
 
     Returns a new dataframe containing only the exported formats.
     """
-    output = copy.deepcopy(output)
-
     sympy_format = []
     lambda_format = []
     jax_format = []
     torch_format = []
 
-    for _, eqn_row in output.iterrows():
+    sympy_symbols = create_sympy_symbols(feature_names_in)
+
+    for equation in output["equation"]:
         eqn = pysr2sympy(
-            eqn_row["equation"],
+            equation,
             feature_names_in=feature_names_in,
             extra_sympy_mappings=extra_sympy_mappings,
         )
         sympy_format.append(eqn)
 
         # NumPy:
-        sympy_symbols = create_sympy_symbols(feature_names_in)
         lambda_format.append(
             sympy2numpy(
                 eqn,

pysr/export_numpy.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import warnings
+from functools import cached_property
 
 import numpy as np
 import pandas as pd
@@ -29,6 +30,13 @@ def __init__(self, eqn, sympy_symbols, selection=None):
     def __repr__(self):
         return f"PySRFunction(X=>{self._sympy})"
 
+    def __getstate__(self):
+        # The lambdified function is generated code that cannot be pickled,
+        # so drop the cached copy; it is rebuilt on first use after loading.
+        state = self.__dict__.copy()
+        state.pop("_lambda", None)
+        return state
+
     def __call__(self, X):
         expected_shape = (X.shape[0],)
         if isinstance(X, pd.DataFrame):
@@ -49,6 +57,8 @@ def __call__(self, X):
 
         return self._lambda(*X.T) * np.ones(expected_shape)
 
-    @property
+    @cached_property
     def _lambda(self):
+        # `lambdify` generates and compiles Python source, which costs far
+        # more than evaluating the result, so build it once per equation.
         return lambdify(self._sympy_symbols, self._sympy)

pysr/expression_specs.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import copy
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from textwrap import dedent
@@ -279,7 +278,6 @@ def _search_output_to_callable_expressions(
     search_output,
     i: int | None,
 ) -> pd.DataFrame:
-    equations = copy.deepcopy(equations)
     _, all_out_hof = search_output
     out_hof = all_out_hof[i] if i is not None else all_out_hof
     expressions = []

pysr/sr.py

@@ -2468,7 +2468,7 @@ def _run(
         np_dtype = (
             None
             if type_spec_runtime is not None
-            else self._get_precision_mapped_dtype(np.array(X))
+            else self._get_precision_mapped_dtype(np.asarray(X))
         )
 
         if self.elementwise_loss is not None:
@@ -2693,19 +2693,20 @@ def _run(
 
         # Convert data to desired precision
 
-        # This converts the data into a Julia array:
+        # This converts the data into a Julia array. `jl_numpy_array` copies
+        # into Julia memory, so only convert dtypes here rather than copying:
         if type_spec_runtime is not None:
             jl_X = type_spec_to_julia_array(type_spec_runtime, X, transpose=True)
             jl_y = type_spec_to_julia_array(type_spec_runtime, y)
         else:
-            jl_X = jl_numpy_array(np.array(X, dtype=np_dtype).T)
-            numeric_y = np.array(y, dtype=np_dtype)
+            jl_X = jl_numpy_array(np.asarray(X, dtype=np_dtype).T)
+            numeric_y = np.asarray(y, dtype=np_dtype)
             jl_y = jl_numpy_array(numeric_y.T if numeric_y.ndim > 1 else numeric_y)
         if weights is not None:
             if len(weights.shape) == 1:
-                jl_weights = jl_numpy_array(np.array(weights, dtype=np_dtype))
+                jl_weights = jl_numpy_array(np.asarray(weights, dtype=np_dtype))
             else:
-                jl_weights = jl_numpy_array(np.array(weights, dtype=np_dtype).T)
+                jl_weights = jl_numpy_array(np.asarray(weights, dtype=np_dtype).T)
         else:
             jl_weights = None
 
@@ -3134,7 +3135,6 @@ def latex(
             raise ValueError(
                 f"`expression_spec={self.expression_spec_}` does not support latex export."
             )
-        self.refresh()
         sympy_representation = self.sympy(index=index)
         if self.nout_ > 1:
             output = []
@@ -3411,8 +3411,7 @@ def idx_model_selection(equations: pd.DataFrame, model_selection: str):
         chosen_idx = equations["loss"].idxmin()
     elif model_selection == "best":
         threshold = 1.5 * equations["loss"].min()
-        filtered_equations = equations.query(f"loss <= {threshold}")
-        chosen_idx = filtered_equations["score"].idxmax()
+        chosen_idx = equations.loc[equations["loss"] <= threshold, "score"].idxmax()
     elif model_selection == "score":
         chosen_idx = equations["score"].idxmax()
     else:

pysr/test/test_main.py

@@ -1592,6 +1592,24 @@ def test_best_lambda(self):
         for f in [self.model.predict, self.equations_.iloc[-1]["lambda_format"]]:
             np.testing.assert_almost_equal(f(X), y, decimal=3)
 
+    def test_lambda_format_compiles_once_and_stays_picklable(self):
+        f = self.equations_.iloc[-1]["lambda_format"]
+        with mock.patch("pysr.export_numpy.lambdify", wraps=sympy.lambdify) as lambdify:
+            for _ in range(3):
+                f(self.X)
+                self.model.predict(self.X)
+        self.assertEqual(lambdify.call_count, 1)
+
+        restored = pkl.loads(pkl.dumps(f))
+        np.testing.assert_array_equal(restored(self.X), f(self.X))
+
+    def test_latex_refreshes_once(self):
+        with mock.patch.object(
+            self.model, "refresh", wraps=self.model.refresh
+        ) as refresh:
+            self.model.latex()
+        self.assertEqual(refresh.call_count, 1)
+
     def test_all_selection_strategies(self):
         equations = pd.DataFrame(
             dict(