diff --git a/qlib/contrib/eva/alpha.py b/qlib/contrib/eva/alpha.py index 2c2f9ea7e18..c7cdaf780b0 100644 --- a/qlib/contrib/eva/alpha.py +++ b/qlib/contrib/eva/alpha.py @@ -130,8 +130,8 @@ def pred_autocorr(pred: pd.Series, lag=1, inst_col="instrument", date_col="datet :param lag: """ if isinstance(pred, pd.DataFrame): - pred = pred.iloc[:, 0] get_module_logger("pred_autocorr").warning(f"Only the first column in {pred.columns} of `pred` is kept") + pred = pred.iloc[:, 0] pred_ustk = pred.sort_index().unstack(inst_col) corr_s = {} for (idx, cur), (_, prev) in zip(pred_ustk.iterrows(), pred_ustk.shift(lag).iterrows()): diff --git a/tests/misc/test_pred_autocorr_dataframe.py b/tests/misc/test_pred_autocorr_dataframe.py new file mode 100644 index 00000000000..67f2c48c4e2 --- /dev/null +++ b/tests/misc/test_pred_autocorr_dataframe.py @@ -0,0 +1,24 @@ +import logging + +import pandas as pd + +from qlib.contrib.eva.alpha import pred_autocorr + + +def test_pred_autocorr_dataframe_uses_first_column(caplog, monkeypatch): + index = pd.MultiIndex.from_product( + [list("ABC"), pd.date_range("2024-01-02", periods=3, freq="B")], + names=["instrument", "datetime"], + ) + first = pd.Series([1, 2, 3, 2, 3, 1, 3, 1, 2], index=index, name="score") + # The second column must have different autocorrelations to detect selecting it. + second = pd.Series([1, 2, 4, 1, 5, 2, 3, 1, 9], index=index) + predictions = pd.DataFrame({"score": first, "unused": second}) + + monkeypatch.setattr(logging.getLogger("qlib"), "propagate", True) + pd.testing.assert_series_equal(pred_autocorr(predictions), pred_autocorr(first)) + assert ( + "qlib.pred_autocorr", + logging.WARNING, + f"Only the first column in {predictions.columns} of `pred` is kept", + ) in caplog.record_tuples