fix(corr_interactive_plot): add MultiIndex support

akanz1 · Dec 21, 2023 · 90703c8 · 90703c8
1 parent 6328744
commit 90703c8
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/src/klib/describe.py b/src/klib/describe.py
@@ -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,
     }