Contrast bars added to forest plot

Author: JAnns98

dabest/forest_plot.py

@@ -11,6 +11,8 @@
 import seaborn as sns
 from typing import List, Optional, Union
 import numpy as np
+import matplotlib.axes as axes
+import matplotlib.patches as mpatches
 
 # %% ../nbs/API/forest_plot.ipynb 6
 def load_plot_data(
@@ -266,6 +268,7 @@ def get_kwargs(
         marker_kwargs,
         errorbar_kwargs,
         delta_text_kwargs,
+        contrast_bars_kwargs,
         marker_size
     ):
     from .misc_tools import merge_two_dicts
@@ -336,7 +339,21 @@ def get_kwargs(
     else:
         delta_text_kwargs = merge_two_dicts(default_delta_text_kwargs, delta_text_kwargs)
 
-    return violin_kwargs, zeroline_kwargs, marker_kwargs, errorbar_kwargs, delta_text_kwargs
+    # Contrast bars kwargs.
+    default_contrast_bars_kwargs = {
+                    "color": None, 
+                    "zorder":-3,
+                    'alpha': 0.15
+    }
+    if contrast_bars_kwargs is None:
+        contrast_bars_kwargs = default_contrast_bars_kwargs
+    else:
+        contrast_bars_kwargs = merge_two_dicts(default_contrast_bars_kwargs, contrast_bars_kwargs)
+
+
+    return (violin_kwargs, zeroline_kwargs, marker_kwargs, errorbar_kwargs, 
+            delta_text_kwargs, contrast_bars_kwargs)
+
 
 
 def color_palette(
@@ -393,6 +410,9 @@ def forest_plot(
     delta_text: bool = True,
     delta_text_kwargs: dict = None,
 
+    contrast_bars: bool = True,
+    contrast_bars_kwargs: dict = None,
+
     violin_kwargs: Optional[dict] = None,
     zeroline_kwargs: Optional[dict] = None,
     marker_kwargs: Optional[dict] = None,
@@ -512,14 +532,15 @@ def forest_plot(
         fig, ax = plt.subplots(figsize=fig_size)
 
     # Get Kwargs
-    (violin_kwargs, zeroline_kwargs, 
-    marker_kwargs, errorbar_kwargs, delta_text_kwargs) = get_kwargs(
+    (violin_kwargs, zeroline_kwargs, marker_kwargs, 
+     errorbar_kwargs, delta_text_kwargs, contrast_bars_kwargs) = get_kwargs(
                                                             violin_kwargs = violin_kwargs,
                                                             zeroline_kwargs = zeroline_kwargs,
                                                             horizontal = horizontal,
                                                             marker_kwargs = marker_kwargs,
                                                             errorbar_kwargs = errorbar_kwargs,
                                                             delta_text_kwargs = delta_text_kwargs,
+                                                            contrast_bars_kwargs = contrast_bars_kwargs,
                                                             marker_size = marker_size
     )
 
@@ -661,6 +682,19 @@ def forest_plot(
             delta_text = np.format_float_positional(delta, precision=2, sign=True, trim="k", min_digits=2)
             ax.text(x, y, delta_text, color=delta_text_colors[idx], zorder=5, **delta_text_kwargs)
 
+    # Contrast bars
+    if contrast_bars:
+        _bar_color = contrast_bars_kwargs.pop('color')
+        if _bar_color is not None:
+            bar_colors = [_bar_color] * number_of_curves_to_plot
+        else:
+            bar_colors = violin_colors
+        for x, y in zip(np.arange(1, number_of_curves_to_plot + 1), differences):
+            if horizontal:
+                ax.add_patch(mpatches.Rectangle((0, x-0.25), y, 0.25, color=bar_colors[x-1], **contrast_bars_kwargs))
+            else:
+                ax.add_patch(mpatches.Rectangle((x, 0), 0.25, y, color=bar_colors[x-1], **contrast_bars_kwargs))
+
     ## Invert Y-axis if horizontal 
     if horizontal:
         ax.invert_yaxis()

nbs/API/forest_plot.ipynb

@@ -63,7 +63,9 @@
     "# %matplotlib inline\n",
     "import seaborn as sns\n",
     "from typing import List, Optional, Union\n",
-    "import numpy as np"
+    "import numpy as np\n",
+    "import matplotlib.axes as axes\n",
+    "import matplotlib.patches as mpatches"
    ]
   },
   {
@@ -326,6 +328,7 @@
     "        marker_kwargs,\n",
     "        errorbar_kwargs,\n",
     "        delta_text_kwargs,\n",
+    "        contrast_bars_kwargs,\n",
     "        marker_size\n",
     "    ):\n",
     "    from .misc_tools import merge_two_dicts\n",
@@ -396,7 +399,21 @@
     "    else:\n",
     "        delta_text_kwargs = merge_two_dicts(default_delta_text_kwargs, delta_text_kwargs)\n",
     "\n",
-    "    return violin_kwargs, zeroline_kwargs, marker_kwargs, errorbar_kwargs, delta_text_kwargs\n",
+    "    # Contrast bars kwargs.\n",
+    "    default_contrast_bars_kwargs = {\n",
+    "                    \"color\": None, \n",
+    "                    \"zorder\":-3,\n",
+    "                    'alpha': 0.15\n",
+    "    }\n",
+    "    if contrast_bars_kwargs is None:\n",
+    "        contrast_bars_kwargs = default_contrast_bars_kwargs\n",
+    "    else:\n",
+    "        contrast_bars_kwargs = merge_two_dicts(default_contrast_bars_kwargs, contrast_bars_kwargs)\n",
+    "\n",
+    "\n",
+    "    return (violin_kwargs, zeroline_kwargs, marker_kwargs, errorbar_kwargs, \n",
+    "            delta_text_kwargs, contrast_bars_kwargs)\n",
+    "\n",
     "\n",
     "\n",
     "def color_palette(\n",
@@ -453,6 +470,9 @@
     "    delta_text: bool = True,\n",
     "    delta_text_kwargs: dict = None,\n",
     "\n",
+    "    contrast_bars: bool = True,\n",
+    "    contrast_bars_kwargs: dict = None,\n",
+    "\n",
     "    violin_kwargs: Optional[dict] = None,\n",
     "    zeroline_kwargs: Optional[dict] = None,\n",
     "    marker_kwargs: Optional[dict] = None,\n",
@@ -572,14 +592,15 @@
     "        fig, ax = plt.subplots(figsize=fig_size)\n",
     "\n",
     "    # Get Kwargs\n",
-    "    (violin_kwargs, zeroline_kwargs, \n",
-    "    marker_kwargs, errorbar_kwargs, delta_text_kwargs) = get_kwargs(\n",
+    "    (violin_kwargs, zeroline_kwargs, marker_kwargs, \n",
+    "     errorbar_kwargs, delta_text_kwargs, contrast_bars_kwargs) = get_kwargs(\n",
     "                                                            violin_kwargs = violin_kwargs,\n",
     "                                                            zeroline_kwargs = zeroline_kwargs,\n",
     "                                                            horizontal = horizontal,\n",
     "                                                            marker_kwargs = marker_kwargs,\n",
     "                                                            errorbar_kwargs = errorbar_kwargs,\n",
     "                                                            delta_text_kwargs = delta_text_kwargs,\n",
+    "                                                            contrast_bars_kwargs = contrast_bars_kwargs,\n",
     "                                                            marker_size = marker_size\n",
     "    )\n",
     "                                            \n",
@@ -721,6 +742,19 @@
     "            delta_text = np.format_float_positional(delta, precision=2, sign=True, trim=\"k\", min_digits=2)\n",
     "            ax.text(x, y, delta_text, color=delta_text_colors[idx], zorder=5, **delta_text_kwargs)\n",
     "\n",
+    "    # Contrast bars\n",
+    "    if contrast_bars:\n",
+    "        _bar_color = contrast_bars_kwargs.pop('color')\n",
+    "        if _bar_color is not None:\n",
+    "            bar_colors = [_bar_color] * number_of_curves_to_plot\n",
+    "        else:\n",
+    "            bar_colors = violin_colors\n",
+    "        for x, y in zip(np.arange(1, number_of_curves_to_plot + 1), differences):\n",
+    "            if horizontal:\n",
+    "                ax.add_patch(mpatches.Rectangle((0, x-0.25), y, 0.25, color=bar_colors[x-1], **contrast_bars_kwargs))\n",
+    "            else:\n",
+    "                ax.add_patch(mpatches.Rectangle((x, 0), 0.25, y, color=bar_colors[x-1], **contrast_bars_kwargs))\n",
+    "\n",
     "    ## Invert Y-axis if horizontal \n",
     "    if horizontal:\n",
     "        ax.invert_yaxis()\n",

nbs/tests/mpl_image_tests/baseline_images/test_500_deltadelta_forest.png

@@ -401,4 +401,14 @@ def test_520_minimeta_with_deltas_and_delta_text_kwargs_forest():
                 delta_text_kwargs={'color': 'black','fontsize': 8, 'rotation': 45, 'va': 'bottom',
                                    'x_coordinates': [1.4, 2.4, 3.4, 4.4, 5.4, 6.4], 
                                    'y_coordinates': [0.6, 0.1, -2, -1.5, -1.5, -1.5]}
+            )
+
+@pytest.mark.mpl_image_compare(tolerance=8)
+def test_521_minimeta_with_deltas_with_contrast_bars_kwargs_forest():
+    plt.rcdefaults()
+    return forest_plot(
+                contrasts_mini_meta, 
+                idx=[(0, 3),(0, 3),(0, 3)],
+                labels=['Contrast A1', 'Mini_Meta A', 'Contrast B1', 'Mini_Meta B', 'Contrast C1', 'Mini_Meta C'],
+                contrast_bars_kwargs={'color': 'red', 'alpha': 0.4}
             )