MilesCranmer commited on
Commit
b41237e
1 Parent(s): 70dedf9

Fix variable names when using select_k_features

Browse files
Files changed (1) hide show
  1. pysr/sr.py +1 -1
pysr/sr.py CHANGED
@@ -215,7 +215,7 @@ def pysr(X=None, y=None, weights=None,
215
  X = X[:, selection]
216
 
217
  if use_custom_variable_names:
218
- variable_names = variable_names[selection]
219
 
220
  if populations is None:
221
  populations = procs
 
215
  X = X[:, selection]
216
 
217
  if use_custom_variable_names:
218
+ variable_names = [variable_names[selection[i]] for i in range(len(selection))]
219
 
220
  if populations is None:
221
  populations = procs