MilesCranmer commited on
Commit
14dfd82
1 Parent(s): 7f2c133

Test multi output

Browse files
Files changed (1) hide show
  1. test/test.py +4 -3
test/test.py CHANGED
@@ -17,14 +17,15 @@ equations = pysr(X, y, **default_test_kwargs)
17
  print(equations)
18
  assert equations.iloc[-1]['MSE'] < 1e-4
19
 
20
- print("Test 2 - test custom operator")
21
- y = X[:, 0]**2
22
  equations = pysr(X, y,
23
  unary_operators=["sq(x) = x^2"], binary_operators=["plus"],
24
  extra_sympy_mappings={'square': lambda x: x**2},
25
  **default_test_kwargs)
26
  print(equations)
27
- assert equations.iloc[-1]['MSE'] < 1e-4
 
28
 
29
  X = np.random.randn(100, 1)
30
  y = X[:, 0] + 3.0
 
17
  print(equations)
18
  assert equations.iloc[-1]['MSE'] < 1e-4
19
 
20
+ print("Test 2 - test custom operator, and multiple outputs")
21
+ y = X[:, [0, 1]]**2
22
  equations = pysr(X, y,
23
  unary_operators=["sq(x) = x^2"], binary_operators=["plus"],
24
  extra_sympy_mappings={'square': lambda x: x**2},
25
  **default_test_kwargs)
26
  print(equations)
27
+ assert equations[0].iloc[-1]['MSE'] < 1e-4
28
+ assert equations[1].iloc[-1]['MSE'] < 1e-4
29
 
30
  X = np.random.randn(100, 1)
31
  y = X[:, 0] + 3.0