Merge pull request #910 from MilesCranmer/fix-inv-pickling

fix: pickling of inv

Author: MilesCranmer

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pysr"
-version = "1.5.5"
+version = "1.5.6"
 authors = [
     {name = "Miles Cranmer", email = "[email protected]"},
 ]

pysr/export_sympy.py

@@ -12,6 +12,7 @@
 
 sympy_mappings = {
     "div": lambda x, y: x / y,
+    "inv": lambda x: 1 / x,
     "mult": lambda x, y: x * y,
     "sqrt": lambda x: sympy.sqrt(x),
     "sqrt_abs": lambda x: sympy.sqrt(abs(x)),

pysr/test/test_main.py

@@ -31,6 +31,7 @@
     load_all_packages,
 )
 from pysr.export_latex import sympy2latex
+from pysr.export_sympy import pysr2sympy
 from pysr.feature_selection import _handle_feature_selection, run_feature_selection
 from pysr.julia_helpers import init_julia
 from pysr.sr import (
@@ -903,6 +904,28 @@ def test_feature_selection_handler(self):
 class TestMiscellaneous(unittest.TestCase):
     """Test miscellaneous functions."""
 
+    def test_pickle_inv_sympy_expression(self):
+        """Test that sympy expressions with the inv operator can be pickled and unpickled correctly."""
+        expr_str = "inv(x0) + x1"
+        sympy_expr = pysr2sympy(expr_str, feature_names_in=["x0", "x1"])
+
+        # Evaluate the original expression at a test point
+        test_vals = {sympy.Symbol("x0"): 2.0, sympy.Symbol("x1"): 3.0}
+        original_result = float(sympy_expr.subs(test_vals))
+
+        # Pickle and unpickle the sympy expression
+        serialized = pkl.dumps(sympy_expr)
+        deserialized_expr = pkl.loads(serialized)
+
+        # Evaluate the unpickled expression at the same test point
+        unpickled_result = float(deserialized_expr.subs(test_vals))
+
+        # Verify the results match
+        self.assertEqual(original_result, unpickled_result)
+
+        # Check that the same operator mapping was used (1/x for inv)
+        self.assertEqual(original_result, 0.5 + 3.0)  # 1/2 + 3 = 3.5
+
     def test_pickle_with_temp_equation_file(self):
         """If we have a temporary equation file, unpickle the estimator."""
         model = PySRRegressor(