MilesCranmer commited on
Commit
a812e16
1 Parent(s): 0ddc60f

Use literal for optimizer_algorithm type hint as well

Browse files
Files changed (1) hide show
  1. pysr/sr.py +2 -2
pysr/sr.py CHANGED
@@ -171,7 +171,7 @@ def _check_assertions(
171
 
172
 
173
  # Class validation constants
174
- VALID_OPTIMIZER_ALGORITHMS = ["NelderMead", "BFGS"]
175
 
176
 
177
  class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
@@ -705,7 +705,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
705
  topn: int = 12,
706
  should_simplify: Optional[bool] = None,
707
  should_optimize_constants: bool = True,
708
- optimizer_algorithm: str = "BFGS",
709
  optimizer_nrestarts: int = 2,
710
  optimize_probability: float = 0.14,
711
  optimizer_iterations: int = 8,
 
171
 
172
 
173
  # Class validation constants
174
+ VALID_OPTIMIZER_ALGORITHMS = ["BFGS", "NelderMead"]
175
 
176
 
177
  class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
 
705
  topn: int = 12,
706
  should_simplify: Optional[bool] = None,
707
  should_optimize_constants: bool = True,
708
+ optimizer_algorithm: Literal["BFGS", "NelderMead"] = "BFGS",
709
  optimizer_nrestarts: int = 2,
710
  optimize_probability: float = 0.14,
711
  optimizer_iterations: int = 8,