MilesCranmer commited on
Commit
abd0cfa
1 Parent(s): 4ca54a5

Set `y_variable_names` in output

Browse files
Files changed (1) hide show
  1. pysr/sr.py +5 -0
pysr/sr.py CHANGED
@@ -1782,6 +1782,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
1782
  None if parallelism in ["serial", "multithreading"] else int(self.procs)
1783
  )
1784
 
 
 
 
 
1785
  # Call to Julia backend.
1786
  # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl
1787
  self.raw_julia_state_ = SymbolicRegression.equation_search(
@@ -1795,6 +1799,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
1795
  and self.pretty_feature_names_in_ is not None
1796
  else self.feature_names_in_.tolist()
1797
  ),
 
1798
  X_units=self.X_units_,
1799
  y_units=self.y_units_,
1800
  options=options,
 
1782
  None if parallelism in ["serial", "multithreading"] else int(self.procs)
1783
  )
1784
 
1785
+ y_variable_names = None
1786
+ if len(y.shape) > 1:
1787
+ y_variable_names = [f"y{_subscriptify(i)}" for i in range(y.shape[1])]
1788
+
1789
  # Call to Julia backend.
1790
  # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl
1791
  self.raw_julia_state_ = SymbolicRegression.equation_search(
 
1799
  and self.pretty_feature_names_in_ is not None
1800
  else self.feature_names_in_.tolist()
1801
  ),
1802
+ y_variable_names=y_variable_names,
1803
  X_units=self.X_units_,
1804
  y_units=self.y_units_,
1805
  options=options,