diff --git a/dabest/misc_tools.py b/dabest/misc_tools.py index 4060b7ca..57c9e2fb 100644 --- a/dabest/misc_tools.py +++ b/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() diff --git a/nbs/API/misc_tools.ipynb b/nbs/API/misc_tools.ipynb index 3e28644e..8eb2b557 100644 --- a/nbs/API/misc_tools.ipynb +++ b/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", diff --git a/nbs/tests/test_color_palette.py b/nbs/tests/test_color_palette.py new file mode 100644 index 00000000..6b74342f --- /dev/null +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 533ba62c..289d1cf1 100644 --- a/pyproject.toml +++ b/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' diff --git a/setup.py b/setup.py index bb8943b0..a9d41ea3 100644 --- a/setup.py +++ b/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 = {