{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7a10993",
   "metadata": {},
   "source": [
    "# Perturbation-space heatmaps\n",
    "\n",
    "This notebook illustrates how to use `demovuln` outputs to visualize population reduction across slices of the perturbation space.\n",
    "\n",
    "The example uses a simple three-stage model with one juvenile stage, one subadult stage, and one adult stage. The model is intentionally minimal, but the same workflow applies to empirical projection matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb6e0f4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:12.138677Z",
     "iopub.status.busy": "2026-05-19T08:06:12.138374Z",
     "iopub.status.idle": "2026-05-19T08:06:12.628443Z",
     "shell.execute_reply": "2026-05-19T08:06:12.626933Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from demovuln import MatrixPopulationModel, PerturbationGrid, run_grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf4006c3",
   "metadata": {},
   "source": [
    "## Define a three-stage projection matrix\n",
    "\n",
    "The model contains one fecundity row and explicit stage definitions. Here, the adult source stage is the third column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57445627",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:12.631571Z",
     "iopub.status.busy": "2026-05-19T08:06:12.631152Z",
     "iopub.status.idle": "2026-05-19T08:06:12.638960Z",
     "shell.execute_reply": "2026-05-19T08:06:12.637605Z"
    }
   },
   "outputs": [],
   "source": [
    "A = np.array([\n",
    "    [0.0, 0.0, 1.2],\n",
    "    [0.35, 0.55, 0.0],\n",
    "    [0.0, 0.20, 0.85],\n",
    "])\n",
    "\n",
    "model = MatrixPopulationModel(\n",
    "    A,\n",
    "    adult_stages=[2],\n",
    "    juvenile_stages=[0, 1],\n",
    "    name=\"three_stage_example\",\n",
    ")\n",
    "\n",
    "print(f\"Dominant eigenvalue: {model.lambda_:.3f}\")\n",
    "print(\"Stable stage distribution:\", model.stable_distribution())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ad3a0f3",
   "metadata": {},
   "source": [
    "## Helper function for heatmaps\n",
    "\n",
    "The table returned by `run_grid` is transformed into a two-dimensional matrix using `pandas.pivot_table`. The plotting function below keeps each heatmap in its own figure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "531efd8e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:12.641801Z",
     "iopub.status.busy": "2026-05-19T08:06:12.641489Z",
     "iopub.status.idle": "2026-05-19T08:06:12.648049Z",
     "shell.execute_reply": "2026-05-19T08:06:12.646332Z"
    }
   },
   "outputs": [],
   "source": [
    "def plot_heatmap(table, *, x, y, value=\"population_reduction\", title=None):\n",
    "    pivot = table.pivot_table(index=y, columns=x, values=value, aggfunc=\"mean\")\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(7, 4.5))\n",
    "    image = ax.imshow(\n",
    "        pivot.values,\n",
    "        origin=\"lower\",\n",
    "        aspect=\"auto\",\n",
    "        extent=[\n",
    "            pivot.columns.min(),\n",
    "            pivot.columns.max(),\n",
    "            pivot.index.min(),\n",
    "            pivot.index.max(),\n",
    "        ],\n",
    "    )\n",
    "    ax.set_xlabel(x)\n",
    "    ax.set_ylabel(y)\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    cbar = fig.colorbar(image, ax=ax)\n",
    "    cbar.set_label(\"Population reduction (%)\")\n",
    "    fig.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f4b36df",
   "metadata": {},
   "source": [
    "## Magnitude-duration slices at fixed period\n",
    "\n",
    "We first fix the perturbation period and vary magnitude and duration. This is analogous to taking a two-dimensional slice through the full perturbation space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec002e09",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:12.651234Z",
     "iopub.status.busy": "2026-05-19T08:06:12.650914Z",
     "iopub.status.idle": "2026-05-19T08:06:12.944213Z",
     "shell.execute_reply": "2026-05-19T08:06:12.942467Z"
    }
   },
   "outputs": [],
   "source": [
    "fixed_period_grid = PerturbationGrid(\n",
    "    magnitudes=np.linspace(0, 0.8, 17),\n",
    "    durations=range(0, 7),\n",
    "    periods=[8],\n",
    ")\n",
    "\n",
    "fixed_period_results = {}\n",
    "\n",
    "for target in [\"adult_survival\", \"juvenile_survival\", \"fecundity\"]:\n",
    "    fixed_period_results[target] = run_grid(\n",
    "        model,\n",
    "        target=target,\n",
    "        grid=fixed_period_grid,\n",
    "        t_max=48,\n",
    "        recovery_steps=12,\n",
    "    )\n",
    "    print(f\"{target}: Φ = {fixed_period_results[target].vulnerability:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "905bbb1a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:12.947149Z",
     "iopub.status.busy": "2026-05-19T08:06:12.946855Z",
     "iopub.status.idle": "2026-05-19T08:06:13.609676Z",
     "shell.execute_reply": "2026-05-19T08:06:13.608238Z"
    }
   },
   "outputs": [],
   "source": [
    "for target, result in fixed_period_results.items():\n",
    "    plot_heatmap(\n",
    "        result.table,\n",
    "        x=\"magnitude\",\n",
    "        y=\"duration\",\n",
    "        title=f\"{target}: magnitude-duration slice at period = 8\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef8fa820",
   "metadata": {},
   "source": [
    "## Duration-period slices at fixed magnitude\n",
    "\n",
    "We can also fix perturbation magnitude and examine how population reduction depends on duration and recurrence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3675cd15",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:13.612596Z",
     "iopub.status.busy": "2026-05-19T08:06:13.612270Z",
     "iopub.status.idle": "2026-05-19T08:06:13.870988Z",
     "shell.execute_reply": "2026-05-19T08:06:13.869670Z"
    }
   },
   "outputs": [],
   "source": [
    "fixed_magnitude = 0.25\n",
    "duration_period_grid = PerturbationGrid(\n",
    "    magnitudes=[fixed_magnitude],\n",
    "    durations=range(0, 9),\n",
    "    periods=range(1, 13),\n",
    ")\n",
    "\n",
    "adult_duration_period = run_grid(\n",
    "    model,\n",
    "    target=\"adult_survival\",\n",
    "    grid=duration_period_grid,\n",
    "    t_max=60,\n",
    "    recovery_steps=12,\n",
    "    skip_infeasible=True,\n",
    ")\n",
    "\n",
    "plot_heatmap(\n",
    "    adult_duration_period.table,\n",
    "    x=\"period\",\n",
    "    y=\"duration\",\n",
    "    title=f\"Adult survival: duration-period slice at magnitude = {fixed_magnitude}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a63dcadb",
   "metadata": {},
   "source": [
    "## Integrated vulnerability by demographic target\n",
    "\n",
    "Finally, we compute integrated vulnerability for several demographic targets over the same perturbation grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a7653a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:13.873649Z",
     "iopub.status.busy": "2026-05-19T08:06:13.873403Z",
     "iopub.status.idle": "2026-05-19T08:06:17.157302Z",
     "shell.execute_reply": "2026-05-19T08:06:17.155685Z"
    }
   },
   "outputs": [],
   "source": [
    "full_grid = PerturbationGrid(\n",
    "    magnitudes=np.linspace(0, 0.8, 17),\n",
    "    durations=range(0, 7),\n",
    "    periods=range(1, 13),\n",
    ")\n",
    "\n",
    "target_summary = []\n",
    "\n",
    "for target in [\"adult_survival\", \"juvenile_survival\", \"fecundity\", \"all\"]:\n",
    "    result = run_grid(\n",
    "        model,\n",
    "        target=target,\n",
    "        grid=full_grid,\n",
    "        t_max=60,\n",
    "        recovery_steps=12,\n",
    "    )\n",
    "    target_summary.append({\"target\": target, \"vulnerability\": result.vulnerability})\n",
    "\n",
    "target_summary = pd.DataFrame(target_summary)\n",
    "target_summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "683f38a4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-19T08:06:17.160054Z",
     "iopub.status.busy": "2026-05-19T08:06:17.159762Z",
     "iopub.status.idle": "2026-05-19T08:06:17.283584Z",
     "shell.execute_reply": "2026-05-19T08:06:17.281962Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(7, 4))\n",
    "ax.bar(target_summary[\"target\"], target_summary[\"vulnerability\"])\n",
    "ax.set_xlabel(\"Perturbed demographic target\")\n",
    "ax.set_ylabel(\"Integrated vulnerability Φ (%)\")\n",
    "ax.set_title(\"Integrated vulnerability across perturbation targets\")\n",
    "ax.tick_params(axis=\"x\", rotation=30)\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5a41f89",
   "metadata": {},
   "source": [
    "## Notes\n",
    "\n",
    "These examples are deliberately small enough to run quickly. For empirical analyses, the same functions can be used with larger projection matrices and denser perturbation grids. The computational cost scales with the number of simulated perturbation regimes and the projection horizon."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyPI",
   "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.14.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}