test: simplify custom complexity test

MilesCranmer · Jun 16, 2024 · be4be4d · be4be4d
1 parent 09617a6
commit be4be4d
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/pysr/test/test.py b/pysr/test/test.py
@@ -181,18 +181,18 @@ def test_custom_variable_complexity(self):
         for case in (1, 2):
             y = self.X[:, [0, 1]] ** 2
             model = PySRRegressor(
-                binary_operators=["*", "+"],
+                binary_operators=["*"],
                 verbosity=0,
                 **self.default_test_kwargs,
                 early_stop_condition=f"stop_if(l, c) = l < 1e-8 && c <= {7 if case == 1 else 5}",
             )
             if case == 1:
-                complexity_of_variables = [2, 3] + [
-                    100 for _ in range(self.X.shape[1] - 2)
-                ]
+                complexity_of_variables = [2, 3]
             elif case == 2:
                 complexity_of_variables = 2
-            model.fit(self.X, y, complexity_of_variables=complexity_of_variables)
+            model.fit(
+                self.X[:, [0, 1]], y, complexity_of_variables=complexity_of_variables
+            )
             equations = model.equations_
             self.assertLessEqual(equations[0].iloc[-1]["loss"], 1e-4)
             self.assertLessEqual(equations[1].iloc[-1]["loss"], 1e-4)