MilesCranmer commited on
Commit
7701310
1 Parent(s): c52d151

Allow control of file loading verbosity

Browse files
Files changed (1) hide show
  1. pysr/sr.py +12 -6
pysr/sr.py CHANGED
@@ -903,6 +903,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
903
  feature_names_in=None,
904
  selection_mask=None,
905
  nout=1,
 
906
  **pysr_kwargs,
907
  ):
908
  """
@@ -932,6 +933,8 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
932
  Number of outputs of the model.
933
  Not needed if loading from a pickle file.
934
  Default is `1`.
 
 
935
  **pysr_kwargs : dict
936
  Any other keyword arguments to initialize the PySRRegressor object.
937
  These will overwrite those stored in the pickle file.
@@ -946,9 +949,11 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
946
  pkl_filename = _csv_filename_to_pkl_filename(equation_file)
947
 
948
  # Try to load model from <equation_file>.pkl
949
- print(f"Checking if {pkl_filename} exists...")
 
950
  if os.path.exists(pkl_filename):
951
- print(f"Loading model from {pkl_filename}")
 
952
  assert binary_operators is None
953
  assert unary_operators is None
954
  assert n_features_in is None
@@ -968,10 +973,11 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
968
  return model
969
 
970
  # Else, we re-create it.
971
- print(
972
- f"{pkl_filename} does not exist, "
973
- "so we must create the model from scratch."
974
- )
 
975
  assert binary_operators is not None or unary_operators is not None
976
  assert n_features_in is not None
977
 
 
903
  feature_names_in=None,
904
  selection_mask=None,
905
  nout=1,
906
+ verbosity=1,
907
  **pysr_kwargs,
908
  ):
909
  """
 
933
  Number of outputs of the model.
934
  Not needed if loading from a pickle file.
935
  Default is `1`.
936
+ verbosity : int
937
+ What verbosity level to use. 0 means minimal print statements.
938
  **pysr_kwargs : dict
939
  Any other keyword arguments to initialize the PySRRegressor object.
940
  These will overwrite those stored in the pickle file.
 
949
  pkl_filename = _csv_filename_to_pkl_filename(equation_file)
950
 
951
  # Try to load model from <equation_file>.pkl
952
+ if verbosity > 0:
953
+ print(f"Checking if {pkl_filename} exists...")
954
  if os.path.exists(pkl_filename):
955
+ if verbosity > 0:
956
+ print(f"Loading model from {pkl_filename}")
957
  assert binary_operators is None
958
  assert unary_operators is None
959
  assert n_features_in is None
 
973
  return model
974
 
975
  # Else, we re-create it.
976
+ if verbosity > 0:
977
+ print(
978
+ f"{pkl_filename} does not exist, "
979
+ "so we must create the model from scratch."
980
+ )
981
  assert binary_operators is not None or unary_operators is not None
982
  assert n_features_in is not None
983