Adapt the color palette for Sankey Diagram

Author: Jacobluke-

dabest/plot_tools.py

@@ -422,6 +422,8 @@ def single_sankey(left, right, xpos=0, leftWeight=None, rightWeight=None,
 
     '''
     Make a single Sankey diagram showing proportion flow from left to right
+    Original code from: https://github.com/anazalea/pySankey
+    Changes are added to normalize each diagram's height to be 1
 
     Keywords
     --------
@@ -501,7 +503,6 @@ def single_sankey(left, right, xpos=0, leftWeight=None, rightWeight=None,
     else:
         check_data_matches_labels(leftLabels, dataFrame['right'], 'right')
 
-    #TODO: Align with the given method of setting color palette
     # If no colorDict given, make one
     if colorDict is None:
         colorDict = {}
@@ -512,7 +513,7 @@ def single_sankey(left, right, xpos=0, leftWeight=None, rightWeight=None,
     else:
         missing = [label for label in allLabels if label not in colorDict.keys()]
         if missing:
-            msg = "The colorDict parameter is missing values for the following labels : "
+            msg = "The palette parameter is missing values for the following labels : "
             msg += '{}'.format(', '.join(missing))
             raise ValueError(msg)
 
@@ -601,15 +602,15 @@ def normalize_dict(nested_dict, target):
     # Plot vertical bars for each label
     for leftLabel in leftLabels:
         ax.fill_between(
-            [leftpos + (-0.02 * xMax), leftpos],
+            [leftpos + (-0.05 * xMax), leftpos],
             2 * [leftWidths_norm[leftLabel]["bottom"]],
             2 * [leftWidths_norm[leftLabel]["bottom"] + leftWidths_norm[leftLabel]["left"]],
             color=colorDict[leftLabel],
             alpha=0.99,
         )
     for rightLabel in rightLabels:
         ax.fill_between(
-            [xMax + leftpos, leftpos + (1.02 * xMax)], 
+            [xMax + leftpos, leftpos + (1.05 * xMax)], 
             2 * [rightWidths_norm[rightLabel]['bottom']],
             2 * [rightWidths_norm[rightLabel]['bottom'] + rightWidths_norm[rightLabel]['right']],
             color=colorDict[rightLabel],
@@ -643,18 +644,47 @@ def normalize_dict(nested_dict, target):
                 )
 
 def sankeydiag(data, xvar, yvar, left_idx, right_idx, 
-                leftLabels=None, rightLabels=None, 
+                leftLabels=None, rightLabels=None,  
+                palette=None,
                 ax=None, width=0.5, rightColor=False,
                 align='center', alpha=0.65, **kwargs):
     '''
     Read in melted pd.DataFrame, and draw multiple sankey diagram on a single axes
     using the value in column yvar according to the value in column xvar
     left_idx in the column xvar is on the left side of each sankey diagram
     right_idx in the column xvar is on the right side of each sankey diagram
+
+    Keywords
+    --------
+    data: pd.DataFrame
+        input data, melted dataframe created by dabest.load()
+    left_idx: str
+        the value in column xvar that is on the left side of each sankey diagram
+    right_idx: str
+        the value in column xvar that is on the right side of each sankey diagram
+        if len(left_idx) == 1, it will be broadcasted to the same length as right_idx
+        otherwise it should have the same length as right_idx
+    leftLabels: list
+        labels for the left side of the diagram. The diagram will be sorted by these labels.
+    rightLabels: list
+        labels for the right side of the diagram. The diagram will be sorted by these labels.
+    palette: str or dict
+    ax: matplotlib axes to be drawn on
+    width: float
+        the width of each sankey diagram
+    align: str
+        the alignment of each sankey diagram, can be 'center' or 'left'
+    alpha: float
+        the transparency of each strip
+    rightColor: bool
+        if True, each strip of the diagram will be colored according to the corresponding left labels
+    colorDict: dictionary of colors for each label
+        input format: {'label': 'color'}
     '''
 
     import numpy as np
     import pandas as pd
+    import seaborn as sns
     import matplotlib.pyplot as plt
 
     if "width" in kwargs:
@@ -671,6 +701,8 @@ def sankeydiag(data, xvar, yvar, left_idx, right_idx,
 
     if ax is None:
         fig, ax = plt.subplots()
+
+    allLabels = data[yvar].unique()
 
     # Check if all the elements in left_idx and right_idx are in xvar column
     if not all(elem in data[xvar].unique() for elem in left_idx):
@@ -679,12 +711,34 @@ def sankeydiag(data, xvar, yvar, left_idx, right_idx,
         raise ValueError(f"{right_idx} not found in {xvar} column")
 
     xpos = 0
-    broadcasted_left = np.broadcast_to(left_idx, len(right_idx))
+
+    # For baseline comparison, broadcast left_idx to the same length as right_idx
+    # so that the left of sankey diagram will be the same
+    # For sequential comparison, left_idx and right_idx can have anything different 
+    # but should have the same length
+    if len(left_idx) == 1:
+        broadcasted_left = np.broadcast_to(left_idx, len(right_idx))
+    elif len(left_idx) != len(right_idx):
+        raise ValueError(f"left_idx and right_idx should have the same length")
+    else:
+        broadcasted_left = left_idx
+
+    if isinstance(palette, dict):
+        if not all(key in allLabels for key in palette.keys()):
+            raise ValueError(f"keys in palette should be in {yvar} column")
+        else: 
+            plot_palette = palette
+    elif isinstance(palette, str):
+        plot_palette = {}
+        colorPalette = sns.color_palette(palette, len(allLabels))
+        for i, label in enumerate(allLabels):
+            plot_palette[label] = colorPalette[i]
 
     for left, right in zip(broadcasted_left, right_idx):
         single_sankey(data[data[xvar]==left][yvar], data[data[xvar]==right][yvar], 
-                        xpos=xpos, ax=ax, width=width, leftLabels=leftLabels, 
-                        rightLabels=rightLabels, rightColor=rightColor, 
+                        xpos=xpos, ax=ax, colorDict=plot_palette, width=width, 
+                        leftLabels=leftLabels, rightLabels=rightLabels, 
+                        rightColor=rightColor, 
                         align=align, alpha=alpha)
         xpos += 1
 

dabest/plotter.py

@@ -31,6 +31,7 @@ def EffectSizeDataFramePlotter(EffectSizeDataFrame, **plot_kwargs):
         swarmplot_kwargs=None,
         violinplot_kwargs=None,
         slopegraph_kwargs=None,
+        sankey_kwargs=None,
         reflines_kwargs=None,
         group_summary_kwargs=None,
         legend_kwargs=None,
@@ -248,6 +249,11 @@ def EffectSizeDataFramePlotter(EffectSizeDataFrame, **plot_kwargs):
 
         contrast_colors = [sns.desaturate(c, contrast_desat) for c in unsat_colors]
         plot_palette_contrast = dict(zip(names.categories, contrast_colors))
+
+        # For Sankey Diagram plot, no need to worry about the color, each bar will have the same two colors
+        # default color palette will be set to "hls"
+        plot_palette_sankey = None
+        
     else:
         swarm_colors = [sns.desaturate(c, swarm_desat) for c in unsat_colors]
         plot_palette_raw = dict(zip(names, swarm_colors))
@@ -258,6 +264,8 @@ def EffectSizeDataFramePlotter(EffectSizeDataFrame, **plot_kwargs):
         contrast_colors = [sns.desaturate(c, contrast_desat) for c in unsat_colors]
         plot_palette_contrast = dict(zip(names, contrast_colors))
 
+        plot_palette_sankey = custom_pal
+
     # Infer the figsize.
     fig_size   = plot_kwargs["fig_size"]
     if fig_size is None:
@@ -445,9 +453,10 @@ def EffectSizeDataFramePlotter(EffectSizeDataFrame, **plot_kwargs):
 
             # Replace the paired proportional plot with sankey diagram
             sankey = sankeydiag(plot_data, xvar=xvar, yvar=yvar, 
-                                ax=rawdata_axes, 
                                 left_idx=sankey_control_group, 
                                 right_idx=sankey_test_group,
+                                palette=plot_palette_sankey,
+                                ax=rawdata_axes, 
                                 **sankey_kwargs)
 
     else:
@@ -561,9 +570,6 @@ def EffectSizeDataFramePlotter(EffectSizeDataFrame, **plot_kwargs):
         ticks_to_plot = np.arange(1, len(temp_all_plot_groups), 2).tolist()
         ticks_to_skip_contrast = np.cumsum([(len(t)-1)*2 for t in idx])[:-1].tolist()
         ticks_to_skip_contrast.insert(0, 0)
-    # elif is_paired == "sequential" and proportional == True:
-    #     ticks_to_skip = []
-    #     ticks_to_plot = np.arange(0, len(all_plot_groups)-1).tolist()
     else:
         ticks_to_skip = np.cumsum([len(t) for t in idx])[:-1].tolist()
         ticks_to_skip.insert(0, 0)
@@ -590,7 +596,7 @@ def EffectSizeDataFramePlotter(EffectSizeDataFrame, **plot_kwargs):
         current_ci_high   = results.bca_high[j]
 
         #TODO: Warnings for bias-corrected and accelerated confidence intervals
-        
+
         # Create the violinplot.
         # New in v0.2.6: drop negative infinities before plotting.
         v = contrast_axes.violinplot(current_bootstrap[~np.isinf(current_bootstrap)],