nbs/API/forest_plot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Forest plot\n",
    "\n",
    "> Creating forest plots from contrast objects.\n",
    "\n",
    "- order: 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp forest_plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from __future__ import annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import *\n",
    "import nbdev\n",
    "nbdev.nbdev_export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import dabest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "import seaborn as sns\n",
    "from typing import List, Optional, Union\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def load_plot_data(\n",
    "    contrasts: List, effect_size: str = \"mean_diff\", contrast_type: str = \"delta2\"\n",
    ") -> List:\n",
    "    \"\"\"\n",
    "    Loads plot data based on specified effect size and contrast type.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    contrasts : List\n",
    "        List of contrast objects.\n",
    "    effect_size: str\n",
    "        Type of effect size ('mean_diff', 'median_diff', etc.).\n",
    "    contrast_type: str\n",
    "        Type of contrast ('delta2', 'mini_meta').\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    List: Contrast plot data based on specified parameters.\n",
    "    \"\"\"\n",
    "    effect_attr_map = {\n",
    "        \"mean_diff\": \"mean_diff\",\n",
    "        \"median_diff\": \"median_diff\",\n",
    "        \"cliffs_delta\": \"cliffs_delta\",\n",
    "        \"cohens_d\": \"cohens_d\",\n",
    "        \"hedges_g\": \"hedges_g\",\n",
    "        \"delta_g\": \"delta_g\"\n",
    "    }\n",
    "\n",
    "    contrast_attr_map = {\"delta2\": \"delta_delta\", \"mini_meta\": \"mini_meta_delta\"}\n",
    "\n",
    "    effect_attr = effect_attr_map.get(effect_size)\n",
    "    contrast_attr = contrast_attr_map.get(contrast_type)\n",
    "\n",
    "    if not effect_attr:\n",
    "        raise ValueError(f\"Invalid effect_size: {effect_size}\") \n",
    "    if not contrast_attr:\n",
    "        raise ValueError(f\"Invalid contrast_type: {contrast_type}. Available options: [`delta2`, `mini_meta`]\")\n",
    "\n",
    "    return [\n",
    "        getattr(getattr(contrast, effect_attr), contrast_attr) for contrast in contrasts\n",
    "    ]\n",
    "\n",
    "\n",
    "def extract_plot_data(contrast_plot_data, contrast_type):\n",
    "    \"\"\"Extracts bootstrap, difference, and confidence intervals based on contrast labels.\"\"\"\n",
    "    if contrast_type == \"mini_meta\":\n",
    "        attribute_suffix = \"weighted_delta\"\n",
    "    else:\n",
    "        attribute_suffix = \"delta_delta\"\n",
    "\n",
    "    bootstraps = [\n",
    "        getattr(result, f\"bootstraps_{attribute_suffix}\")\n",
    "        for result in contrast_plot_data\n",
    "    ]\n",
    "    \n",
    "    differences = [result.difference for result in contrast_plot_data]\n",
    "    bcalows = [result.bca_low for result in contrast_plot_data]\n",
    "    bcahighs = [result.bca_high for result in contrast_plot_data]\n",
    "    \n",
    "    return bootstraps, differences, bcalows, bcahighs\n",
    "\n",
    "\n",
    "def forest_plot(\n",
    "    contrasts: List,\n",
    "    selected_indices: Optional[List] = None,\n",
    "    contrast_type: str = \"delta2\",\n",
    "    xticklabels: Optional[List] = None,\n",
    "    effect_size: str = \"mean_diff\",\n",
    "    contrast_labels: List[str] = None,\n",
    "    ylabel: str = \"value\",\n",
    "    plot_elements_to_extract: Optional[List] = None,\n",
    "    title: str = \"ΔΔ Forest\",\n",
    "    custom_palette: Optional[Union[dict, list, str]] = None,\n",
    "    fontsize: int = 20,\n",
    "    violin_kwargs: Optional[dict] = None,\n",
    "    marker_size: int = 20,\n",
    "    ci_line_width: float = 2.5,\n",
    "    zero_line_width: int = 1,\n",
    "    remove_spines: bool = True,\n",
    "    ax: Optional[plt.Axes] = None,\n",
    "    additional_plotting_kwargs: Optional[dict] = None,\n",
    "    rotation_for_xlabels: int = 45,\n",
    "    alpha_violin_plot: float = 0.4,\n",
    "    horizontal: bool = False  # New argument for horizontal orientation\n",
    ")-> plt.Figure:\n",
    "    \"\"\"  \n",
    "    Custom function that generates a forest plot from given contrast objects, suitable for a range of data analysis types, including those from packages like DABEST-python.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    contrasts : List\n",
    "        List of contrast objects.\n",
    "    selected_indices : Optional[List], default=None\n",
    "        Indices of specific contrasts to plot, if not plotting all.\n",
    "    analysis_type : str\n",
    "         the type of analysis (e.g., 'delta2', 'minimeta').\n",
    "    xticklabels : Optional[List], default=None\n",
    "        Custom labels for the x-axis ticks.\n",
    "    effect_size : str\n",
    "        Type of effect size to plot (e.g., 'mean_diff', 'median_diff').\n",
    "    contrast_labels : List[str]\n",
    "        Labels for each contrast.\n",
    "    ylabel : str\n",
    "        Label for the y-axis, describing the plotted data or effect size.\n",
    "    plot_elements_to_extract : Optional[List], default=None\n",
    "        Elements to extract for detailed plot customization.\n",
    "    title : str\n",
    "        Plot title, summarizing the visualized data.\n",
    "    ylim : Tuple[float, float]\n",
    "        Limits for the y-axis.\n",
    "    custom_palette : Optional[Union[dict, list, str]], default=None\n",
    "        Custom color palette for the plot.\n",
    "    fontsize : int\n",
    "        Font size for text elements in the plot.\n",
    "    violin_kwargs : Optional[dict], default=None\n",
    "        Additional arguments for violin plot customization.\n",
    "    marker_size : int\n",
    "        Marker size for plotting mean differences or effect sizes.\n",
    "    ci_line_width : float\n",
    "        Width of confidence interval lines.\n",
    "    zero_line_width : int\n",
    "        Width of the line indicating zero effect size.\n",
    "    remove_spines : bool, default=False\n",
    "        If True, removes top and right plot spines.\n",
    "    ax : Optional[plt.Axes], default=None\n",
    "        Matplotlib Axes object for the plot; creates new if None.\n",
    "        additional_plotting_kwargs : Optional[dict], default=None\n",
    "        Further customization arguments for the plot.\n",
    "    rotation_for_xlabels : int, default=0\n",
    "        Rotation angle for x-axis labels, improving readability.\n",
    "    alpha_violin_plot : float, default=1.0\n",
    "        Transparency level for violin plots.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    plt.Figure\n",
    "        The matplotlib figure object with the generated forest plot.\n",
    "    \"\"\"\n",
    "    from .plot_tools import halfviolin\n",
    "\n",
    "    # Validate inputs\n",
    "    if contrasts is None:\n",
    "        raise ValueError(\"The `contrasts` parameter cannot be None\")\n",
    "    \n",
    "    if not isinstance(contrasts, list) or not contrasts:\n",
    "        raise ValueError(\"The `contrasts` argument must be a non-empty list.\")\n",
    "    \n",
    "    if selected_indices is not None and not isinstance(selected_indices, (list, type(None))):\n",
    "        raise TypeError(\"The `selected_indices` must be a list of integers or `None`.\")\n",
    "    \n",
    "    if not isinstance(contrast_type, str):\n",
    "        raise TypeError(\"The `contrast_type` argument must be a string.\")\n",
    "    \n",
    "    if xticklabels is not None and not all(isinstance(label, str) for label in xticklabels):\n",
    "        raise TypeError(\"The `xticklabels` must be a list of strings or `None`.\")\n",
    "    \n",
    "    if not isinstance(effect_size, str):\n",
    "        raise TypeError(\"The `effect_size` argument must be a string.\")\n",
    "    \n",
    "    if contrast_labels is not None and not all(isinstance(label, str) for label in contrast_labels):\n",
    "        raise TypeError(\"The `contrast_labels` must be a list of strings or `None`.\")\n",
    "    \n",
    "    if contrast_labels is not None and len(contrast_labels) != len(contrasts):\n",
    "        raise ValueError(\"`contrast_labels` must match the number of `contrasts` if provided.\")\n",
    "    \n",
    "    if not isinstance(ylabel, str):\n",
    "        raise TypeError(\"The `ylabel` argument must be a string.\")\n",
    "    \n",
    "    if custom_palette is not None and not isinstance(custom_palette, (dict, list, str, type(None))):\n",
    "        raise TypeError(\"The `custom_palette` must be either a dictionary, list, string, or `None`.\")\n",
    "    \n",
    "    if not isinstance(fontsize, (int, float)):\n",
    "        raise TypeError(\"`fontsize` must be an integer or float.\")\n",
    "    \n",
    "    if not isinstance(marker_size, (int, float)) or marker_size <= 0:\n",
    "        raise TypeError(\"`marker_size` must be a positive integer or float.\")\n",
    "    \n",
    "    if not isinstance(ci_line_width, (int, float)) or ci_line_width <= 0:\n",
    "        raise TypeError(\"`ci_line_width` must be a positive integer or float.\")\n",
    "    \n",
    "    if not isinstance(zero_line_width, (int, float)) or zero_line_width <= 0:\n",
    "        raise TypeError(\"`zero_line_width` must be a positive integer or float.\")\n",
    "    \n",
    "    if not isinstance(remove_spines, bool):\n",
    "        raise TypeError(\"`remove_spines` must be a boolean value.\")\n",
    "    \n",
    "    if ax is not None and not isinstance(ax, plt.Axes):\n",
    "        raise TypeError(\"`ax` must be a `matplotlib.axes.Axes` instance or `None`.\")\n",
    "    \n",
    "    if not isinstance(rotation_for_xlabels, (int, float)) or not 0 <= rotation_for_xlabels <= 360:\n",
    "        raise TypeError(\"`rotation_for_xlabels` must be an integer or float between 0 and 360.\")\n",
    "    \n",
    "    if not isinstance(alpha_violin_plot, float) or not 0 <= alpha_violin_plot <= 1:\n",
    "        raise TypeError(\"`alpha_violin_plot` must be a float between 0 and 1.\")\n",
    "    \n",
    "    if not isinstance(horizontal, bool):\n",
    "        raise TypeError(\"`horizontal` must be a boolean value.\")\n",
    "\n",
    "    # Load plot data\n",
    "    contrast_plot_data = load_plot_data(contrasts, effect_size, contrast_type)\n",
    "\n",
    "    # Extract data for plotting\n",
    "    bootstraps, differences, bcalows, bcahighs = extract_plot_data(\n",
    "        contrast_plot_data, contrast_type\n",
    "    )\n",
    "    # Adjust figure size based on orientation\n",
    "    all_groups_count = len(contrasts)\n",
    "    if horizontal:\n",
    "        fig_size = (4, 1.5 * all_groups_count)\n",
    "    else:\n",
    "        fig_size = (1.5 * all_groups_count, 4)\n",
    "\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(figsize=fig_size)\n",
    "    else:\n",
    "        fig = ax.figure\n",
    "\n",
    "    # Adjust violin plot orientation based on the 'horizontal' argument\n",
    "    violin_kwargs = violin_kwargs or {\n",
    "        \"widths\": 0.5,\n",
    "        \"showextrema\": False,\n",
    "        \"showmedians\": False,\n",
    "    }\n",
    "    violin_kwargs[\"vert\"] = not horizontal\n",
    "    v = ax.violinplot(bootstraps, **violin_kwargs)\n",
    "\n",
    "    # Adjust the halfviolin function call based on 'horizontal'\n",
    "    if horizontal:\n",
    "        half = \"top\"\n",
    "    else:\n",
    "        half = \"right\"  # Assuming \"right\" is the default or another appropriate value\n",
    "\n",
    "    # Assuming halfviolin has been updated to accept a 'half' parameter\n",
    "    halfviolin(v, alpha=alpha_violin_plot, half=half)\n",
    "    \n",
    "    # Handle the custom color palette\n",
    "    if custom_palette:\n",
    "        if isinstance(custom_palette, dict):\n",
    "            violin_colors = [\n",
    "                custom_palette.get(c, sns.color_palette()[0]) for c in contrasts\n",
    "            ]\n",
    "        elif isinstance(custom_palette, list):\n",
    "            violin_colors = custom_palette[: len(contrasts)]\n",
    "        elif isinstance(custom_palette, str):\n",
    "            if custom_palette in plt.colormaps():\n",
    "                violin_colors = sns.color_palette(custom_palette, len(contrasts))\n",
    "            else:\n",
    "                raise ValueError(\n",
    "                    f\"The specified `custom_palette` {custom_palette} is not a recognized Matplotlib palette.\"\n",
    "                )\n",
    "    else:\n",
    "        violin_colors = sns.color_palette()[: len(contrasts)]\n",
    "\n",
    "    for patch, color in zip(v[\"bodies\"], violin_colors):\n",
    "        patch.set_facecolor(color)\n",
    "        patch.set_alpha(alpha_violin_plot)\n",
    "\n",
    "    # Flipping the axes for plotting based on 'horizontal'\n",
    "    for k in range(1, len(contrasts) + 1):\n",
    "        if horizontal:\n",
    "            ax.plot(differences[k - 1], k, \"k.\", markersize=marker_size)  # Flipped axes\n",
    "            ax.plot([bcalows[k - 1], bcahighs[k - 1]], [k, k], \"k\", linewidth=ci_line_width)  # Flipped axes\n",
    "        else:\n",
    "            ax.plot(k, differences[k - 1], \"k.\", markersize=marker_size)\n",
    "            ax.plot([k, k], [bcalows[k - 1], bcahighs[k - 1]], \"k\", linewidth=ci_line_width)\n",
    "\n",
    "    # Adjusting labels, ticks, and limits based on 'horizontal'\n",
    "    if horizontal:\n",
    "        ax.set_yticks(range(1, len(contrasts) + 1))\n",
    "        ax.set_yticklabels(contrast_labels,  rotation=rotation_for_xlabels, fontsize=fontsize)\n",
    "        ax.set_xlabel(ylabel, fontsize=fontsize)\n",
    "    else:\n",
    "        ax.set_xticks(range(1, len(contrasts) + 1))\n",
    "        ax.set_xticklabels(contrast_labels, rotation=rotation_for_xlabels, fontsize=fontsize)\n",
    "        ax.set_ylabel(ylabel, fontsize=fontsize)\n",
    "\n",
    "    # Setting the title and adjusting spines as before\n",
    "    ax.set_title(title, fontsize=fontsize)\n",
    "    if remove_spines:\n",
    "        for spine in ax.spines.values():\n",
    "            spine.set_visible(False)\n",
    "\n",
    "    # Apply additional customizations if provided\n",
    "    if additional_plotting_kwargs:\n",
    "        ax.set(**additional_plotting_kwargs)\n",
    "\n",
    "    return fig"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}