Add early_stop_condition to stop earlier

pysr/sr.py

@@ -420,6 +420,7 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
         cluster_manager=None,
         skip_mutation_failures=True,
         max_evals=None,
+        early_stop_condition=None,
         # To support deprecated kwargs:
         **kwargs,
     ):
@@ -562,6 +563,8 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
         :type skip_mutation_failures: bool
         :param max_evals: Limits the total number of evaluations of expressions to this number.
         :type max_evals: int
+        :param early_stop_condition: Stop the search early if this loss is reached.
+        :type early_stop_condition: float
         :param kwargs: Supports deprecated keyword arguments. Other arguments will result
         in an error
         :type kwargs: dict
@@ -749,6 +752,7 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
                 cluster_manager=cluster_manager,
                 skip_mutation_failures=skip_mutation_failures,
                 max_evals=max_evals,
+                early_stop_condition=early_stop_condition,
             ),
         }
 
@@ -1313,8 +1317,9 @@ class PySRRegressor(BaseEstimator, RegressorMixin):
             progress=self.params["progress"],
             timeout_in_seconds=self.params["timeout_in_seconds"],
             crossoverProbability=self.params["crossover_probability"],
-            max_evals=self.params["max_evals"],
             skip_mutation_failures=self.params["skip_mutation_failures"],
+            max_evals=self.params["max_evals"],
+            earlyStopCondition=self.params["early_stop_condition"],
         )
 
         np_dtype = {16: np.float16, 32: np.float32, 64: np.float64}[