MilesCranmer commited on
Commit
f1ac704
1 Parent(s): 1f01976

Allow both `bkup` and `csv` file

Browse files
Files changed (1) hide show
  1. pysr/sr.py +8 -5
pysr/sr.py CHANGED
@@ -1834,10 +1834,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
1834
  if self.nout_ > 1:
1835
  all_outputs = []
1836
  for i in range(1, self.nout_ + 1):
1837
- df = pd.read_csv(
1838
- str(self.equation_file_) + f".out{i}" + ".bkup",
1839
- sep="|",
1840
- )
1841
  # Rename Complexity column to complexity:
1842
  df.rename(
1843
  columns={
@@ -1850,7 +1850,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
1850
 
1851
  all_outputs.append(df)
1852
  else:
1853
- all_outputs = [pd.read_csv(str(self.equation_file_) + ".bkup", sep="|")]
 
 
 
1854
  all_outputs[-1].rename(
1855
  columns={
1856
  "Complexity": "complexity",
 
1834
  if self.nout_ > 1:
1835
  all_outputs = []
1836
  for i in range(1, self.nout_ + 1):
1837
+ cur_filename = str(self.equation_file_) + f".out{i}" + ".bkup"
1838
+ if not os.path.exists(cur_filename):
1839
+ cur_filename = str(self.equation_file_) + f".out{i}"
1840
+ df = pd.read_csv(cur_filename, sep="|")
1841
  # Rename Complexity column to complexity:
1842
  df.rename(
1843
  columns={
 
1850
 
1851
  all_outputs.append(df)
1852
  else:
1853
+ filename = str(self.equation_file_) + ".bkup"
1854
+ if not os.path.exists(filename):
1855
+ filename = str(self.equation_file_)
1856
+ all_outputs = [pd.read_csv(filename, sep="|")]
1857
  all_outputs[-1].rename(
1858
  columns={
1859
  "Complexity": "complexity",