Added custom_pal list functionality to sankey plots. All types of custom_pal should now work with all types of plots

Author: JAnns98

dabest/misc_tools.py

@@ -549,9 +549,15 @@ def get_color_palette(
             filled.extend([True] * (len(idx[i]) - 1))
 
     if color_col is not None:
-        names = color_groups if not color_by_subgroups else idx
+        if sankey:
+            names = [1, 0]
+        else:
+            names = color_groups if not color_by_subgroups else idx
     else:
-        names = all_plot_groups if not color_by_subgroups else idx
+        if sankey:
+            names = [1, 0]
+        else:
+            names = all_plot_groups if not color_by_subgroups else idx
 
     n_groups = len(color_groups)
     custom_pal = plot_kwargs["custom_palette"]
@@ -584,7 +590,16 @@ def get_color_palette(
             unsat_colors = groups_in_palette.values()
 
         elif isinstance(custom_pal, list):
-            if len(custom_pal) < n_groups:
+            if sankey:
+                if len(custom_pal) != 2:
+                    raise ValueError("To specify a custom palette for a Sankey diagram, you must provide exactly two colors.")
+                else:
+                    groups_in_palette = {
+                        k: custom_pal[k] for k in [1, 0]
+                    }
+                    names = groups_in_palette.keys()
+                    unsat_colors = groups_in_palette.values()
+            elif len(custom_pal) < n_groups:
                 err1 = "The specified `custom_palette` has fewer colors than the number of groups."
                 err2 = " Please specify a custom palette with at least {} colors.".format(n_groups)
                 raise ValueError(err1 + err2)
@@ -618,11 +633,6 @@ def get_color_palette(
             plot_palette_raw = dict(zip(categories, swarm_colors))
             plot_palette_contrast = dict(zip(categories, contrast_colors))
             plot_palette_bar = dict(zip(categories, bar_color))
-
-        # 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]
         contrast_colors = [sns.desaturate(c, contrast_desat) for c in unsat_colors]
@@ -640,8 +650,12 @@ def get_color_palette(
             plot_palette_raw = dict(zip(names, swarm_colors))
             plot_palette_contrast = dict(zip(names, contrast_colors))
             plot_palette_bar = dict(zip(names, bar_color))
+            plot_palette_sankey = dict(zip(names, unsat_colors))
 
-        plot_palette_sankey = custom_pal
+    # For Sankey Diagram plot, each bar will have the same two colors if custom_pal is None
+    # default color palette will be set to "hls"
+    if custom_pal is None:
+        plot_palette_sankey = None
 
     return (color_col, bootstraps_color_by_group, n_groups, filled, plot_palette_raw, bar_color, 
             plot_palette_bar, plot_palette_contrast, plot_palette_sankey)

nbs/API/misc_tools.ipynb

@@ -602,9 +602,15 @@
     "            filled.extend([True] * (len(idx[i]) - 1))\n",
     "\n",
     "    if color_col is not None:\n",
-    "        names = color_groups if not color_by_subgroups else idx\n",
+    "        if sankey:\n",
+    "            names = [1, 0]\n",
+    "        else:\n",
+    "            names = color_groups if not color_by_subgroups else idx\n",
     "    else:\n",
-    "        names = all_plot_groups if not color_by_subgroups else idx\n",
+    "        if sankey:\n",
+    "            names = [1, 0]\n",
+    "        else:\n",
+    "            names = all_plot_groups if not color_by_subgroups else idx\n",
     "\n",
     "    n_groups = len(color_groups)\n",
     "    custom_pal = plot_kwargs[\"custom_palette\"]\n",
@@ -637,7 +643,16 @@
     "            unsat_colors = groups_in_palette.values()\n",
     "\n",
     "        elif isinstance(custom_pal, list):\n",
-    "            if len(custom_pal) < n_groups:\n",
+    "            if sankey:\n",
+    "                if len(custom_pal) != 2:\n",
+    "                    raise ValueError(\"To specify a custom palette for a Sankey diagram, you must provide exactly two colors.\")\n",
+    "                else:\n",
+    "                    groups_in_palette = {\n",
+    "                        k: custom_pal[k] for k in [1, 0]\n",
+    "                    }\n",
+    "                    names = groups_in_palette.keys()\n",
+    "                    unsat_colors = groups_in_palette.values()\n",
+    "            elif len(custom_pal) < n_groups:\n",
     "                err1 = \"The specified `custom_palette` has fewer colors than the number of groups.\"\n",
     "                err2 = \" Please specify a custom palette with at least {} colors.\".format(n_groups)\n",
     "                raise ValueError(err1 + err2)\n",
@@ -671,11 +686,6 @@
     "            plot_palette_raw = dict(zip(categories, swarm_colors))\n",
     "            plot_palette_contrast = dict(zip(categories, contrast_colors))\n",
     "            plot_palette_bar = dict(zip(categories, bar_color))\n",
-    "\n",
-    "        # For Sankey Diagram plot, no need to worry about the color, each bar will have the same two colors\n",
-    "        # default color palette will be set to \"hls\"\n",
-    "        plot_palette_sankey = None\n",
-    "\n",
     "    else:\n",
     "        swarm_colors = [sns.desaturate(c, swarm_desat) for c in unsat_colors]\n",
     "        contrast_colors = [sns.desaturate(c, contrast_desat) for c in unsat_colors]\n",
@@ -693,8 +703,12 @@
     "            plot_palette_raw = dict(zip(names, swarm_colors))\n",
     "            plot_palette_contrast = dict(zip(names, contrast_colors))\n",
     "            plot_palette_bar = dict(zip(names, bar_color))\n",
+    "            plot_palette_sankey = dict(zip(names, unsat_colors))\n",
     "\n",
-    "        plot_palette_sankey = custom_pal\n",
+    "    # For Sankey Diagram plot, each bar will have the same two colors if custom_pal is None\n",
+    "    # default color palette will be set to \"hls\"\n",
+    "    if custom_pal is None:\n",
+    "        plot_palette_sankey = None\n",
     "\n",
     "    return (color_col, bootstraps_color_by_group, n_groups, filled, plot_palette_raw, bar_color, \n",
     "            plot_palette_bar, plot_palette_contrast, plot_palette_sankey)\n",

nbs/tests/mpl_image_tests/baseline_images/test_143_sankey_change_palette_c.png

@@ -416,6 +416,9 @@ def test_141_sankey_change_palette_a():
 def test_142_sankey_change_palette_b():
     return multi_groups_paired.mean_diff.plot(custom_palette={1: 'red', 0: 'blue'})
 
+@pytest.mark.mpl_image_compare(tolerance=8)
+def test_143_sankey_change_palette_c():
+    return multi_groups_paired.mean_diff.plot(custom_palette=['red', 'blue'])
 
 @pytest.mark.mpl_image_compare(tolerance=8)
 def test_136_style_sheets():