MilesCranmer commited on
Commit
1b17efe
1 Parent(s): 896e650

Test custom operator with underscore

Browse files
Files changed (1) hide show
  1. test/test.py +3 -2
test/test.py CHANGED
@@ -39,8 +39,8 @@ class TestPipeline(unittest.TestCase):
39
  def test_multioutput_custom_operator_quiet(self):
40
  y = self.X[:, [0, 1]] ** 2
41
  model = PySRRegressor(
42
- unary_operators=["sq(x) = x^2"],
43
- extra_sympy_mappings={"sq": lambda x: x**2},
44
  binary_operators=["plus"],
45
  verbosity=0,
46
  **self.default_test_kwargs,
@@ -49,6 +49,7 @@ class TestPipeline(unittest.TestCase):
49
  model.fit(self.X, y)
50
  equations = model.equations
51
  print(equations)
 
52
  self.assertLessEqual(equations[0].iloc[-1]["loss"], 1e-4)
53
  self.assertLessEqual(equations[1].iloc[-1]["loss"], 1e-4)
54
 
 
39
  def test_multioutput_custom_operator_quiet(self):
40
  y = self.X[:, [0, 1]] ** 2
41
  model = PySRRegressor(
42
+ unary_operators=["square_op(x) = x^2"],
43
+ extra_sympy_mappings={"square_op": lambda x: x**2},
44
  binary_operators=["plus"],
45
  verbosity=0,
46
  **self.default_test_kwargs,
 
49
  model.fit(self.X, y)
50
  equations = model.equations
51
  print(equations)
52
+ self.assertIn("square_op", equations.sympy())
53
  self.assertLessEqual(equations[0].iloc[-1]["loss"], 1e-4)
54
  self.assertLessEqual(equations[1].iloc[-1]["loss"], 1e-4)
55