MilesCranmer commited on
Commit
58834e8
1 Parent(s): fcf3f92

Add PySRRegressor test

Browse files
Files changed (1) hide show
  1. test/test.py +7 -6
test/test.py CHANGED
@@ -1,7 +1,7 @@
1
  import unittest
2
  from unittest.mock import patch
3
  import numpy as np
4
- from pysr import pysr, get_hof, best, best_tex, best_callable, best_row
5
  from pysr.sr import run_feature_selection, _handle_feature_selection, _yesno
6
  import sympy
7
  from sympy import lambdify
@@ -78,18 +78,19 @@ class TestPipeline(unittest.TestCase):
78
  best_callable()[1](self.X), self.X[:, 1] ** 2, decimal=4
79
  )
80
 
81
- def test_empty_operators_single_input(self):
82
  X = np.random.randn(100, 1)
83
  y = X[:, 0] + 3.0
84
- equations = pysr(
85
- X,
86
- y,
87
  unary_operators=[],
88
  binary_operators=["plus"],
89
  **self.default_test_kwargs,
90
  )
 
91
 
92
- self.assertLessEqual(equations.iloc[-1]["MSE"], 1e-4)
 
93
 
94
  def test_noisy(self):
95
 
 
1
  import unittest
2
  from unittest.mock import patch
3
  import numpy as np
4
+ from pysr import pysr, get_hof, best, best_tex, best_callable, best_row, PySRRegressor
5
  from pysr.sr import run_feature_selection, _handle_feature_selection, _yesno
6
  import sympy
7
  from sympy import lambdify
 
78
  best_callable()[1](self.X), self.X[:, 1] ** 2, decimal=4
79
  )
80
 
81
+ def test_empty_operators_single_input_sklearn(self):
82
  X = np.random.randn(100, 1)
83
  y = X[:, 0] + 3.0
84
+ regressor = PySRRegressor(
85
+ model_selection="accuracy",
 
86
  unary_operators=[],
87
  binary_operators=["plus"],
88
  **self.default_test_kwargs,
89
  )
90
+ regressor.fit(X, y)
91
 
92
+ self.assertLessEqual(regressor.equations.iloc[-1]["MSE"], 1e-4)
93
+ np.testing_assert_almost_equal(regressor.predict(X), y, decimal=1)
94
 
95
  def test_noisy(self):
96