MilesCranmer commited on
Commit
5fac847
1 Parent(s): 2e2f15a

Fix variable_names in test

Browse files
Files changed (2) hide show
  1. pysr/sr.py +5 -1
  2. test/test.py +1 -2
pysr/sr.py CHANGED
@@ -541,7 +541,11 @@ Tried to activate project {julia_project} but failed."""
541
  Main.y,
542
  weights=Main.weights,
543
  niterations=int(niterations),
544
- varMap=([variable_names[i] for i in selection] if selection else None),
 
 
 
 
545
  options=options,
546
  numprocs=int(cprocs),
547
  multithreading=bool(multithreading),
 
541
  Main.y,
542
  weights=Main.weights,
543
  niterations=int(niterations),
544
+ varMap=(
545
+ variable_names
546
+ if selection is None
547
+ else [variable_names[i] for i in selection]
548
+ ),
549
  options=options,
550
  numprocs=int(cprocs),
551
  multithreading=bool(multithreading),
test/test.py CHANGED
@@ -209,10 +209,9 @@ class TestFeatureSelection(unittest.TestCase):
209
  X = np.random.randn(20000, 5)
210
  y = X[:, 2] ** 2 + X[:, 3] ** 2
211
  var_names = [f"x{i}" for i in range(5)]
212
- selected_X, selected_var_names, selection = _handle_feature_selection(
213
  X,
214
  select_k_features=2,
215
- use_custom_variable_names=True,
216
  variable_names=var_names,
217
  y=y,
218
  )
 
209
  X = np.random.randn(20000, 5)
210
  y = X[:, 2] ** 2 + X[:, 3] ** 2
211
  var_names = [f"x{i}" for i in range(5)]
212
+ selected_X, selection = _handle_feature_selection(
213
  X,
214
  select_k_features=2,
 
215
  variable_names=var_names,
216
  y=y,
217
  )