notebooks/synonymous_mut_rates.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f53f9e29-d668-4f1e-83bb-5de1420f3953",
   "metadata": {},
   "source": [
    "# Mutation spectrum and rates at four-fold synonymous sites\n",
    "Analyze the mutation rates and spectrum only at four-fold synonymous sites, as defined in relation to the founder sequence for that clade."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eb11c2f-9f43-4ead-bfb1-43b6f7dd8b89",
   "metadata": {},
   "source": [
    "Import Python modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d349183-450d-415c-b85d-3b81f777ef5f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import os\n",
    "\n",
    "import altair as alt\n",
    "\n",
    "import numpy\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import yaml\n",
    "\n",
    "import scipy\n",
    "\n",
    "import sklearn.decomposition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e4d7926-7220-427d-a7ef-17f3a1e94d71",
   "metadata": {},
   "source": [
    "Get input variables from `snakemake`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8ff9999-26d6-449c-8a81-9ad39fde2d9d",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "if \"snakemake\" in locals() or \"snakemake\" in globals():\n",
    "    mutation_counts_csv = snakemake.input.mutation_counts_csv\n",
    "    clade_founder_nts_csv = snakemake.input.clade_founder_nts_csv\n",
    "    rates_by_clade_csv = snakemake.output.rates_by_clade\n",
    "    rates_plot = snakemake.output.rates_plot\n",
    "    synonymous_spectra_min_counts = snakemake.params.synonymous_spectra_min_counts\n",
    "    subset_order = list(snakemake.params.sample_subsets)\n",
    "    clade_synonyms = snakemake.params.clade_synonyms\n",
    "else:\n",
    "    # if running interactively\n",
    "    mutation_counts_csv = \"../results/mutation_counts/aggregated.csv\"\n",
    "    clade_founder_nts_csv = \"../results/clade_founder_nts/clade_founder_nts.csv\"\n",
    "    rates_by_clade_csv = \"../results/synonymous_mut_rates/rates_by_clade.csv\"\n",
    "    rates_plot = \"../results/synonymous_mut_rates/rates_plot.html\"\n",
    "    with open(\"config.yaml\" if os.path.isfile(\"config.yaml\") else \"../config.yaml\") as f:\n",
    "        config = yaml.safe_load(f)\n",
    "    synonymous_spectra_min_counts = config[\"synonymous_spectra_min_counts\"]\n",
    "    subset_order = list(config[\"sample_subsets\"])\n",
    "    clade_synonyms = config[\"clade_synonyms\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a85ab2bd-e4b7-4d7b-b780-df216bda2b51",
   "metadata": {},
   "source": [
    "Read the mutation counts and assign mutation types:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5178e432-6e87-43de-9bcb-c9d1984f81d7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mutation_counts = pd.read_csv(mutation_counts_csv).assign(\n",
    "    mut_type=lambda x: x[\"nt_mutation\"].map(lambda m: f\"{m[0]}to{m[-1]}\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "702373e8-da95-47d5-b073-cefdaebba5dc",
   "metadata": {},
   "source": [
    "For each clade plot the top mutations as a fraction of all mutations in that clade, just using the \"all\" subset.\n",
    "\n",
    "You can mouseover points to highlight mutations (which will highlight all mutations at that site on all facets), and click the legend to show/hide excluded or non-excluded mutations.\n",
    "\n",
    "This plot is useful to look at to identifier apparent outlier sites with aberrantly high mutation counts that can then be specified for exclusion (note those specifications are done in the pipeline `config.yaml` file, and also all reversions from clade founder to reference may be excluded):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "788dc870-b517-463e-9750-cb72b2eb15f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_n = 100  # plot this many per clade\n",
    "\n",
    "mutation_freqs = (\n",
    "    mutation_counts.query(\"subset == 'all'\")\n",
    "    .sort_values([\"clade\", \"count\"], ascending=False)\n",
    "    .groupby(\"clade\")\n",
    "    .head(n=top_n)\n",
    "    .assign(\n",
    "        freq=lambda x: x[\"count\"] / x.groupby(\"clade\")[\"count\"].transform(\"sum\"),\n",
    "        rank=lambda x: x.groupby(\"clade\")[\"freq\"].rank(ascending=False, method=\"first\"),\n",
    "        exclude=lambda x: x[\"exclude\"].map({True: \"yes\", False: \"no\"}),\n",
    "    )\n",
    ")\n",
    "\n",
    "select_exclude = alt.selection_multi(\n",
    "    fields=[\"exclude\"],\n",
    "    bind=\"legend\",\n",
    "    init=[{\"exclude\": \"yes\"}, {\"exclude\": \"no\"}],\n",
    ")\n",
    "\n",
    "select_site = alt.selection_single(\n",
    "    fields=[\"nt_site\"],\n",
    "    on=\"mouseover\",\n",
    "    empty=\"none\",\n",
    ")\n",
    "\n",
    "mutation_freqs_chart = (\n",
    "    alt.Chart(mutation_freqs)\n",
    "    .encode(\n",
    "        x=\"rank\",\n",
    "        y=\"freq\",\n",
    "        strokeWidth=alt.condition(select_site, alt.value(2), alt.value(0)),\n",
    "        color=alt.Color(\"exclude\", scale=alt.Scale(domain=[\"yes\", \"no\"])),\n",
    "        shape=alt.Shape(\"synonymous\"),\n",
    "        size=alt.condition(select_site, alt.value(50), alt.value(25)),\n",
    "        tooltip=[\"nt_site\", \"nt_mutation\", \"count\", \"freq\"],\n",
    "    )\n",
    "    .mark_point(filled=True, stroke=\"black\")\n",
    "    .properties(width=200, height=100)\n",
    "    .facet(\"clade\", columns=4)\n",
    "    .add_selection(select_exclude, select_site)\n",
    "    .transform_filter(select_exclude)\n",
    ")\n",
    "\n",
    "mutation_freqs_chart"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ee3f6bd-f62b-4e37-ad69-5a567d392804",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-09-08T20:58:54.989586Z",
     "iopub.status.busy": "2022-09-08T20:58:54.988909Z",
     "iopub.status.idle": "2022-09-08T20:58:54.997156Z",
     "shell.execute_reply": "2022-09-08T20:58:54.996168Z",
     "shell.execute_reply.started": "2022-09-08T20:58:54.989428Z"
    },
    "tags": []
   },
   "source": [
    "Tally mutation type counts among **only four-fold synonymous** mutations for each clade and subset, also removing any mutations specified for exclusion:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "244341c6-7add-45c7-91f6-d86c8f5cf684",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mut_type_counts = (\n",
    "    mutation_counts.query(\"synonymous\")\n",
    "    .query(\"four_fold_degenerate\")\n",
    "    .query(\"not exclude\")\n",
    "    .groupby([\"clade\", \"subset\", \"mut_type\"], as_index=False)\n",
    "    .aggregate({\"count\": \"sum\"})\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81167003-67ad-4f2c-835c-bcad7ac07ac3",
   "metadata": {},
   "source": [
    "Now also repeat these mutation type counts tally, but any mutations in the top 10 most frequent observed mutation for any clade, not doing any subsetting (just taking subset \"all\"):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "918c077e-fd9a-4371-a98f-8c7eb566e36d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "exclude_top_n = 5  # exclude mutations in this top rank for any clade\n",
    "\n",
    "mut_type_counts_exclude_top = (\n",
    "    mutation_counts.query(\"synonymous\")\n",
    "    .query(\"four_fold_degenerate\")\n",
    "    .query(\"not exclude\")\n",
    "    .query(\"subset == 'all'\")\n",
    "    .assign(\n",
    "        clade_rank=lambda x: x.groupby(\"clade\")[\"count\"].rank(\n",
    "            ascending=False, method=\"min\"\n",
    "        ),\n",
    "        highest_rank=lambda x: x.groupby(\"nt_mutation\")[\"clade_rank\"].transform(\"min\"),\n",
    "    )\n",
    "    .query(\"highest_rank > @exclude_top_n\")\n",
    "    .groupby([\"clade\", \"mut_type\"], as_index=False)\n",
    "    .aggregate({\"count\": \"sum\"})\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2778893d-ffe4-4571-9a8c-19d297cb2e87",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-09-08T21:05:52.421686Z",
     "iopub.status.busy": "2022-09-08T21:05:52.421174Z",
     "iopub.status.idle": "2022-09-08T21:05:52.432474Z",
     "shell.execute_reply": "2022-09-08T21:05:52.431513Z",
     "shell.execute_reply.started": "2022-09-08T21:05:52.421656Z"
    },
    "tags": []
   },
   "source": [
    "Plot total mutation counts for each clade and subset on a log scale.\n",
    "Also draw a line at our minimum cutoff: we only keep subsets above this cutoff:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7adad884-7f81-49a1-8d93-0db32fdd33d7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "clade_counts = mut_type_counts.groupby([\"clade\", \"subset\"], as_index=False).aggregate(\n",
    "    {\"count\": \"sum\"}\n",
    ")\n",
    "\n",
    "clade_counts_chart = (\n",
    "    alt.Chart(clade_counts)\n",
    "    .encode(\n",
    "        x=\"clade\",\n",
    "        y=alt.Y(\"count\", title=\"total mutations\", scale=alt.Scale(type=\"log\")),\n",
    "        tooltip=[\"clade\", \"subset\", \"count\"],\n",
    "        color=\"subset\",\n",
    "    )\n",
    "    .mark_circle(size=50, opacity=0.7)\n",
    "    .properties(width=alt.Step(18), height=175)\n",
    ")\n",
    "\n",
    "# draw cutoff line\n",
    "cutoff = (\n",
    "    alt.Chart(pd.DataFrame({\"y\": [synonymous_spectra_min_counts]}))\n",
    "    .encode(y=\"y\")\n",
    "    .mark_rule(strokeDash=[2, 2])\n",
    ")\n",
    "\n",
    "(clade_counts_chart + cutoff).configure_axis(grid=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e1c6a8f-a811-482a-ab05-054b46f0f87d",
   "metadata": {},
   "source": [
    "For genome partitioning, we subdivide the genome into halves based on the first and last site with an observed mutation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4722f3b7-a653-40e2-ac4b-95a4ccb9603c",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_partitions = 2\n",
    "\n",
    "min_site = mutation_counts[\"nt_site\"].min()\n",
    "max_site = mutation_counts[\"nt_site\"].max() + 1\n",
    "partition_bounds = numpy.linspace(min_site, max_site, n_partitions + 1)\n",
    "\n",
    "\n",
    "def assign_partition(r):\n",
    "    \"\"\"Assign nucleotide mutation to its partition.\"\"\"\n",
    "    for i in range(1, n_partitions + 1):\n",
    "        if partition_bounds[i - 1] <= r < partition_bounds[i]:\n",
    "            return f\"partition {i}\"\n",
    "\n",
    "\n",
    "mutation_counts = mutation_counts.assign(\n",
    "    partition=lambda x: x[\"nt_site\"].map(assign_partition)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb0195dd-de1d-4049-aded-994dbfd984a2",
   "metadata": {},
   "source": [
    "Get PCA of mutation spectrum, using only filtered synonymous mutation counts for non-excluded mutations for clades/subsets/partitions with adequate counts.\n",
    "\n",
    "We do the PCA on four different ways of partitioning the data:\n",
    "\n",
    " 1. Just looking at the \"all\" subset for each clade across entire genome.\n",
    " 2. Looking at the \"all\" subset for each clade across entire genome, but excluding any mutation that is among the top ranked most counts for any clade.\n",
    " 3. Looking at all subsets for each clade across entire genome.\n",
    " 4. Looking at the \"all\" subset along partitions of the genome. \n",
    "\n",
    "We standardize the vectors of mutation fractions before doing the PCA.\n",
    "\n",
    "There is also an option, currently not used, to do a composite log ratio (CLR) transform the fraction of all mutations that are each type before performing the PCA, using the formula [here](https://www.rdocumentation.org/packages/compositions/versions/2.0-4/topics/clr) (see [this paper](https://www.jstor.org/stable/2335943)).\n",
    "\n",
    "In the plots below, you can mouseover the points for details and click on clades in legends (shift click for multiple clades) to highlight just points for the selected clade(s).\n",
    "You can also use the scroll bar to only show points with at least the indicated number of total synonymous mutation counts (after filtering):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56a9263f-f805-4d97-8372-8952777adb41",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def clr(frac_vec):\n",
    "    \"\"\"https://www.rdocumentation.org/packages/compositions/versions/2.0-4/topics/clr\"\"\"\n",
    "    assert all(frac_vec > 0), \"does not currently handle zeros in frac_vec\"\n",
    "    return numpy.log(frac_vec) - 1 / len(frac_vec) * numpy.log(frac_vec).sum()\n",
    "\n",
    "\n",
    "for title, subsets, partition, exclude_top, clr_transform in [\n",
    "    (\"all samples, whole genome\", [\"all\"], False, False, False),\n",
    "    (\"all samples, whole genome, +/- top mutations\", [\"all\"], False, True, False),\n",
    "    (\"by region, whole genome\", subset_order, False, False, False),\n",
    "    (\"all samples, partitioned genome\", [\"all\"], True, False, False),\n",
    "]:\n",
    "\n",
    "    filtered_mutation_counts = (\n",
    "        mutation_counts.query(\"synonymous\")\n",
    "        .query(\"four_fold_degenerate\")\n",
    "        .query(\"not exclude\")\n",
    "        .query(\"subset in @subsets\")\n",
    "    )\n",
    "\n",
    "    if partition:\n",
    "        filtered_mutation_counts = pd.concat(\n",
    "            [\n",
    "                filtered_mutation_counts.assign(partition=\"all\"),\n",
    "                filtered_mutation_counts,\n",
    "            ]\n",
    "        )\n",
    "    else:\n",
    "        filtered_mutation_counts = filtered_mutation_counts.assign(partition=\"all\")\n",
    "\n",
    "    filtered_mut_type_counts = filtered_mutation_counts.groupby(\n",
    "        [\"clade\", \"subset\", \"partition\", \"mut_type\"], as_index=False\n",
    "    ).aggregate({\"count\": \"sum\"})\n",
    "\n",
    "    if exclude_top:\n",
    "        assert all(filtered_mut_type_counts[\"partition\"] == \"all\")\n",
    "        assert all(filtered_mut_type_counts[\"subset\"] == \"all\")\n",
    "        filtered_mut_type_counts = pd.concat(\n",
    "            [\n",
    "                filtered_mut_type_counts.assign(excluded=\"no\"),\n",
    "                mut_type_counts_exclude_top.assign(\n",
    "                    partition=\"all\",\n",
    "                    subset=\"all\",\n",
    "                    excluded=\"yes\",\n",
    "                ),\n",
    "            ]\n",
    "        )\n",
    "    else:\n",
    "        filtered_mut_type_counts = filtered_mut_type_counts.assign(excluded=\"no\")\n",
    "\n",
    "    mut_type_freqs = filtered_mut_type_counts.assign(\n",
    "        total_count=lambda x: (\n",
    "            x.groupby([\"clade\", \"subset\", \"partition\", \"excluded\"])[\"count\"].transform(\n",
    "                \"sum\"\n",
    "            )\n",
    "        ),\n",
    "        freq=lambda x: x[\"count\"] / x[\"total_count\"],\n",
    "    ).query(\"total_count >= @synonymous_spectra_min_counts\")\n",
    "\n",
    "    mut_type_freqs = mut_type_freqs.pivot_table(\n",
    "        index=[\"clade\", \"subset\", \"partition\", \"excluded\", \"total_count\"],\n",
    "        values=\"freq\",\n",
    "        columns=\"mut_type\",\n",
    "        fill_value=0,\n",
    "    )\n",
    "\n",
    "    if clr_transform:\n",
    "        display(mut_type_freqs)\n",
    "        mut_type_freqs = mut_type_freqs.apply(clr, axis=1)\n",
    "        display(mut_type_freqs)\n",
    "\n",
    "    scaled_freqs = sklearn.preprocessing.StandardScaler().fit_transform(\n",
    "        mut_type_freqs.values\n",
    "    )\n",
    "    pca = sklearn.decomposition.PCA(n_components=2)\n",
    "    pca_coords = pca.fit_transform(scaled_freqs)\n",
    "    assert len(pca_coords) == len(mut_type_freqs)\n",
    "\n",
    "    mut_type_freqs_pca_coords = mut_type_freqs.reset_index().assign(\n",
    "        principal_component_1=pca_coords[:, 0],\n",
    "        principal_component_2=pca_coords[:, 1],\n",
    "        log10_total_count=lambda x: numpy.log(x[\"total_count\"]) / numpy.log(10),\n",
    "    )\n",
    "\n",
    "    # percent variance explained by each component\n",
    "    pca_var = 100 * pca.explained_variance_ratio_\n",
    "\n",
    "    total_count_selection = alt.selection_single(\n",
    "        fields=[\"log10_total_count\"],\n",
    "        init={\n",
    "            \"log10_total_count\": numpy.log(synonymous_spectra_min_counts)\n",
    "            / numpy.log(10)\n",
    "        },\n",
    "        bind=alt.binding_range(\n",
    "            name=\"minimum log10 total counts\",\n",
    "            min=int(mut_type_freqs_pca_coords[\"log10_total_count\"].min()),\n",
    "            max=mut_type_freqs_pca_coords[\"log10_total_count\"].max(),\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    clade_selection = alt.selection_multi(fields=[\"clade\"], bind=\"legend\")\n",
    "\n",
    "    tooltip = [\"clade\", \"total_count\"]\n",
    "\n",
    "    plot_size = 300  # scaled by component variance explained\n",
    "\n",
    "    pca_chart = (\n",
    "        alt.Chart(mut_type_freqs_pca_coords)\n",
    "        .encode(\n",
    "            y=alt.Y(\n",
    "                \"principal_component_1\",\n",
    "                title=f\"PC1 ({pca_var[0]:.0f}% variance)\",\n",
    "                scale=alt.Scale(nice=False, padding=10),\n",
    "                axis=alt.Axis(labels=False, ticks=False),\n",
    "            ),\n",
    "            x=alt.X(\n",
    "                \"principal_component_2\",\n",
    "                title=f\"PC2 ({pca_var[1]:.0f}% variance)\",\n",
    "                scale=alt.Scale(nice=False, padding=10),\n",
    "                axis=alt.Axis(labels=False, ticks=False),\n",
    "            ),\n",
    "            color=alt.Color(\n",
    "                \"clade\",\n",
    "                scale=alt.Scale(\n",
    "                    scheme=\"viridis\",\n",
    "                    domain=sorted(mut_type_freqs_pca_coords[\"clade\"].unique()),\n",
    "                ),\n",
    "            ),\n",
    "            strokeWidth=alt.condition(clade_selection, alt.value(1.5), alt.value(0)),\n",
    "            opacity=alt.condition(clade_selection, alt.value(0.9), alt.value(0.45)),\n",
    "            size=alt.condition(clade_selection, alt.value(65), alt.value(45)),\n",
    "        )\n",
    "        .mark_point(filled=True, stroke=\"black\")\n",
    "        .add_selection(total_count_selection, clade_selection)\n",
    "        .transform_filter(\n",
    "            total_count_selection.log10_total_count <= alt.datum.log10_total_count\n",
    "        )\n",
    "        .configure_axis(grid=False)\n",
    "        .configure_legend(columns=2)\n",
    "        .properties(\n",
    "            height=plot_size,\n",
    "            width=plot_size * pca_var[1] / pca_var[0],\n",
    "            title=title,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    if len(subsets) > 1:\n",
    "        subset_selection = alt.selection_multi(fields=[\"subset\"], bind=\"legend\")\n",
    "        pca_chart = (\n",
    "            pca_chart.encode(\n",
    "                shape=alt.Shape(\n",
    "                    \"subset\", sort=subset_order, scale=alt.Scale(domain=subset_order)\n",
    "                )\n",
    "            )\n",
    "            .add_selection(subset_selection)\n",
    "            .transform_filter(subset_selection)\n",
    "        )\n",
    "        tooltip.append(\"subset\")\n",
    "\n",
    "    if partition:\n",
    "        partition_selection = alt.selection_multi(fields=[\"partition\"], bind=\"legend\")\n",
    "        pca_chart = (\n",
    "            pca_chart.encode(\n",
    "                shape=alt.Shape(\n",
    "                    \"partition\",\n",
    "                    scale=alt.Scale(\n",
    "                        domain=mut_type_freqs_pca_coords[\"partition\"].unique()\n",
    "                    ),\n",
    "                ),\n",
    "            )\n",
    "            .add_selection(partition_selection)\n",
    "            .transform_filter(partition_selection)\n",
    "        )\n",
    "        tooltip.append(\"partition\")\n",
    "\n",
    "    if exclude_top:\n",
    "        exclude_selection = alt.selection_multi(fields=[\"excluded\"], bind=\"legend\")\n",
    "        pca_chart = (\n",
    "            pca_chart.encode(\n",
    "                shape=alt.Shape(\n",
    "                    \"excluded\",\n",
    "                    title=f\"excluding top {exclude_top_n} per clade\",\n",
    "                    scale=alt.Scale(\n",
    "                        domain=mut_type_freqs_pca_coords[\"excluded\"].unique()\n",
    "                    ),\n",
    "                ),\n",
    "            )\n",
    "            .add_selection(exclude_selection)\n",
    "            .transform_filter(exclude_selection)\n",
    "        )\n",
    "        tooltip.append(\n",
    "            alt.Tooltip(\"excluded\", title=f\"excluding top {exclude_top_n} per clade\")\n",
    "        )\n",
    "\n",
    "    pca_chart = pca_chart.encode(tooltip=tooltip)\n",
    "\n",
    "    display(pca_chart)\n",
    "    print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e065e3a2-862f-44fe-ae8a-61dce578bdd4",
   "metadata": {},
   "source": [
    "Compute statistical significance of differences between clades.\n",
    "We just do this on \"all\" sequences for a clade, not partitioning the genomes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d861e0a9-b905-4313-aacf-0ae5d306ec05",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_mut_type_counts = (\n",
    "    mut_type_counts.query(\"subset == 'all'\")\n",
    "    .drop(columns=\"subset\")\n",
    "    .assign(total_count=lambda x: x.groupby(\"clade\")[\"count\"].transform(\"sum\"))\n",
    "    .query(\"total_count >= @synonymous_spectra_min_counts\")\n",
    "    .drop(columns=\"total_count\")\n",
    ")\n",
    "\n",
    "wide_all_mut_type_counts = all_mut_type_counts.pivot_table(\n",
    "    index=\"mut_type\",\n",
    "    columns=\"clade\",\n",
    "    values=\"count\",\n",
    "    fill_value=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a1b3c05-86cc-497b-8e8f-069a04d60191",
   "metadata": {},
   "source": [
    "Now run chi2 test.\n",
    "Also, Bonferroni correct the P-values (this is conservative, but is fine as these P-values are so tiny):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0305e071-1620-4fe2-becd-c41ed7539fb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_p = 1e-20  # plot P-values less than this as this\n",
    "\n",
    "records = []\n",
    "for clade1, clade2 in itertools.combinations(wide_all_mut_type_counts.columns, 2):\n",
    "    chi2, p, dof, _ = scipy.stats.chi2_contingency(\n",
    "        wide_all_mut_type_counts[[clade1, clade2]]\n",
    "    )\n",
    "    records.append((clade1, clade2, p, chi2))\n",
    "\n",
    "chi2_stats = pd.DataFrame(records, columns=[\"clade_1\", \"clade_2\", \"p\", \"chi2\"]).assign(\n",
    "    p=lambda x: x[\"p\"].clip(lower=min_p),\n",
    "    bonferroni_p=lambda x: (x[\"p\"] * len(x)).clip(upper=1),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1f7188a-e817-46d2-89bc-f90403598963",
   "metadata": {},
   "source": [
    "Plot the Bonferroni corrected P-values.\n",
    "Note since counts are very large, many comparisons will be highly significant:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f632d6a-da0a-4ca3-b1a9-0e7bfaf4b235",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_chart = (\n",
    "    alt.Chart(chi2_stats)\n",
    "    .encode(\n",
    "        x=alt.X(\"clade_1\", title=None),\n",
    "        y=alt.Y(\"clade_2\", title=None),\n",
    "        fill=alt.Fill(\n",
    "            \"bonferroni_p\",\n",
    "            title=\"Bonferroni corrected P-value\",\n",
    "            scale=alt.Scale(type=\"log\", scheme=\"yelloworangered\", reverse=True),\n",
    "            legend=alt.Legend(orient=\"top\"),\n",
    "        ),\n",
    "        tooltip=[\n",
    "            \"clade_1\",\n",
    "            \"clade_2\",\n",
    "            alt.Tooltip(\"p\", format=\".2g\"),\n",
    "            alt.Tooltip(\"bonferroni_p\", format=\".2g\"),\n",
    "            alt.Tooltip(\"chi2\", format=\".2g\"),\n",
    "        ],\n",
    "    )\n",
    "    .mark_rect(stroke=\"black\")\n",
    "    .properties(width=alt.Step(14), height=alt.Step(14))\n",
    ")\n",
    "\n",
    "p_chart"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39b1a1ff-6dfc-4101-9872-da16d3cc38a1",
   "metadata": {},
   "source": [
    "Now get the mutation counts and fractions for each clade, both with and without excluding top mutations.\n",
    "Then also compute a normalized **rate** for each mutation type, which is the fraction of all mutations that are of that type divided by the overall fraction of 4-fold synonymous sites sites that are the parental nucleotide in the mutation.\n",
    "\n",
    "Note that a **caveat** is that for these rates we do not adjust the composition to account for any excluded 4-fold synonymous sites.\n",
    "This is probably not currently a big problem, but could become a concern if a lot of sites are excluded:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66013519-7dff-468b-94dc-9f7e67552751",
   "metadata": {},
   "outputs": [],
   "source": [
    "parent_composition = (\n",
    "    pd.read_csv(clade_founder_nts_csv)\n",
    "    .query(\"four_fold_degenerate\")\n",
    "    .groupby([\"clade\", \"nt\"], as_index=False)\n",
    "    .aggregate(parent_nt_count=pd.NamedAgg(\"site\", \"nunique\"))\n",
    "    .assign(\n",
    "        parent_nt_frac=lambda x: (\n",
    "            x[\"parent_nt_count\"]\n",
    "            / x.groupby(\"clade\")[\"parent_nt_count\"].transform(\"sum\")\n",
    "        ),\n",
    "    )\n",
    ")\n",
    "\n",
    "mut_type_count_frac_rate = (\n",
    "    pd.concat(\n",
    "        [\n",
    "            mut_type_counts[[\"subset\", \"clade\", \"mut_type\", \"count\"]]\n",
    "            .assign(exclude_top_mutations=\"no\"),\n",
    "            mut_type_counts_exclude_top\n",
    "            .assign(exclude_top_mutations=\"yes\")\n",
    "            .assign(subset=\"all\")\n",
    "        ]\n",
    "    )\n",
    "    .assign(\n",
    "        total_count=lambda x: x.groupby([\"subset\", \"clade\", \"exclude_top_mutations\"])[\n",
    "            \"count\"\n",
    "        ].transform(\"sum\"),\n",
    "        fraction=lambda x: x[\"count\"] / x[\"total_count\"],\n",
    "        parent_nt=lambda x: x[\"mut_type\"].str[0],\n",
    "    )\n",
    "    .query(\"total_count >= @synonymous_spectra_min_counts\")\n",
    "    .merge(\n",
    "        parent_composition[[\"clade\", \"nt\", \"parent_nt_frac\"]].rename(\n",
    "            columns={\"nt\": \"parent_nt\"}\n",
    "        ),\n",
    "        how=\"left\",\n",
    "        validate=\"many_to_one\",\n",
    "    )\n",
    "    .assign(rate=lambda x: x[\"fraction\"] / x[\"parent_nt_frac\"])\n",
    "    .sort_values([\"subset\", \"clade\", \"mut_type\"], ascending=[False, True, True])\n",
    "    .reset_index(drop=True)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5981a02-4822-4c3b-bd43-7b2043191820",
   "metadata": {},
   "source": [
    "Plot the rates. The chart below is interactive, and you can click to select mutation types, clades, etc:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e834892-2b7b-453e-a61c-d5d6b98302b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "melted_df = (\n",
    "    mut_type_count_frac_rate.assign(\n",
    "        mut_type=lambda x: x[\"mut_type\"].str.replace(\"to\", \" -> \")\n",
    "    )\n",
    "    .query(\"subset == 'all'\")\n",
    "    .melt(\n",
    "        id_vars=[\"clade\", \"mut_type\", \"exclude_top_mutations\", \"count\"],\n",
    "        value_vars=[\"fraction\", \"rate\"],\n",
    "    )\n",
    "    .assign(\n",
    "        variable=lambda x: x[\"variable\"].map(\n",
    "            {\n",
    "                \"fraction\": \"fraction of all mutations\",\n",
    "                \"rate\": \"relative rate of mutation\",\n",
    "            }\n",
    "        ),\n",
    "        clade=lambda x: x[\"clade\"].map(lambda c:\n",
    "            f\"{c} ({clade_synonyms[c]})\" if c in clade_synonyms else c\n",
    "        ),\n",
    "    )\n",
    ")\n",
    "\n",
    "mut_type_selection = alt.selection_multi(fields=[\"mut_type\"])\n",
    "\n",
    "mut_type_selection_bar = (\n",
    "    alt.Chart(melted_df[[\"mut_type\"]].drop_duplicates())\n",
    "    .encode(\n",
    "        x=alt.X(\"mut_type\", title=\"click / shift-click to select mutation types\"),\n",
    "        color=alt.condition(\n",
    "            mut_type_selection, alt.value(\"darkgray\"), alt.value(\"white\")\n",
    "        ),\n",
    "    )\n",
    "    .mark_rect(stroke=\"black\")\n",
    "    .add_selection(mut_type_selection)\n",
    "    .properties(width=alt.Step(15))\n",
    ")\n",
    "\n",
    "frac_rate_chart_base = (\n",
    "    alt.Chart(melted_df)\n",
    "    .encode(\n",
    "        x=alt.X(\"mut_type\", title=\"mutation type\"),\n",
    "        y=alt.Y(\"value\", title=None),\n",
    "        color=alt.Color(\"clade\", scale=alt.Scale(scheme=\"viridis\")),\n",
    "        strokeWidth=alt.condition(clade_selection, alt.value(0.5), alt.value(0)),\n",
    "        opacity=alt.condition(clade_selection, alt.value(0.9), alt.value(0.2)),\n",
    "        size=alt.condition(clade_selection, alt.value(50), alt.value(40)),\n",
    "        column=alt.Column(\n",
    "            \"variable\",\n",
    "            title=None,\n",
    "            header=alt.Header(\n",
    "                labelOrient=\"left\", labelFontSize=11, labelFontStyle=\"bold\"\n",
    "            ),\n",
    "        ),\n",
    "        tooltip=[\n",
    "            alt.Tooltip(c, format=\".3g\") if melted_df[c].dtype == float else c\n",
    "            for c in melted_df.columns.tolist()\n",
    "        ],\n",
    "    )\n",
    "    .mark_point(filled=True, stroke=\"black\")\n",
    "    .add_selection(clade_selection, mut_type_selection)\n",
    "    .transform_filter(mut_type_selection)\n",
    "    .resolve_scale(y=\"independent\")\n",
    "    .properties(width=alt.Step(15), height=200)\n",
    ")\n",
    "\n",
    "frac_rate_chart = (\n",
    "    (frac_rate_chart_base & mut_type_selection_bar)\n",
    "    .configure_axis(grid=False)\n",
    "    .configure_legend(columns=2)\n",
    "    .configure(padding=20)\n",
    ")\n",
    "\n",
    "print(f\"Saving to {rates_plot}\")\n",
    "frac_rate_chart.save(rates_plot)\n",
    "\n",
    "frac_rate_chart"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12416fc3-a6a2-4f35-a25c-e04129ce6f10",
   "metadata": {},
   "source": [
    "Get rates to write to file.\n",
    "Get the rates by clade **without** excluding the top mutations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "374f820a-54f6-4907-86f2-f78b8d5286a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "rates_by_clade = mut_type_count_frac_rate.query(\"exclude_top_mutations == 'no'\").drop(\n",
    "    columns=\"exclude_top_mutations\"\n",
    ")\n",
    "\n",
    "print(f\"Writing rates to {rates_by_clade_csv}\")\n",
    "rates_by_clade.to_csv(rates_by_clade_csv, index=False, float_format=\"%.5g\")\n",
    "\n",
    "rates_by_clade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7fbdab4-a46a-4fd0-a553-8172bf0f42ef",
   "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
}