Add counts to prop plots feature and added tests and tutorial details for this feature

Author: JAnns98

dabest/_effsize_objects.py

@@ -1037,6 +1037,9 @@ def plot(
 
         es_marker_kwargs=None,
         es_errorbar_kwargs=None,
+
+        prop_sample_counts=False,
+        prop_sample_counts_kwargs=None
     ):
         """
         Creates an estimation plot for the effect size of interest.
@@ -1126,6 +1129,9 @@ def plot(
             Pass any keyword arguments accepted by the seaborn `swarmplot`
             command here, as a dict. If None, the following keywords are
             passed to sns.swarmplot : {'size':`raw_marker_size`}.
+        barplot_kwargs : dict, default None
+            By default, the keyword arguments passed are:
+            {"estimator": np.mean, "errorbar": plot_kwargs["ci"]}
         violinplot_kwargs : dict, default None
             Pass any keyword arguments accepted by the matplotlib `
             pyplot.violinplot` command here, as a dict. If None, the following
@@ -1241,6 +1247,13 @@ def plot(
             Pass relevant keyword arguments to the effectsize errorbar plotting. If none, the following keywords are passed:
             {'color': 'black', 'lw': 2, 'linestyle': '-', 'alpha': 1,'zorder': 1,}
 
+        prop_sample_counts: bool, default False
+            Show the sample counts for each group in proportional plots
+        prop_sample_counts_kwargs: dict, default None
+            Pass relevant keyword arguments. If None, the following keywords are passed:
+            {'color': 'k', 'zorder': 5, 'ha': 'center', 'va': 'center'}
+        
+
         Returns
         -------
         A :class:`matplotlib.figure.Figure` with 2 Axes, if ``ax = None``.

dabest/_modidx.py

@@ -106,6 +106,8 @@
                                                                                     'dabest/plot_tools.py'),
                                    'dabest.plot_tools.SwarmPlot._swarm': ('API/plot_tools.html#swarmplot._swarm', 'dabest/plot_tools.py'),
                                    'dabest.plot_tools.SwarmPlot.plot': ('API/plot_tools.html#swarmplot.plot', 'dabest/plot_tools.py'),
+                                   'dabest.plot_tools.add_counts_to_prop_plots': ( 'API/plot_tools.html#add_counts_to_prop_plots',
+                                                                                   'dabest/plot_tools.py'),
                                    'dabest.plot_tools.barplotter': ('API/plot_tools.html#barplotter', 'dabest/plot_tools.py'),
                                    'dabest.plot_tools.check_data_matches_labels': ( 'API/plot_tools.html#check_data_matches_labels',
                                                                                     'dabest/plot_tools.py'),

dabest/misc_tools.py

@@ -200,7 +200,7 @@ def get_kwargs(plot_kwargs, ytick_color):
     # Barplot kwargs
     default_barplot_kwargs = {
         "estimator": np.mean, 
-        "errorbar": plot_kwargs["ci"]
+        "errorbar": plot_kwargs["ci"],
     }
     if plot_kwargs["barplot_kwargs"] is None:
         barplot_kwargs = default_barplot_kwargs
@@ -435,11 +435,23 @@ def get_kwargs(plot_kwargs, ytick_color):
     else:
         es_errorbar_kwargs = merge_two_dicts(default_es_errorbar_kwargs, plot_kwargs['es_errorbar_kwargs'])
 
+    # Prop sample counts kwargs
+    default_prop_sample_counts_kwargs = {
+                                        'color': 'k', 
+                                        'zorder': 5, 
+                                        'ha': 'center', 
+                                        'va': 'center'
+    }
+    if plot_kwargs['prop_sample_counts_kwargs'] is None:
+        prop_sample_counts_kwargs = default_prop_sample_counts_kwargs
+    else:
+        prop_sample_counts_kwargs = merge_two_dicts(default_prop_sample_counts_kwargs, plot_kwargs['prop_sample_counts_kwargs'])
+
     # Return the kwargs.
     return (swarmplot_kwargs, barplot_kwargs, sankey_kwargs, violinplot_kwargs, slopegraph_kwargs, 
             reflines_kwargs, legend_kwargs, group_summaries_kwargs, redraw_axes_kwargs, delta_dot_kwargs,
             delta_text_kwargs, summary_bars_kwargs, swarm_bars_kwargs, contrast_bars_kwargs, table_kwargs, gridkey_kwargs,
-            es_marker_kwargs, es_errorbar_kwargs)
+            es_marker_kwargs, es_errorbar_kwargs, prop_sample_counts_kwargs)
 
 
 def get_color_palette(plot_kwargs, plot_data, xvar, show_pairs, idx, all_plot_groups):

dabest/plot_tools.py

@@ -9,8 +9,8 @@
 __all__ = ['halfviolin', 'get_swarm_spans', 'error_bar', 'check_data_matches_labels', 'normalize_dict', 'width_determine',
            'single_sankey', 'sankeydiag', 'summary_bars_plotter', 'contrast_bars_plotter', 'swarm_bars_plotter',
            'delta_text_plotter', 'DeltaDotsPlotter', 'slopegraph_plotter', 'plot_minimeta_or_deltadelta_violins',
-           'effect_size_curve_plotter', 'gridkey_plotter', 'barplotter', 'table_for_horizontal_plots', 'swarmplot',
-           'SwarmPlot']
+           'effect_size_curve_plotter', 'gridkey_plotter', 'barplotter', 'table_for_horizontal_plots',
+           'add_counts_to_prop_plots', 'swarmplot', 'SwarmPlot']
 
 # %% ../nbs/API/plot_tools.ipynb 4
 import math
@@ -1919,9 +1919,8 @@ def barplotter(xvar, yvar, all_plot_groups, rawdata_axes, plot_data, bar_color,
     horizontal : bool
         If the plot is horizontal.
     """
