MilesCranmer commited on
Commit
41f0271
1 Parent(s): 06a77ea

test: code coverage for complexity_of_variables

Browse files
Files changed (1) hide show
  1. pysr/test/test.py +11 -0
pysr/test/test.py CHANGED
@@ -199,6 +199,17 @@ class TestPipeline(unittest.TestCase):
199
  self.assertEqual(model.get_best()[0]["complexity"], 2)
200
  self.assertEqual(model.get_best()[1]["complexity"], 3 if case == 1 else 2)
201
 
 
 
 
 
 
 
 
 
 
 
 
202
  def test_multioutput_weighted_with_callable_temp_equation(self):
203
  X = self.X.copy()
204
  y = X[:, [0, 1]] ** 2
 
199
  self.assertEqual(model.get_best()[0]["complexity"], 2)
200
  self.assertEqual(model.get_best()[1]["complexity"], 3 if case == 1 else 2)
201
 
202
+ def test_error_message_custom_variable_complexity(self):
203
+ X = np.ones((10, 2))
204
+ y = np.ones((10,))
205
+ model = PySRRegressor()
206
+ with self.assertRaises(ValueError) as cm:
207
+ model.fit(X, y, complexity_of_variables=[1, 2, 3])
208
+
209
+ self.assertIn(
210
+ "number of elements in `complexity_of_variables`", str(cm.exception)
211
+ )
212
+
213
  def test_multioutput_weighted_with_callable_temp_equation(self):
214
  X = self.X.copy()
215
  y = X[:, [0, 1]] ** 2