Sankey custom palette fix

Fixed the functionality to supply a custom_palette dict for the colors associated with values 0 and 1 {0: color, 1: color}

Author: JAnns98

dabest/misc_tools.py

@@ -486,6 +486,7 @@ def get_color_palette(
         idx: list, 
         all_plot_groups: list,
         delta2: bool,
+        sankey: bool
     ):
     """
     Create the color palette to be used in the plotter function.
@@ -506,6 +507,8 @@ def get_color_palette(
         A list of all the group names.
     delta2 : bool
         A boolean flag to determine if the plot will have a delta-delta effect size.
+    sankey : bool
+        A boolean flag to determine if the plot is for a Sankey diagram.
     """
     # Create color palette that will be shared across subplots.
     color_col = plot_kwargs["color_col"]
@@ -560,6 +563,10 @@ def get_color_palette(
                 groups_in_palette = {
                     k: custom_pal[k] for k in color_groups
                 }
+            elif sankey:
+                groups_in_palette = {
+                    k: custom_pal[k] for k in [1, 0]
+                }
             elif color_col is None:
                 groups_in_palette = {
                     k: custom_pal[k] for k in all_plot_groups if k in color_groups

dabest/plotter.py

@@ -152,7 +152,8 @@ def effectsize_df_plotter(effectsize_df: object, **plot_kwargs) -> matplotlib.fi
                                                                                         show_pairs = show_pairs,
                                                                                         idx = idx,
                                                                                         all_plot_groups = all_plot_groups,
-                                                                                        delta2 = effectsize_df.delta2
+                                                                                        delta2 = effectsize_df.delta2,
+                                                                                        sankey = True if proportional and show_pairs else False,
                                                                                         )
 
     # Initialise the figure.

nbs/API/misc_tools.ipynb

@@ -539,6 +539,7 @@
     "        idx: list, \n",
     "        all_plot_groups: list,\n",
     "        delta2: bool,\n",
+    "        sankey: bool\n",
     "    ):\n",
     "    \"\"\"\n",
     "    Create the color palette to be used in the plotter function.\n",
@@ -559,6 +560,8 @@
     "        A list of all the group names.\n",
     "    delta2 : bool\n",
     "        A boolean flag to determine if the plot will have a delta-delta effect size.\n",
+    "    sankey : bool\n",
+    "        A boolean flag to determine if the plot is for a Sankey diagram.\n",
     "    \"\"\"\n",
     "    # Create color palette that will be shared across subplots.\n",
     "    color_col = plot_kwargs[\"color_col\"]\n",
@@ -613,6 +616,10 @@
     "                groups_in_palette = {\n",
     "                    k: custom_pal[k] for k in color_groups\n",
     "                }\n",
+    "            elif sankey:\n",
+    "                groups_in_palette = {\n",
+    "                    k: custom_pal[k] for k in [1, 0]\n",
+    "                }\n",
     "            elif color_col is None:\n",
     "                groups_in_palette = {\n",
     "                    k: custom_pal[k] for k in all_plot_groups if k in color_groups\n",

nbs/API/plotter.ipynb

@@ -209,7 +209,8 @@
     "                                                                                        show_pairs = show_pairs,\n",
     "                                                                                        idx = idx,\n",
     "                                                                                        all_plot_groups = all_plot_groups,\n",
-    "                                                                                        delta2 = effectsize_df.delta2\n",
+    "                                                                                        delta2 = effectsize_df.delta2,\n",
+    "                                                                                        sankey = True if proportional and show_pairs else False,\n",
     "                                                                                        )\n",
     "\n",
     "    # Initialise the figure.\n",

nbs/tests/mpl_image_tests/baseline_images/test_141_sankey_change_palette_a.png

@@ -407,6 +407,16 @@ def test_139_multi_2group_show_sample_counts_and_kwargs():
 def test_140_multi_groups_paired_show_sample_counts_with_sankey_off():
     return multi_groups_paired.mean_diff.plot(prop_sample_counts=True, sankey_kwargs={"sankey": False})
 
+
+@pytest.mark.mpl_image_compare(tolerance=8)
+def test_141_sankey_change_palette_a():
+    return multi_groups_paired.mean_diff.plot(custom_palette="Dark2")
+
+@pytest.mark.mpl_image_compare(tolerance=8)
+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_136_style_sheets():
     # Perform this test last so we don't have to reset the plot style.