MilesCranmer commited on
Commit
0e86456
1 Parent(s): 4f7e6cf

Add custom operators to unit test

Browse files
Files changed (1) hide show
  1. pysr/test/test.py +12 -1
pysr/test/test.py CHANGED
@@ -922,10 +922,21 @@ class TestDimensionalConstraints(unittest.TestCase):
922
  def test_dimensional_constraints(self):
923
  y = np.cos(self.X[:, 0])
924
  model = PySRRegressor(
925
- unary_operators=["cos"],
 
 
 
 
 
926
  **self.default_test_kwargs,
927
  early_stop_condition=1e-8,
928
  select_k_features=3,
 
 
 
 
 
 
929
  )
930
  model.fit(self.X, y, X_units=["m", "m", "m", "m", "m"], y_units="m")
931
 
 
922
  def test_dimensional_constraints(self):
923
  y = np.cos(self.X[:, 0])
924
  model = PySRRegressor(
925
+ binary_operators=[
926
+ "my_add(x, y) = x + y",
927
+ "my_sub(x, y) = x - y",
928
+ "my_mul(x, y) = x * y",
929
+ ],
930
+ unary_operators=["my_cos(x) = cos(x)"],
931
  **self.default_test_kwargs,
932
  early_stop_condition=1e-8,
933
  select_k_features=3,
934
+ extra_sympy_mappings={
935
+ "my_cos": sympy.cos,
936
+ "my_add": lambda x, y: x + y,
937
+ "my_sub": lambda x, y: x - y,
938
+ "my_mul": lambda x, y: x * y,
939
+ },
940
  )
941
  model.fit(self.X, y, X_units=["m", "m", "m", "m", "m"], y_units="m")
942