{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bf0ddd61-d8e3-4a6a-bd9e-b5fc38e32e09",
   "metadata": {},
   "source": [
    "# Doughnuts\n",
    "Pie charts with a doughnut shape where the thickness of the doughnut indicates the total\n",
    "number of points in the pie chart relative to a reference total. Based on Figs. 1/2 in\n",
    "[Roberts et al. (2020) - Impact of Model Resolution on Tropical Cyclone Simulation Using the HighResMIP–PRIMAVERA Multimodel Ensemble](https://doi.org/10.1175/JCLI-D-19-0639.1)\n",
    "\n",
    "## Default Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abfc2ce2-d86d-4de6-88a8-215887cb185a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import huracanpy\n",
    "\n",
    "fig, axes = plt.subplots(2, 2)\n",
    "\n",
    "# Default doughnut\n",
    "# thickness = 2/3 of the pie\n",
    "huracanpy.plot.doughnut([1, 2], 3, ax=axes[0, 0])\n",
    "\n",
    "# Less data than reference total - thinner doughnut\n",
    "huracanpy.plot.doughnut([1, 2], 6, ax=axes[0, 1])\n",
    "\n",
    "# More data than reference total - thicker doughnut\n",
    "huracanpy.plot.doughnut([1, 2], 1.5, ax=axes[1, 0])\n",
    "\n",
    "# Much more data\n",
    "huracanpy.plot.doughnut([1, 2], 0.1, ax=axes[1, 1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd196ad-c1e4-427e-9bf2-dc7feae99824",
   "metadata": {},
   "source": [
    "## IBTrACS Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd7705b6-72a2-4367-8190-170e02bcfa24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "ibtracs = huracanpy.load(source=\"ibtracs\", ibtracs_online=False)\n",
    "genesis_points = ibtracs.groupby(\"track_id\").first()\n",
    "basins, counts = np.unique(genesis_points.basin, return_counts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b47e8bed-b67d-4ec1-a4bb-c762b9499bdd",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "wedges, labels, percentage_labels = huracanpy.plot.doughnut(\n",
    "    counts, sum(counts), labels=basins, autopct=\"%.1f%%\"\n",
    ")\n",
    "\n",
    "# Set the centre text to the average number of storms per year\n",
    "# I'm not sure if the first and last year in the data are full years, so this may be an\n",
    "# underestimate, but I'll keep it simple for the example\n",
    "per_year = sum(counts) / len(set(ibtracs.time.dt.year.values))\n",
    "plt.text(0, 0, f\"{per_year:.1f}\", ha=\"center\", va=\"center\")\n",
    "\n",
    "plt.title(\"IBTrACS - Storm distibution by basin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48e78572-4849-4df4-893f-fc0e27f41b44",
   "metadata": {},
   "source": [
    "## Recreating the doughnuts from Roberts et al. (2020)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecdf385c-cf34-4a29-9b2d-a6450007729b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "\n",
    "# Create the input data as you might expect to produce it with huracanpy/xarray\n",
    "# A 2d array of average number of storms per year as a function of model and basin\n",
    "da = xr.DataArray(\n",
    "    data=[\n",
    "        np.array([16, 53, 19, 5, 7]) * 0.521,\n",
    "        np.array([17, 46, 23, 7, 7]) * 0.873,\n",
    "        np.array([16, 46, 16, 17, 5]) * 0.509,\n",
    "        np.array([17, 46, 18, 15, 5]) * 0.581,\n",
    "        np.array([13, 45, 20, 18, 4]) * 0.249,\n",
    "        np.array([14, 44, 20, 18, 4]) * 0.42,\n",
    "        np.array([22, 42, 12, 18, 5]) * 0.662,\n",
    "        np.array([23, 44, 9, 20, 4]) * 0.654,\n",
    "        np.array([21, 33, 19, 22, 5]) * 0.14,\n",
    "        np.array([19, 35, 19, 22, 4]) * 0.133,\n",
    "        np.array([11, 32, 21, 28, 8]) * 0.296,\n",
    "        np.array([14, 39, 21, 18, 8]) * 0.649,\n",
    "        np.array([17, 40, 28, 11, 4]) * 0.648,\n",
    "        np.array([19, 39, 26, 10, 6]) * 0.627,\n",
    "        np.array([18, 45, 20, 12, 5]) * 0.495,\n",
    "        np.array([18, 42, 23, 13, 4]) * 0.617,\n",
    "        np.array([21, 38, 24, 11, 6]) * 0.729,\n",
    "        np.array([21, 42, 28, 8, 1]) * 0.536,\n",
    "    ],\n",
    "    coords=dict(\n",
    "        model=[\n",
    "            \"HadGEM3-GC31-LM\",\n",
    "            \"HadGEM3-GC31-HM\",\n",
    "            \"ECMWF-IFS-LR\",\n",
    "            \"ECMWF-IFS-HR\",\n",
    "            \"EC-Earth3P-LR\",\n",
    "            \"EC-Earth3P-HR\",\n",
    "            \"CNRM-CM6-1\",\n",
    "            \"CNRM-CM6-1-HR\",\n",
    "            \"MPI-ESM1-2-HR\",\n",
    "            \"MPI-ESM1-2-XR\",\n",
    "            \"CMCC-CM2-HR4\",\n",
    "            \"CMCC-CM2-VHR4\",\n",
    "            \"MERRA2\",\n",
    "            \"JRA55\",\n",
    "            \"ERAI\",\n",
    "            \"ERA5\",\n",
    "            \"CFSR2\",\n",
    "            \"Obs\",\n",
    "        ],\n",
    "        basin=[\"na\", \"wp\", \"ep\", \"ni\", \"other\"],\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64a80fef-aea5-4a8e-a807-a8a9b3b0189d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from string import ascii_lowercase\n",
    "\n",
    "# Specific parameters to plt.pie used by Roberts et al. (2020). See\n",
    "# https://github.com/eerie-project/storm_track_analysis/blob/main/assess/tc_assessment.py#L384\n",
    "pie_kwargs = dict(\n",
    "    startangle=90,\n",
    "    autopct=\"%1.0f%%\",\n",
    "    pctdistance=0.7,\n",
    "    labels=da.basin.values,\n",
    "    labeldistance=1.0,\n",
    "    colors=[\"#ff6666\", \"#ffcc99\", \"#cc9966\", \"#cc6666\", \"#66b3ff\"],\n",
    ")\n",
    "# The second value in the centre is the number of southern hemisphere storms\n",
    "# This isn't related to the data in the doughnut so I've just put a list of values here\n",
    "sh_values = [\n",
    "    68.3,\n",
    "    95.0,\n",
    "    48.4,\n",
    "    53.9,\n",
    "    33.4,\n",
    "    39.3,\n",
    "    64.3,\n",
    "    60.1,\n",
    "    20.9,\n",
    "    20.3,\n",
    "    24.2,\n",
    "    48.5,\n",
    "    53.9,\n",
    "    48.3,\n",
    "    46.0,\n",
    "    51.5,\n",
    "    53.3,\n",
    "    20.8,\n",
    "]\n",
    "\n",
    "fig, axes = plt.subplots(3, 6, figsize=(20, 10))\n",
    "axes = axes.flatten()\n",
    "\n",
    "# Thickness of doughnuts relative to the \"Obs\" doughnut\n",
    "reference_total = da.sel(model=\"Obs\").values.sum()\n",
    "\n",
    "# One plot for each model. Loop over array per model\n",
    "for n, model in enumerate(da.model.values):\n",
    "    da_ = da.sel(model=model)\n",
    "\n",
    "    huracanpy.plot.doughnut(da_.values, reference_total, ax=axes[n], **pie_kwargs)\n",
    "    axes[n].text(\n",
    "        0, 0, f\"{da_.values.sum():.1f}\\n{sh_values[n]}\", ha=\"center\", va=\"center\"\n",
    "    )\n",
    "    axes[n].set_title(f\"({ascii_lowercase[n]}) {model}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}