notebooks/summarize_expected_vs_actual.py.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "68b91993-e2fa-44be-bcef-68130c72491d",
   "metadata": {},
   "source": [
    "# Summarize actual versus expected counts\n",
    "Plot average per-site mutations for each nucleotide type across all clades (summed) and sites (averaged).\n",
    "\n",
    "Import Python modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33c0415-444b-47c5-801d-783c8c0f8f6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import altair as alt\n",
    "\n",
    "import numpy\n",
    "\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe263d47-d19a-450b-9b98-224ba2d2940b",
   "metadata": {},
   "source": [
    "Now get variables from `snakemake`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0855debd-3616-4572-8757-ebd54249fdd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"snakemake\" not in globals() and \"snakemake\" not in locals():\n",
    "    # variables set manually for interactive debugging\n",
    "    counts_csv = \"../results/expected_vs_actual_mut_counts/expected_vs_actual_mut_counts.csv\"\n",
    "    chartfile = \"../results/expected_vs_actual_mut_counts/avg_counts.html\"\n",
    "    \n",
    "else:\n",
    "    # get variables from `snakemake` when running pipeline\n",
    "    counts_csv = snakemake.input.csv\n",
    "    chartfile = snakemake.output.chart"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9301be24-0dab-499d-8bd1-86932920de9f",
   "metadata": {},
   "source": [
    "Get average per-site expected, synonymous, nonsynonymous, and stop counts across all clades and sites.\n",
    "We ignore noncoding sites:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02afc21c-046e-4e95-abee-8fa63520417c",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = pd.read_csv(counts_csv, low_memory=False).query(\"not noncoding\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fefec82c-0aa4-4bd5-ac15-76cfe12088d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_counts = (\n",
    "    counts\n",
    "    .query(\"subset == 'all'\")\n",
    "    .query(\"not exclude\")\n",
    "    .assign(\n",
    "        mut_type=lambda x: numpy.where(\n",
    "            x[\"clade_founder_aa\"] == x[\"mutant_aa\"],\n",
    "            \"synonymous\",\n",
    "            numpy.where(x[\"mutant_aa\"] == \"*\", \"stop\", \"nonsynonymous\")\n",
    "        ),\n",
    "        mut=lambda x: x[\"nt_mutation\"].str[0] + \" to \" + x[\"nt_mutation\"].str[-1],\n",
    "    )\n",
    "    .groupby([\"mut_type\", \"mut\"], as_index=False)\n",
    "    .aggregate(\n",
    "        total_expected_count=pd.NamedAgg(\"expected_count\", \"sum\"),\n",
    "        total_actual_count=pd.NamedAgg(\"actual_count\", \"sum\"),\n",
    "        n_mutations=pd.NamedAgg(\"nt_mutation\", \"nunique\"),\n",
    "    )\n",
    "    .assign(\n",
    "        expected=lambda x: x[\"total_expected_count\"] / x[\"n_mutations\"],\n",
    "        actual=lambda x: x[\"total_actual_count\"] / x[\"n_mutations\"],\n",
    "    )\n",
    "    .melt(\n",
    "        id_vars=[\"mut_type\", \"mut\"],\n",
    "        value_vars=[\"expected\", \"actual\"],\n",
    "        var_name=\"count_type\",\n",
    "        value_name=\"counts\",\n",
    "    )\n",
    "    .assign(\n",
    "        mut_type=lambda x: numpy.where(\n",
    "            x[\"count_type\"] == \"actual\", x[\"mut_type\"], \"expected (4-fold degenerate)\",\n",
    "        ),\n",
    "    )\n",
    "    .groupby([\"mut_type\", \"mut\"], as_index=False)\n",
    "    .aggregate({\"counts\": \"mean\"})\n",
    ")\n",
    "\n",
    "avg_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04904b8a-a7b8-4c0c-ae8e-3cb1e24f6aa2",
   "metadata": {},
   "outputs": [],
   "source": [
    "height = 120\n",
    "\n",
    "mut_type_order = [\"expected (4-fold degenerate)\", \"synonymous\", \"nonsynonymous\", \"stop\"]\n",
    "\n",
    "mut_type_selection = alt.selection_multi(fields=[\"mut_type\"], bind=\"legend\")\n",
    "\n",
    "avg_count_chart = (\n",
    "    alt.Chart(avg_counts)\n",
    "    .encode(\n",
    "        x=alt.X(\n",
    "            \"mut_type\",\n",
    "            title=None,\n",
    "            axis=alt.Axis(labels=False, ticks=True, values=[\"synonymous\"]),\n",
    "            scale=alt.Scale(\n",
    "                domain=[\"dummy\", *mut_type_order],\n",
    "            ),\n",
    "        ),\n",
    "        y=alt.Y(\"counts\", title=\"average counts per site\"),\n",
    "        column=alt.Column(\n",
    "            \"mut\",\n",
    "            title=None,\n",
    "            header=alt.Header(\n",
    "                labelOrient=\"bottom\",\n",
    "                labelAngle=-90,\n",
    "                labelAlign=\"right\",\n",
    "                labelBaseline=\"middle\",\n",
    "                labelPadding=height + 8,\n",
    "            ),\n",
    "            spacing=0,\n",
    "        ),\n",
    "        color=alt.Color(\n",
    "            \"mut_type\",\n",
    "            title=\"type of count\",\n",
    "            sort=mut_type_order,\n",
    "            scale=alt.Scale(\n",
    "                range=[\"#999999\", \"#009E73\", \"#56B4E9\", \"#E69F00\"],\n",
    "            ),\n",
    "            legend=alt.Legend(orient=\"bottom\", titleOrient=\"left\"),\n",
    "        ),\n",
    "        tooltip=[\n",
    "            alt.Tooltip(\"mut_type\", title=\"type of count\"),\n",
    "            alt.Tooltip(\"mut\", title=\"mutation type\"),\n",
    "            alt.Tooltip(\"counts\", title=\"average counts\", format=\".2f\"),\n",
    "        ],\n",
    "        opacity=alt.condition(mut_type_selection, alt.value(1), alt.value(0.25)),\n",
    "    )\n",
    "    .mark_bar()\n",
    "    .configure_axis(grid=False)\n",
    "    .configure_view(stroke=None)\n",
    "    .properties(height=height, width=38)\n",
    "    .add_selection(mut_type_selection)\n",
    ")\n",
    "\n",
    "avg_count_chart.save(chartfile)\n",
    "\n",
    "avg_count_chart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1cd9828-2030-4cde-8f55-0cafdcfec22e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}