-    
-    x_label, y_label = rawdata_axes.get_xlabel(), rawdata_axes.get_ylabel()
 
+    x_label, y_label = rawdata_axes.get_xlabel(), rawdata_axes.get_ylabel()
     if horizontal:
         x_var, y_var, orient = np.ones(len(all_plot_groups)), all_plot_groups, "h"
     else:
@@ -1969,7 +1968,7 @@ def barplotter(xvar, yvar, all_plot_groups, rawdata_axes, plot_data, bar_color,
             centre = x + width / 2.0
             bar.set_x(centre - bar_width / 2.0)
             bar.set_width(bar_width)
-    
+
     # reset the x and y labels
     rawdata_axes.set_xlabel(x_label)
     rawdata_axes.set_ylabel(y_label)
@@ -2047,6 +2046,61 @@ def table_for_horizontal_plots(effectsize_df, ax, contrast_axes, ticks_to_plot,
     ax.set_xlabel(label, fontsize=fontsize_label) # Set the x-axis label - hardcoded for now
     sns.despine(ax=ax, left=True, bottom=True)
 
+
+def add_counts_to_prop_plots(plot_data, xvar, yvar, rawdata_axes, horizontal, is_paired, prop_sample_counts_kwargs):
+    """
+    Add counts to the proportion plots.
+
+    Parameters
+    ----------
+    plot_data : object (Dataframe)
+        Dataframe of the plot data.
+    xvar : str
+        Column name of the x variable.
+    yvar : str
+        Column name of the y variable.
+    rawdata_axes : object
+        Matplotlib axis object to plot on.
+    horizontal : bool
+        If the plot is horizontal.
+    is_paired : bool
+        Whether the data is paired.
+    prop_sample_counts_kwargs : dict
+        Keyword arguments for the sample counts.
+    """
+
+    # Group orders
+    if isinstance(plot_data[xvar].dtype, pd.CategoricalDtype):
+        sample_size_text_order = pd.unique(plot_data[xvar]).categories
+    else:
+        sample_size_text_order = pd.unique(plot_data[xvar])
+
+    # Get the sample size values
+    ones, zeros = plot_data[plot_data[yvar] == 1], plot_data[plot_data[yvar] == 0]
+
+    sample_size_val1 = ones.groupby(xvar, observed=False)[yvar].count().reindex(index=sample_size_text_order)
+    sample_size_val0 = zeros.groupby(xvar, observed=False)[yvar].count().reindex(index=sample_size_text_order)
+
+    fontsize = 8 if horizontal else 10
+    fontsize -= 2 if is_paired else 0
+
+    if "fontsize" not in prop_sample_counts_kwargs.keys():
+        fontsize = 8 if horizontal else 10
+        fontsize -= 2 if is_paired else 0
+        prop_sample_counts_kwargs.update({'fontsize': fontsize})
+
+    for sample_text_x, sample_text_y0, sample_text_y1 in zip(
+                                                            np.arange(0,len(sample_size_text_order)+1,1), 
+                                                            sample_size_val0,
+                                                            sample_size_val1,
+                                                            ):
+        if horizontal:
+            rawdata_axes.text(0.05, sample_text_x, sample_text_y1, **prop_sample_counts_kwargs)
+            rawdata_axes.text(0.95, sample_text_x, sample_text_y0, **prop_sample_counts_kwargs)
+        else:
+            rawdata_axes.text(sample_text_x, 0.05, sample_text_y1, **prop_sample_counts_kwargs)
+            rawdata_axes.text(sample_text_x, 0.95, sample_text_y0, **prop_sample_counts_kwargs)
+
 # %% ../nbs/API/plot_tools.ipynb 6
 def swarmplot(
     data: pd.DataFrame,

dabest/plotter.py

@@ -61,7 +61,8 @@ def effectsize_df_plotter(effectsize_df, **plot_kwargs):
         delta_text=True, delta_text_kwargs=None,
         delta_dot=True, delta_dot_kwargs=None,
         horizontal=False, horizontal_table_kwargs=None,
-        es_marker_kwargs=None, es_errorbar_kwargs=None
+        es_marker_kwargs=None, es_errorbar_kwargs=None,
+        prop_sample_counts=False, prop_sample_counts_kwargs=None
     """
     from .misc_tools import (
         get_params,
@@ -93,6 +94,7 @@ def effectsize_df_plotter(effectsize_df, **plot_kwargs):
         gridkey_plotter,
         barplotter,
         table_for_horizontal_plots,
+        add_counts_to_prop_plots,
     )
 
     warnings.filterwarnings(
@@ -121,13 +123,14 @@ def effectsize_df_plotter(effectsize_df, **plot_kwargs):
                                                                                                 plot_kwargs=plot_kwargs,
                                                                                                 )
 
-    (swarmplot_kwargs, barplot_kwargs, sankey_kwargs, violinplot_kwargs, 
-     slopegraph_kwargs, reflines_kwargs, legend_kwargs, group_summaries_kwargs, 
-     redraw_axes_kwargs, delta_dot_kwargs, delta_text_kwargs, summary_bars_kwargs, 
-     swarm_bars_kwargs, contrast_bars_kwargs, table_kwargs, gridkey_kwargs, es_marker_kwargs, es_errorbar_kwargs) = get_kwargs(
-                                                                                                                            plot_kwargs=plot_kwargs, 
-                                                                                                                            ytick_color=ytick_color
-                                                                                                                            )
+    (swarmplot_kwargs, barplot_kwargs, sankey_kwargs, 
+     violinplot_kwargs, slopegraph_kwargs, reflines_kwargs, 
+     legend_kwargs, group_summaries_kwargs, redraw_axes_kwargs, delta_dot_kwargs, 
+     delta_text_kwargs, summary_bars_kwargs, swarm_bars_kwargs, contrast_bars_kwargs, 
+     table_kwargs, gridkey_kwargs, es_marker_kwargs, es_errorbar_kwargs, prop_sample_counts_kwargs) = get_kwargs(
+                                                                                                                plot_kwargs=plot_kwargs, 
+                                                                                                                ytick_color=ytick_color
+                                                                                                                )
 
     # We also need to extract the `sankey` and `flow` from the kwargs for plotter.py
     # to use for varying different kinds of paired proportional plots
@@ -318,6 +321,18 @@ def effectsize_df_plotter(effectsize_df, **plot_kwargs):
             horizontal=horizontal,
             )
 
+    # Add counts to prop plots
+    if proportional and plot_kwargs['prop_sample_counts'] and sankey_kwargs["flow"]:
+        add_counts_to_prop_plots(
+                        plot_data=plot_data, 
+                        xvar=xvar, 
+                        yvar=yvar, 
+                        rawdata_axes=rawdata_axes, 
+                        horizontal=horizontal,
+                        is_paired = is_paired,
+                        prop_sample_counts_kwargs=prop_sample_counts_kwargs,
+                        )
+
     # Plot effect sizes and bootstraps.
     plot_groups = (temp_all_plot_groups if (is_paired == "baseline" and show_pairs and two_col_sankey) 
                    else temp_idx if two_col_sankey 

nbs/API/effsize_objects.ipynb

@@ -1196,6 +1196,9 @@
     "\n",
     "        es_marker_kwargs=None,\n",
     "        es_errorbar_kwargs=None,\n",
+    "\n",
+    "        prop_sample_counts=False,\n",
+    "        prop_sample_counts_kwargs=None\n",
     "    ):\n",
     "        \"\"\"\n",
     "        Creates an estimation plot for the effect size of interest.\n",
@@ -1285,6 +1288,9 @@
     "            Pass any keyword arguments accepted by the seaborn `swarmplot`\n",
     "            command here, as a dict. If None, the following keywords are\n",
     "            passed to sns.swarmplot : {'size':`raw_marker_size`}.\n",
+    "        barplot_kwargs : dict, default None\n",
+    "            By default, the keyword arguments passed are:\n",
+    "            {\"estimator\": np.mean, \"errorbar\": plot_kwargs[\"ci\"]}\n",
     "        violinplot_kwargs : dict, default None\n",
     "            Pass any keyword arguments accepted by the matplotlib `\n",
     "            pyplot.violinplot` command here, as a dict. If None, the following\n",
@@ -1400,6 +1406,13 @@
     "            Pass relevant keyword arguments to the effectsize errorbar plotting. If none, the following keywords are passed:\n",
     "            {'color': 'black', 'lw': 2, 'linestyle': '-', 'alpha': 1,'zorder': 1,}\n",
     "\n",
+    "        prop_sample_counts: bool, default False\n",
+    "            Show the sample counts for each group in proportional plots\n",
+    "        prop_sample_counts_kwargs: dict, default None\n",
+    "            Pass relevant keyword arguments. If None, the following keywords are passed:\n",
+    "            {'color': 'k', 'zorder': 5, 'ha': 'center', 'va': 'center'}\n",
+    "        \n",
+    "\n",
     "        Returns\n",
     "        -------\n",
     "        A :class:`matplotlib.figure.Figure` with 2 Axes, if ``ax = None``.\n",

nbs/API/misc_tools.ipynb

@@ -253,7 +253,7 @@
     "    # Barplot kwargs\n",
     "    default_barplot_kwargs = {\n",
     "        \"estimator\": np.mean, \n",
-    "        \"errorbar\": plot_kwargs[\"ci\"]\n",
+    "        \"errorbar\": plot_kwargs[\"ci\"],\n",
     "    }\n",
     "    if plot_kwargs[\"barplot_kwargs\"] is None:\n",
     "        barplot_kwargs = default_barplot_kwargs\n",
@@ -488,11 +488,23 @@
     "    else:\n",
     "        es_errorbar_kwargs = merge_two_dicts(default_es_errorbar_kwargs, plot_kwargs['es_errorbar_kwargs'])\n",
     "\n",
+    "    # Prop sample counts kwargs\n",
+    "    default_prop_sample_counts_kwargs = {\n",
+    "                                        'color': 'k', \n",
+    "                                        'zorder': 5, \n",
+    "                                        'ha': 'center', \n",
+    "                                        'va': 'center'\n",
+    "    }\n",
+    "    if plot_kwargs['prop_sample_counts_kwargs'] is None:\n",
+    "        prop_sample_counts_kwargs = default_prop_sample_counts_kwargs\n",
+    "    else:\n",
+    "        prop_sample_counts_kwargs = merge_two_dicts(default_prop_sample_counts_kwargs, plot_kwargs['prop_sample_counts_kwargs'])\n",
+    "\n",
     "    # Return the kwargs.\n",
     "    return (swarmplot_kwargs, barplot_kwargs, sankey_kwargs, violinplot_kwargs, slopegraph_kwargs, \n",
     "            reflines_kwargs, legend_kwargs, group_summaries_kwargs, redraw_axes_kwargs, delta_dot_kwargs,\n",
     "            delta_text_kwargs, summary_bars_kwargs, swarm_bars_kwargs, contrast_bars_kwargs, table_kwargs, gridkey_kwargs,\n",
-    "            es_marker_kwargs, es_errorbar_kwargs)\n",
+    "            es_marker_kwargs, es_errorbar_kwargs, prop_sample_counts_kwargs)\n",
     "\n",
     "\n",
     "def get_color_palette(plot_kwargs, plot_data, xvar, show_pairs, idx, all_plot_groups):\n",

nbs/API/plot_tools.ipynb

@@ -1969,9 +1969,8 @@
     "    horizontal : bool\n",
     "        If the plot is horizontal.\n",
     "    \"\"\"\n",
-    "    \n",
-    "    x_label, y_label = rawdata_axes.get_xlabel(), rawdata_axes.get_ylabel()\n",
     "\n",
+    "    x_label, y_label = rawdata_axes.get_xlabel(), rawdata_axes.get_ylabel()\n",
     "    if horizontal:\n",
     "        x_var, y_var, orient = np.ones(len(all_plot_groups)), all_plot_groups, \"h\"\n",
     "    else:\n",
@@ -2019,7 +2018,7 @@
     "            centre = x + width / 2.0\n",
     "            bar.set_x(centre - bar_width / 2.0)\n",
     "            bar.set_width(bar_width)\n",
-    "    \n",
+    "\n",
     "    # reset the x and y labels\n",
     "    rawdata_axes.set_xlabel(x_label)\n",
     "    rawdata_axes.set_ylabel(y_label)\n",
@@ -2095,7 +2094,62 @@
     "    ax.set_yticklabels([])\n",
     "    ax.tick_params(left=False, bottom=False)\n",
     "    ax.set_xlabel(label, fontsize=fontsize_label) # Set the x-axis label - hardcoded for now\n",
-    "    sns.despine(ax=ax, left=True, bottom=True)"
+    "    sns.despine(ax=ax, left=True, bottom=True)\n",
+    "\n",
+    "\n",
+    "def add_counts_to_prop_plots(plot_data, xvar, yvar, rawdata_axes, horizontal, is_paired, prop_sample_counts_kwargs):\n",
+    "    \"\"\"\n",
+    "    Add counts to the proportion plots.\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    plot_data : object (Dataframe)\n",
+    "        Dataframe of the plot data.\n",
+    "    xvar : str\n",
+    "        Column name of the x variable.\n",
+    "    yvar : str\n",
+    "        Column name of the y variable.\n",
+    "    rawdata_axes : object\n",
+    "        Matplotlib axis object to plot on.\n",
+    "    horizontal : bool\n",
+    "        If the plot is horizontal.\n",
+    "    is_paired : bool\n",
+    "        Whether the data is paired.\n",
+    "    prop_sample_counts_kwargs : dict\n",
+    "        Keyword arguments for the sample counts.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    # Group orders\n",
+    "    if isinstance(plot_data[xvar].dtype, pd.CategoricalDtype):\n",
+    "        sample_size_text_order = pd.unique(plot_data[xvar]).categories\n",
+    "    else:\n",
+    "        sample_size_text_order = pd.unique(plot_data[xvar])\n",
+    "\n",
+    "    # Get the sample size values\n",
+    "    ones, zeros = plot_data[plot_data[yvar] == 1], plot_data[plot_data[yvar] == 0]\n",
+    "\n",
+    "    sample_size_val1 = ones.groupby(xvar, observed=False)[yvar].count().reindex(index=sample_size_text_order)\n",
+    "    sample_size_val0 = zeros.groupby(xvar, observed=False)[yvar].count().reindex(index=sample_size_text_order)\n",
+    "\n",
+    "    fontsize = 8 if horizontal else 10\n",
+    "    fontsize -= 2 if is_paired else 0\n",
+    "\n",
+    "    if \"fontsize\" not in prop_sample_counts_kwargs.keys():\n",
+    "        fontsize = 8 if horizontal else 10\n",
+    "        fontsize -= 2 if is_paired else 0\n",
+    "        prop_sample_counts_kwargs.update({'fontsize': fontsize})\n",
+    "\n",
+    "    for sample_text_x, sample_text_y0, sample_text_y1 in zip(\n",
+    "                                                            np.arange(0,len(sample_size_text_order)+1,1), \n",
+    "                                                            sample_size_val0,\n",
+    "                                                            sample_size_val1,\n",
+    "                                                            ):\n",
+    "        if horizontal:\n",
+    "            rawdata_axes.text(0.05, sample_text_x, sample_text_y1, **prop_sample_counts_kwargs)\n",
+    "            rawdata_axes.text(0.95, sample_text_x, sample_text_y0, **prop_sample_counts_kwargs)\n",
+    "        else:\n",
+    "            rawdata_axes.text(sample_text_x, 0.05, sample_text_y1, **prop_sample_counts_kwargs)\n",
+    "            rawdata_axes.text(sample_text_x, 0.95, sample_text_y0, **prop_sample_counts_kwargs)"
    ]
   },
   {