Skip to content

BUG: rolling.corr with MultiIndex columns #43261

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Sep 7, 2021
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v1.4.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -356,7 +356,7 @@ Groupby/resample/rolling
- Bug in :meth:`pandas.DataFrame.ewm`, where non-float64 dtypes were silently failing (:issue:`42452`)
- Bug in :meth:`pandas.DataFrame.rolling` operation along rows (``axis=1``) incorrectly omits columns containing ``float16`` and ``float32`` (:issue:`41779`)
- Bug in :meth:`Resampler.aggregate` did not allow the use of Named Aggregation (:issue:`32803`)
-
- Bug in :meth:`DataFrame.rolling.corr` when the :class:`DataFrame` columns was a :class:`MultiIndex` (:issue:`21157`)

Reshaping
^^^^^^^^^
Expand Down
14 changes: 12 additions & 2 deletions pandas/core/window/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,18 @@ def dataframe_from_int_dict(data, frame_template):
# mypy needs to know columns is a MultiIndex, Index doesn't
# have levels attribute
arg2.columns = cast(MultiIndex, arg2.columns)
result.index = MultiIndex.from_product(
arg2.columns.levels + [result_index]
# GH 21157
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thi s is really confusing. you are matching the codes from the original? can you reindex?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll add comment to what's going on here, but unfortunately can't reindex here.

This needs to act like a MultiIndex.product, but instead of taking the product of all original levels + result level independently, the original levels need too be treated like a group first before taking a product

idx_codes, idx_uniques = result_index.factorize()
result_levels = list(arg2.columns.levels) + [idx_uniques]
result_codes = [
np.tile(code, int(len(result) / len(code)))
for code in arg2.columns.codes
] + [np.tile(idx_codes, int(len(result) / len(idx_codes)))]
result_names = list(arg2.columns.names) + [result_index.name]
result.index = MultiIndex(
levels=result_levels,
codes=result_codes,
names=result_names,
)
# GH 34440
num_levels = len(result.index.levels)
Expand Down
16 changes: 16 additions & 0 deletions pandas/tests/window/test_rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -1500,3 +1500,19 @@ def test_rolling_numeric_dtypes():
dtype="float64",
)
tm.assert_frame_equal(result, expected)


def test_multindex_columns_pairwise_func():
# GH 21157
columns = MultiIndex.from_arrays([["M", "N"], ["P", "Q"]], names=["a", "b"])
df = DataFrame(np.ones((5, 2)), columns=columns)
result = df.rolling(3).corr()
expected = DataFrame(
np.nan,
index=MultiIndex.from_arrays(
[np.repeat(np.arange(5), 2), ["M", "N"] * 5, ["P", "Q"] * 5],
names=[None, "a", "b"],
),
columns=columns,
)
tm.assert_frame_equal(result, expected)