Bugfix #218

dabest/misc_tools.py

@@ -546,6 +546,7 @@ def get_color_palette(
 
     # Create color palette that will be shared across subplots.
     color_col = plot_kwargs["color_col"]
+
     if color_col is None:
         color_groups = pd.unique(plot_data[xvar])
         bootstraps_color_by_group = True
@@ -555,11 +556,10 @@ def get_color_palette(
         color_groups = pd.unique(plot_data[color_col])
         bootstraps_color_by_group = False
     if show_pairs:
-        if plot_kwargs["custom_palette"] is not None:
-            if delta2 or sankey:
-                bootstraps_color_by_group = False
-            else:
-                bootstraps_color_by_group = True
+        # When `color_col` is given, the palette is keyed by the `color_col`
+        # categories, so the bootstraps cannot be coloured by the x-axis group.
+        if plot_kwargs["custom_palette"] is not None and color_col is None:
+            bootstraps_color_by_group = not (delta2 or sankey)
         else:
             bootstraps_color_by_group = False
 
@@ -632,7 +632,14 @@ def get_color_palette(
                     k: custom_pal[k] for k in all_plot_groups if k in color_groups
                 }
             else:
-                raise ValueError("The `custom_palette` dictionary is not supported when `color_col` is not None.")
+                missing = [k for k in color_groups if k not in custom_pal]
+                if missing:
+                    err1 = "The `custom_palette` dictionary is missing colors for the "
+                    err2 = "following `{}` groups: {}.".format(color_col, missing)
+                    raise ValueError(err1 + err2)
+                groups_in_palette = {
+                    k: custom_pal[k] for k in color_groups
+                }
 
             names = groups_in_palette.keys()
             unsat_colors = groups_in_palette.values()

nbs/API/misc_tools.ipynb

@@ -597,6 +597,7 @@
     "\n",
     "    # Create color palette that will be shared across subplots.\n",
     "    color_col = plot_kwargs[\"color_col\"]\n",
+    "\n",
     "    if color_col is None:\n",
     "        color_groups = pd.unique(plot_data[xvar])\n",
     "        bootstraps_color_by_group = True\n",
@@ -606,11 +607,10 @@
     "        color_groups = pd.unique(plot_data[color_col])\n",
     "        bootstraps_color_by_group = False\n",
     "    if show_pairs:\n",
-    "        if plot_kwargs[\"custom_palette\"] is not None:\n",
-    "            if delta2 or sankey:\n",
-    "                bootstraps_color_by_group = False\n",
-    "            else:\n",
-    "                bootstraps_color_by_group = True\n",
+    "        # When `color_col` is given, the palette is keyed by the `color_col`\n",
+    "        # categories, so the bootstraps cannot be coloured by the x-axis group.\n",
+    "        if plot_kwargs[\"custom_palette\"] is not None and color_col is None:\n",
+    "            bootstraps_color_by_group = not (delta2 or sankey)\n",
     "        else:\n",
     "            bootstraps_color_by_group = False\n",
     "\n",
@@ -683,7 +683,14 @@
     "                    k: custom_pal[k] for k in all_plot_groups if k in color_groups\n",
     "                }\n",
     "            else:\n",
-    "                raise ValueError(\"The `custom_palette` dictionary is not supported when `color_col` is not None.\")\n",
+    "                missing = [k for k in color_groups if k not in custom_pal]\n",
+    "                if missing:\n",
+    "                    err1 = \"The `custom_palette` dictionary is missing colors for the \"\n",
+    "                    err2 = \"following `{}` groups: {}.\".format(color_col, missing)\n",
+    "                    raise ValueError(err1 + err2)\n",
+    "                groups_in_palette = {\n",
+    "                    k: custom_pal[k] for k in color_groups\n",
+    "                }\n",
     "\n",
     "            names = groups_in_palette.keys()\n",
     "            unsat_colors = groups_in_palette.values()\n",

nbs/tests/test_color_palette.py

@@ -0,0 +1,141 @@
+"""
+Tests for `get_color_palette`, which decides how the raw data, the slopegraph
+and the bootstrap distributions are coloured.
+
+Regression coverage for issue #218: combining `color_col` with a
+`custom_palette` on a paired plot used to raise a `KeyError`, because the
+palette is keyed by the `color_col` categories while the bootstraps were
+still being coloured by the x-axis group.
+"""
+
+import pytest
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+from dabest import load
+from dabest.misc_tools import get_color_palette
+
+
+N = 20
+IDX = ("Control 1", "Test 1")
+ALL_PLOT_GROUPS = list(IDX)
+COLOR_GROUPS = ["Female", "Male"]
+
+
+@pytest.fixture
+def df():
+    np.random.seed(9999)
+    return pd.DataFrame(
+        {
+            "Control 1": np.random.normal(3, 0.4, N),
+            "Test 1": np.random.normal(3.5, 0.5, N),
+            "Gender": ["Female"] * (N // 2) + ["Male"] * (N // 2),
+            "ID": range(1, N + 1),
+        }
+    )
+
+
+@pytest.fixture
+def plot_data(df):
+    return df.melt(
+        id_vars=["Gender", "ID"], var_name="group", value_name="value"
+    )
+
+
+def make_plot_kwargs(color_col=None, custom_palette=None):
+    return {
+        "color_col": color_col,
+        "custom_palette": custom_palette,
+        "empty_circle": False,
+        "raw_desat": 1.0,
+        "contrast_desat": 1.0,
+    }
+
+
+def call(plot_data, color_col=None, custom_palette=None, show_pairs=True):
+    return get_color_palette(
+        plot_kwargs=make_plot_kwargs(color_col, custom_palette),
+        plot_data=plot_data,
+        xvar="group",
+        show_pairs=show_pairs,
+        idx=IDX,
+        all_plot_groups=ALL_PLOT_GROUPS,
+        delta2=False,
+        proportional=False,
+    )
+
+
+def test_paired_color_col_with_list_palette(plot_data):
+    # The palette is keyed by the `color_col` categories, so the bootstraps
+    # must not be coloured by the x-axis group.
+    (color_col, bootstraps_color_by_group, n_groups, _, _,
+     plot_palette_raw, plot_palette_contrast, _) = call(
+        plot_data, color_col="Gender", custom_palette=["red", "blue"]
+    )
+
+    assert color_col == "Gender"
+    assert bootstraps_color_by_group is False
+    assert n_groups == 2
+    assert list(plot_palette_raw.keys()) == COLOR_GROUPS
+    assert list(plot_palette_contrast.keys()) == COLOR_GROUPS
+
+
+def test_paired_color_col_with_dict_palette(plot_data):
+    palette = {"Female": "red", "Male": "blue"}
+    (_, bootstraps_color_by_group, _, _, _,
+     plot_palette_raw, _, _) = call(
+        plot_data, color_col="Gender", custom_palette=palette
+    )
+
+    assert bootstraps_color_by_group is False
+    assert list(plot_palette_raw.keys()) == COLOR_GROUPS
+
+
+def test_paired_dict_palette_missing_color_raises(plot_data):
+    with pytest.raises(ValueError) as excinfo:
+        call(plot_data, color_col="Gender", custom_palette={"Female": "red"})
+
+    assert "missing colors" in str(excinfo.value)
+    assert "Male" in str(excinfo.value)
+
+
+def test_paired_custom_palette_without_color_col(plot_data):
+    # Issue #207: without a `color_col`, a custom palette colours the paired
+    # groups and the bootstraps follow the x-axis group.
+    (color_col, bootstraps_color_by_group, _, _, _,
+     plot_palette_raw, _, _) = call(plot_data, custom_palette=["red", "blue"])
+
+    assert color_col is None
+    assert bootstraps_color_by_group is True
+    assert list(plot_palette_raw.keys()) == ALL_PLOT_GROUPS
+
+
+def test_paired_without_custom_palette(plot_data):
+    _, bootstraps_color_by_group, _, _, _, _, _, _ = call(plot_data)
+    assert bootstraps_color_by_group is False
+
+
+def test_unpaired_color_col_with_custom_palette(plot_data):
+    _, bootstraps_color_by_group, _, _, _, plot_palette_raw, _, _ = call(
+        plot_data,
+        color_col="Gender",
+        custom_palette=["red", "blue"],
+        show_pairs=False,
+    )
+
+    assert bootstraps_color_by_group is False
+    assert list(plot_palette_raw.keys()) == COLOR_GROUPS
+
+
+@pytest.mark.parametrize(
+    "custom_palette",
+    [["red", "blue"], {"Female": "red", "Male": "blue"}, "Dark2"],
+)
+@pytest.mark.parametrize("paired", ["baseline", "sequential"])
+def test_paired_plot_with_color_col_and_custom_palette(df, custom_palette, paired):
+    # Issue #218: this used to raise `KeyError: 'Test 1'`.
+    loaded = load(df, idx=IDX, paired=paired, id_col="ID")
+    fig = loaded.mean_diff.plot(color_col="Gender", custom_palette=custom_palette)
+    assert fig is not None
+    plt.close(fig)

pyproject.toml

@@ -28,7 +28,7 @@ dev = ['pytest~=8.3.4', 'pytest-mpl~=0.17.0']
 version = {attr = "dabest.__version__"}
 
 [tool.setuptools.packages.find]
-include = ["dabest"]
+include = ["dabest", "dabest.*"]
 
 [tool.nbdev]
 branch = 'master'

setup.py

@@ -46,7 +46,7 @@
     extras_require={ 'dev': dev_requirements },
     dependency_links = cfg.get('dep_links','').split(),
     python_requires  = '>=' + cfg['min_python'],
-    long_description = open('README.md').read(),
+    long_description = open('README.md', encoding='utf-8').read(),
     long_description_content_type = 'text/markdown',
     zip_safe = False,
     entry_points = {