notebooks/merge_expected_and_actual_counts.py.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2c0cec4-3606-40b3-8c31-f6a99f195ff0",
   "metadata": {},
   "source": [
    "# Merge expected and actual counts\n",
    "\n",
    "First, read the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45108db9-2934-4c3b-bf77-eef7eb876078",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "import Bio.Seq\n",
    "\n",
    "import numpy\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "expected = (\n",
    "    pd.read_csv(snakemake.input.expected)\n",
    "    .rename(columns={\"site\": \"nt_site\", \"codon\": \"clade_founder_codon\"})\n",
    "    .assign(\n",
    "        nt_mutation=lambda x: (\n",
    "            x[\"clade_founder_nt\"] + x[\"nt_site\"].astype(str) + x[\"mut_type\"].str[-1]\n",
    "        )\n",
    "    )\n",
    ")\n",
    "\n",
    "actual = pd.read_csv(snakemake.input.actual)[\n",
    "    [\n",
    "        \"clade\",\n",
    "        \"subset\",\n",
    "        \"nt_site\",\n",
    "        \"nt_mutation\",\n",
    "        \"codon_change\",\n",
    "        \"count\",\n",
    "        \"count_terminal\",\n",
    "        \"count_non_terminal\",\n",
    "        \"mean_log_size\",\n",
    "    ]\n",
    "].rename(columns={\"count\": \"actual_count\"})\n",
    "\n",
    "muts_to_exclude = pd.read_csv(snakemake.input.muts_to_exclude).rename(\n",
    "    columns={\"site\": \"nt_site\", \"mutation\": \"nt_mutation\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "548d318f-7479-4564-9b6f-de6297913470",
   "metadata": {},
   "source": [
    "Merge expected and actual, putting counts of any expected that are not observed to zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a971df6-d441-4562-b20b-17a2a302320c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "merged = (\n",
    "    expected\n",
    "    .merge(actual, how=\"left\", validate=\"one_to_many\")\n",
    "    .assign(\n",
    "        actual_count=lambda x: x[\"actual_count\"].fillna(0).astype(int),\n",
    "        count_terminal=lambda x: x[\"count_terminal\"].fillna(0).astype(int),\n",
    "        count_non_terminal=lambda x: x[\"count_non_terminal\"].fillna(0).astype(int),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de046433-b7cd-4a8e-897f-f46a5be54ab8",
   "metadata": {},
   "source": [
    "Now exclude any mutations that are not from the founder identity codon, which results when there are multiple substitutions per site so that a clade founder identity is different from the reference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af0e5cac-66de-47c5-9684-7b517bbb522e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# make sure founder nucleotide identity already matches clade founder\n",
    "assert (merged[\"clade_founder_nt\"] == merged[\"nt_mutation\"].str[0]).all()\n",
    "\n",
    "merged = merged.assign(\n",
    "    start_codon=lambda x: x[\"codon_change\"].map(\n",
    "        lambda s: pd.NA if (pd.isnull(s) or s == \"noncoding\") else \";\".join([c[:3] for c in s.split(\";\")])\n",
    "    ),\n",
    "    from_founder_codon=lambda x: (\n",
    "        (x[\"start_codon\"] == x[\"clade_founder_codon\"]) | x[\"start_codon\"].isnull()\n",
    "    ),\n",
    ").query(\"from_founder_codon\")\n",
    "\n",
    "# add back any zero count mutations lost to this process\n",
    "merged = (\n",
    "    merged\n",
    "    .merge(expected, how=\"outer\")\n",
    "    .assign(\n",
    "        actual_count=lambda x: x[\"actual_count\"].fillna(0).astype(int),\n",
    "        count_terminal=lambda x: x[\"count_terminal\"].fillna(0).astype(int),\n",
    "        count_non_terminal=lambda x: x[\"count_non_terminal\"].fillna(0).astype(int),\n",
    "        mean_log_size=lambda x: x[\"mean_log_size\"].fillna(0),\n",
    "    )\n",
    ")\n",
    "assert len(expected) == len(merged)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c1c8fd1-f25e-45d9-8822-993cb30d81ce",
   "metadata": {},
   "source": [
    "Now add the sites to exclude:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45426501-964e-4acb-a0ac-b9bf7f5b9ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "merged = merged.merge(\n",
    "    muts_to_exclude.assign(exclude=True),\n",
    "    on=[\"clade\", \"nt_site\", \"nt_mutation\"],\n",
    "    how=\"left\",\n",
    "    validate=\"many_to_one\",\n",
    ").assign(\n",
    "    exclude=lambda x: x[\"exclude\"].fillna(False),\n",
    "    masked_in_usher=lambda x: x[\"masked_in_usher\"].fillna(False),\n",
    ")\n",
    "\n",
    "merged"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59326383-1066-4395-84f5-a2dc2a4df757",
   "metadata": {},
   "source": [
    "Get the amino-acid mutation translations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d360a5b8-9b15-49ad-b529-9a636b71c70a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "translation_table = {}\n",
    "for nt1, nt2, nt3 in itertools.product(\"ACGT\", \"ACGT\", \"ACGT\"):\n",
    "    codon = f\"{nt1}{nt2}{nt3}\"\n",
    "    translation_table[codon] = str(Bio.Seq.Seq(codon).translate())\n",
    "translation_table[\"noncoding\"] = \"noncoding\"\n",
    "\n",
    "\n",
    "def mutant_codon(row):\n",
    "    codon = row[\"clade_founder_codon\"]\n",
    "    if codon == \"noncoding\":\n",
    "        return \"noncoding\"\n",
    "    i = int(row[\"codon_position\"])\n",
    "    mut = row[\"nt_mutation\"]\n",
    "    assert mut[0] == codon[i - 1]\n",
    "    return codon[: i - 1] + mut[-1] + codon[i:]\n",
    "\n",
    "\n",
    "add_aa_muts = (\n",
    "    merged[\n",
    "        [\"clade\", \"nt_mutation\", \"clade_founder_codon\", \"codon_position\", \"codon_site\"]\n",
    "    ]\n",
    "    .drop_duplicates()\n",
    "    .assign(\n",
    "        clade_founder_codon=lambda x: x[\"clade_founder_codon\"].str.split(\";\"),\n",
    "        codon_position=lambda x: x[\"codon_position\"].str.split(\";\"),\n",
    "        codon_site=lambda x: x[\"codon_site\"].str.split(\";\"),\n",
    "    )\n",
    "    .explode([\"clade_founder_codon\", \"codon_position\", \"codon_site\"])\n",
    "    .assign(\n",
    "        mutant_codon=lambda x: x.apply(mutant_codon, axis=1),\n",
    "        clade_founder_aa=lambda x: x[\"clade_founder_codon\"].map(translation_table),\n",
    "        mutant_aa=lambda x: x[\"mutant_codon\"].map(translation_table),\n",
    "        aa_mutation=lambda x: numpy.where(\n",
    "            x[\"clade_founder_aa\"] == \"noncoding\",\n",
    "            \"noncoding\",\n",
    "            x[\"clade_founder_aa\"] + x[\"codon_site\"].astype(str) + x[\"mutant_aa\"],\n",
    "        ),\n",
    "    )\n",
    "    [[\"clade\", \"nt_mutation\", \"mutant_codon\", \"clade_founder_aa\", \"mutant_aa\", \"aa_mutation\"]]\n",
    "    .groupby([\"clade\", \"nt_mutation\"], as_index=False)\n",
    "    .aggregate(lambda s: \";\".join(s))\n",
    "    .assign(\n",
    "        synonymous=lambda x: (\n",
    "            (x[\"clade_founder_aa\"] == x[\"mutant_aa\"])\n",
    "            & (x[\"clade_founder_aa\"] != \"noncoding\")\n",
    "        ),\n",
    "        noncoding=lambda x: x[\"clade_founder_aa\"] == \"noncoding\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bf0393c-1673-4f36-b276-5e6dafa2687d",
   "metadata": {},
   "source": [
    "Now add in amino acid mutations and clean up data frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0522e718-f88e-4514-8d8e-186ebabdf7e4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# clean up data frame to just have columns of interest\n",
    "merged = (\n",
    "    merged.merge(add_aa_muts, validate=\"many_to_one\", how=\"outer\")[\n",
    "        [\n",
    "            \"clade\",\n",
    "            \"subset\",\n",
    "            \"nt_site\",\n",
    "            \"nt_mutation\",\n",
    "            \"exclude\",\n",
    "            \"masked_in_usher\",\n",
    "            \"expected_count\",\n",
    "            \"actual_count\",\n",
    "            \"clade_founder_nt\",\n",
    "            \"gene\",\n",
    "            \"clade_founder_codon\",\n",
    "            \"clade_founder_aa\",\n",
    "            \"mutant_codon\",\n",
    "            \"mutant_aa\",\n",
    "            \"aa_mutation\",\n",
    "            \"synonymous\",\n",
    "            \"noncoding\",\n",
    "            \"codon_position\",\n",
    "            \"codon_site\",\n",
    "            \"four_fold_degenerate\",\n",
    "            \"count_terminal\",\n",
    "            \"count_non_terminal\",\n",
    "            \"mean_log_size\",\n",
    "        ]\n",
    "    ]\n",
    "    .sort_values([\"clade\", \"subset\", \"nt_site\", \"nt_mutation\"])\n",
    "    .reset_index(drop=True)\n",
    ")\n",
    "\n",
    "assert merged.notnull().all().all()\n",
    "\n",
    "merged.to_csv(snakemake.output.csv, index=False, float_format=\"%.5g\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32831c82-fbca-4166-8ee4-3c9ff8902f26",
   "metadata": {},
   "source": [
    "Just a quick look at correlation between actual and expected counts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed77bd39-5f08-479e-be26-b7ebf98da685",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "(\n",
    "    merged[[\n",
    "        \"exclude\", \"synonymous\", \"noncoding\", \"actual_count\", \"expected_count\", \"count_terminal\", \"count_non_terminal\",\n",
    "    ]]\n",
    "    .groupby([\"exclude\", \"synonymous\", \"noncoding\"])\n",
    "    .corr()\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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}