[FIX] Add MultiIndex support (#175)

* fix(corr_interactive_plot): add MultiIndex support * Update describe.py ignore complexity warnings * Update describe.py * Update describe.py * Update describe.py --------- Co-authored-by: Andreas Kanz <[email protected]>
akanz1 · Jan 5, 2024 · feed65f · feed65f
1 parent 05e0bc2
commit feed65f
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/src/klib/describe.py b/src/klib/describe.py
@@ -428,7 +428,7 @@ def corr_plot(
     return ax
 
 
-def corr_interactive_plot(
+def corr_interactive_plot(  # noqa: C901
     data: pd.DataFrame,
     split: Literal["pos", "neg", "high", "low"] | None = None,
     threshold: float = 0.0,
@@ -591,6 +591,15 @@ def corr_interactive_plot(
 
     vtext = corr.round(2).fillna("") if annot else None
 
+    corr_columns = corr.columns
+    corr_index = corr.index
+
+    if isinstance(corr_columns, pd.MultiIndex):
+        corr_columns = ["-".join(col) for col in corr.columns]
+
+    if isinstance(corr_index, pd.MultiIndex):
+        corr_index = ["-".join(idx) for idx in corr.index]
+
     # Specify kwargs for the heatmap
     kwargs = {
         "colorscale": cmap,
@@ -599,8 +608,8 @@ def corr_interactive_plot(
         "text": vtext,
         "texttemplate": "%{text}",
         "textfont": {"size": 12},
-        "x": corr.columns,
-        "y": corr.index,
+        "x": corr_columns,
+        "y": corr_index,
         "z": corr,
         **kwargs,
     }