Merge pull request #845 from MilesCranmer/fix-greater

fix: comparison operator parsing

Author: MilesCranmer

docs/operators.md

@@ -16,9 +16,12 @@ it can exponentially increase the search space.
 |--------------|------------|----------|
 | `+`          | `max`      | `logical_or`[^2] |
 | `-`          | `min`      | `logical_and`[^3]|
-| `*`          | `greater`[^4] |                 |
-| `/`          | `cond`[^5]   |                 |
-| `^`          | `mod`      |                 |
+| `*`          | `>`[^4]    |                 |
+| `/`          | `>=`       |                 |
+| `^`          | `<`        |                 |
+|              | `<=`       |                 |
+|              | `cond`[^5] |                 |
+|              | `mod`      |                 |
 
 **Unary Operators**
 
@@ -74,5 +77,5 @@ any invalid values over the training dataset.
 [^1]: However, you will need to define a sympy equivalent in `extra_sympy_mapping` if you want to use a function not in the above list.
 [^2]: `logical_or` is equivalent to `(x, y) -> (x > 0 || y > 0) ? 1 : 0`
 [^3]: `logical_and` is equivalent to `(x, y) -> (x > 0 && y > 0) ? 1 : 0`
-[^4]: `greater` is equivalent to `(x, y) -> x > y ? 1 : 0`
+[^4]: `>` is equivalent to `(x, y) -> x > y ? 1 : 0`
 [^5]: `cond` is equivalent to `(x, y) -> x > 0 ? y : 0`

pyproject.toml

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

pysr/export_sympy.py

@@ -55,6 +55,9 @@
     "max": lambda x, y: sympy.Piecewise((y, x < y), (x, True)),
     "min": lambda x, y: sympy.Piecewise((x, x < y), (y, True)),
     "greater": lambda x, y: sympy.Piecewise((1.0, x > y), (0.0, True)),
+    "less": lambda x, y: sympy.Piecewise((1.0, x < y), (0.0, True)),
+    "greater_equal": lambda x, y: sympy.Piecewise((1.0, x >= y), (0.0, True)),
+    "less_equal": lambda x, y: sympy.Piecewise((1.0, x <= y), (0.0, True)),
     "cond": lambda x, y: sympy.Piecewise((y, x > 0), (0.0, True)),
     "logical_or": lambda x, y: sympy.Piecewise((1.0, (x > 0) | (y > 0)), (0.0, True)),
     "logical_and": lambda x, y: sympy.Piecewise((1.0, (x > 0) & (y > 0)), (0.0, True)),

pysr/julia_import.py

@@ -67,5 +67,8 @@ def _import_juliacall():
 # Expose `D` operator:
 jl.seval("using SymbolicRegression: D")
 
+# Expose other operators:
+jl.seval("using SymbolicRegression: less, greater_equal, less_equal")
+
 jl.seval("using Pkg: Pkg")
 Pkg = jl.Pkg

pysr/juliapkg.json

@@ -3,7 +3,7 @@
     "packages": {
         "SymbolicRegression": {
             "uuid": "8254be44-1295-4e6a-a16d-46603ac705cb",
-            "version": "~1.8.0"
+            "version": "~1.9.0"
         },
         "Serialization": {
             "uuid": "9e88b42a-f829-5b0c-bbe9-9e923198166b",

pysr/test/test_main.py

@@ -758,6 +758,25 @@ def test_tensorboard_logger(self):
                 # Verify model still works as expected
                 self.assertLessEqual(model.get_best()["loss"], 1e-4)
 
+    def test_comparison_operator(self):
+        X = self.rstate.randn(100, 2)
+        y = ((X[:, 0] + X[:, 1]) < (X[:, 0] * X[:, 1])).astype(float)
+
+        model = PySRRegressor(
+            binary_operators=["<", "+", "*"],
+            **self.default_test_kwargs,
+            early_stop_condition="stop_if(loss, complexity) = loss < 1e-4 && complexity <= 7",
+        )
+
+        model.fit(X, y)
+
+        best_equation = model.get_best()["equation"]
+        self.assertIn("less", best_equation)
+        self.assertLessEqual(model.get_best()["loss"], 1e-4)
+
+        y_pred = model.predict(X)
+        np.testing.assert_array_almost_equal(y, y_pred, decimal=3)
+
 
 def manually_create_model(equations, feature_names=None):
     if feature_names is None: