{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "50d3669e",
   "metadata": {},
   "source": [
    "# Visualization gallery"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c01e724e-1d4e-41e7-bb76-742f354fd15f",
   "metadata": {},
   "source": [
    "The visualization methods implemented in PEtab Select are demonstrated here. These methods generally visualize the output of a model selection task, so the input is generally a list of already-calibrated models.\n",
    "\n",
    "All dependencies for these plots can be installed with `pip install petab_select[plot]`.\n",
    "\n",
    "In this notebook, some calibrated models that were saved to disk with the `to_yaml` method of a `Models` object, are loaded and used as input here. This is the result of a forward selection with the problem provided in `calibrated_models`. Note that a `Models` object is expect here; if you have a list of models `model_list`, simply use `models = Models(model_list)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca6ce5b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "\n",
    "import petab_select\n",
    "import petab_select.plot\n",
    "\n",
    "models = petab_select.Models.from_yaml(\n",
    "    \"calibrated_models/calibrated_models.yaml\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54532b75-53e4-4670-8e64-21e7adda0c0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "models.df.drop(\n",
    "    columns=[petab_select.Criterion.AIC, petab_select.Criterion.BIC]\n",
    ").style.background_gradient(\n",
    "    cmap=matplotlib.colormaps.get_cmap(\"summer\"),\n",
    "    subset=[petab_select.Criterion.AICC],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaeb0606",
   "metadata": {},
   "source": [
    "To use the plotting methods, we need to first setup an object that contains information common to multiple plotting methods. This can include the models, custom colors and labels, and the criterion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09c9df1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom colors for some models\n",
    "colors = {\n",
    "    \"M-000\": \"lightgreen\",\n",
    "    \"M-001\": \"lightgreen\",\n",
    "}\n",
    "\n",
    "plot_data = petab_select.plot.PlotData(\n",
    "    models=models,\n",
    "    criterion=petab_select.Criterion.AICC,\n",
    "    colors=colors,\n",
    "    relative_criterion=True,\n",
    ")\n",
    "\n",
    "# Change default color\n",
    "petab_select.plot.DEFAULT_NODE_COLOR = \"darkgray\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c73e0bc-5bf8-4c03-a54b-f19ced731322",
   "metadata": {},
   "source": [
    "## UpSet plot\n",
    "\n",
    "This shows models ordered by criterion, with their parameters directly below the bars.\n",
    "\n",
    "A black dot indicates that the parameter (e.g `k2`) is estimated in the model (e.g. the first bar is a model with `k1` and `sigma_x2` estimated)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d99572-f74d-4e25-8237-0aa158eb29f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "petab_select.plot.upset(plot_data=plot_data);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32de6556",
   "metadata": {},
   "source": [
    "## Best model from each iteration\n",
    "\n",
    "This shows strict improvements in the criterion, and the corresponding model, across all iterations of model selection.\n",
    "\n",
    "Since there were no improvements after `M_100`, no other iterations or models are shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56b4a27b",
   "metadata": {},
   "outputs": [],
   "source": [
    "petab_select.plot.line_best_by_iteration(plot_data=plot_data);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e5243e0",
   "metadata": {},
   "source": [
    "## Selection history trajectory\n",
    "\n",
    "This shows the relationship between models across iterations. For example, `M_000` was the predecessor model to `M_001`, `M_010`, and `M_100`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "862a78ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the relative criterion value to each label\n",
    "plot_data.augment_labels(criterion=True)\n",
    "petab_select.plot.graph_history(plot_data=plot_data)\n",
    "# Reset the labels (remove the relative criterion)\n",
    "plot_data.augment_labels()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d87e75ab",
   "metadata": {},
   "source": [
    "## Criterion values of each model\n",
    "\n",
    "This shows the criterion value of every calibrated model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bce41584",
   "metadata": {},
   "outputs": [],
   "source": [
    "petab_select.plot.bar_criterion_vs_models(plot_data=plot_data);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9c2d487",
   "metadata": {},
   "source": [
    "## Criterion values vs. number of estimated parameters\n",
    "\n",
    "This shows all calibrated models.\n",
    "\n",
    "In this example, models with 2 estimated parameters tend to perform best. This is also seen in the UpSet plot above.\n",
    "\n",
    "Jitter is added to distinguish models with the same number of parameters and similar criterion values, according to the optional `max_jitter` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824e2e6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "petab_select.plot.scatter_criterion_vs_n_estimated(\n",
    "    plot_data=plot_data,\n",
    "    # Uncomment to turn off jitter.\n",
    "    # max_jitter=0,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dc7e142",
   "metadata": {},
   "source": [
    "## History as layers in a hierarchical graph\n",
    "\n",
    "This shows the relative change in parameters of each model, compared to its predecessor model.\n",
    "\n",
    "Each column is an iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21157e4d-b2ba-4cb1-95f6-e14052c86959",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Customize the colors\n",
    "# criterion_values = [model.get_criterion(petab_select.Criterion.AICC) for model in models]\n",
    "# norm = matplotlib.colors.Normalize(\n",
    "#     vmin=min(criterion_values),\n",
    "#     vmax=max(criterion_values),\n",
    "# )\n",
    "# cmap = matplotlib.colors.LinearSegmentedColormap.from_list(\"\", [\"green\",\"lightgreen\"])\n",
    "# colorbar_mappable = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap)\n",
    "\n",
    "# Augment labels with the changes in parameters of each model, compared to their predecessor model\n",
    "plot_data.augment_labels(added_parameters=True, removed_parameters=True)\n",
    "\n",
    "petab_select.plot.graph_iteration_layers(\n",
    "    plot_data=plot_data,\n",
    "    draw_networkx_kwargs={\n",
    "        \"arrowstyle\": \"-|>\",\n",
    "        \"node_shape\": \"s\",\n",
    "        \"node_size\": 2500,\n",
    "        \"edgecolors\": \"k\",\n",
    "    },\n",
    "    # colorbar_mappable=colorbar_mappable,\n",
    ");"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}