From a1382fe2f2ebac7a262063eb1e9d2a7389aacd0f Mon Sep 17 00:00:00 2001 From: Mark Colley <39366189+M-Colley@users.noreply.github.com> Date: Wed, 7 Oct 2026 22:04:58 +0100 Subject: [PATCH] perf: cache lambdified equations and avoid redundant export work `CallableEquation._lambda` was a property that re-ran `sympy.lambdify` (code generation plus `exec`) on every call, so every `predict()` and every direct `lambda_format(X)` call paid ~1-3 ms of compilation before evaluating the expression. Cache it with `functools.cached_property` and drop the cached function in `__getstate__`, since generated code cannot be pickled; it is rebuilt lazily after loading. Also remove smaller redundancies on the same paths: - `latex()` refreshed the hall of fame and then called `sympy()`, which refreshes it again. - `idx_model_selection` filtered with `DataFrame.query` on a formatted float on every `predict()`; a boolean mask is ~3x cheaper and does not round-trip the threshold through a string. - `add_export_formats` and `_search_output_to_callable_expressions` deep-copied frames they only read, and the former rebuilt the SymPy symbol list for every row. - `_run` copied `X` once just to inspect its dtype and again before `jl_numpy_array`, which already copies into Julia memory; use `np.asarray` so a copy is only made when the dtype changes. Co-Authored-By: Claude Opus 5.5 --- pysr/export.py | 10 ++++------ pysr/export_numpy.py | 12 +++++++++++- pysr/expression_specs.py | 2 -- pysr/sr.py | 17 ++++++++--------- pysr/test/test_main.py | 18 ++++++++++++++++++ 5 files changed, 41 insertions(+), 18 deletions(-) diff --git a/pysr/export.py b/pysr/export.py index 1e4fcb5e1..8ec5100fc 100644 --- a/pysr/export.py +++ b/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, diff --git a/pysr/export_numpy.py b/pysr/export_numpy.py index c67834985..ff1506891 100644 --- a/pysr/export_numpy.py +++ b/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) diff --git a/pysr/expression_specs.py b/pysr/expression_specs.py index 187516d5d..5b77d3f09 100644 --- a/pysr/expression_specs.py +++ b/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 = [] diff --git a/pysr/sr.py b/pysr/sr.py index 961cb5225..ee11ae192 100644 --- a/pysr/sr.py +++ b/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: diff --git a/pysr/test/test_main.py b/pysr/test/test_main.py index 453eaa50a..12772f9b4 100644 --- a/pysr/test/test_main.py +++ b/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(