notebooks/clade_founder_aas.py.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6f5c9c5d-1626-480c-a441-a6cb34091904",
   "metadata": {},
   "source": [
    "# Get clade founder amino acids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "259b4fde-cd06-4491-a763-65af490c9d48",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-06-13T22:16:13.364786Z",
     "iopub.status.busy": "2023-06-13T22:16:13.364153Z",
     "iopub.status.idle": "2023-06-13T22:16:13.790698Z",
     "shell.execute_reply": "2023-06-13T22:16:13.789934Z",
     "shell.execute_reply.started": "2023-06-13T22:16:13.364740Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import Bio.Seq\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import yaml\n",
    "\n",
    "clade_founder_nts = pd.read_csv(snakemake.input.clade_founder_nts)\n",
    "\n",
    "orf1ab_to_nsps = snakemake.params.orf1ab_to_nsps\n",
    "clade_synonyms = snakemake.params.clade_synonyms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "edb2edcb-85c7-4cba-898a-0896bbb16977",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-06-13T22:21:33.143704Z",
     "iopub.status.busy": "2023-06-13T22:21:33.143146Z",
     "iopub.status.idle": "2023-06-13T22:21:35.425446Z",
     "shell.execute_reply": "2023-06-13T22:21:35.424695Z",
     "shell.execute_reply.started": "2023-06-13T22:21:33.143658Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['ORF1ab', 'S', 'ORF3a', 'E', 'M', 'ORF6', 'ORF7a', 'ORF7b', 'ORF8',\n",
       "       'N', 'ORF9b', 'ORF10', 'nsp1', 'nsp2', 'nsp3', 'nsp4',\n",
       "       'nsp5 (Mpro)', 'nsp6', 'nsp7', 'nsp8', 'nsp9', 'nsp10',\n",
       "       'nsp12 (RdRp)', 'nsp13', 'nsp14', 'nsp15', 'nsp16'], dtype=object)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def clade_label(clade):\n",
    "    if clade in clade_synonyms:\n",
    "        return f\"{clade} ({clade_synonyms[clade]})\"\n",
    "    else:\n",
    "        return clade\n",
    "\n",
    "# codon translation table\n",
    "codon_table = {\n",
    "    f\"{nt1}{nt2}{nt3}\": str(Bio.Seq.Seq(f\"{nt1}{nt2}{nt3}\").translate())\n",
    "    for nt1 in \"ACGT\" for nt2 in \"ACGT\" for nt3 in \"ACGT\"\n",
    "}\n",
    "\n",
    "# get clade founder amino-acids\n",
    "clade_founder_aas = (\n",
    "    clade_founder_nts\n",
    "    .query(\"gene != 'noncoding'\")\n",
    "    [[\"clade\", \"gene\", \"codon\", \"codon_site\"]]\n",
    "    .drop_duplicates()\n",
    "    .assign(\n",
    "        gene=lambda x: x[\"gene\"].str.split(\";\"),\n",
    "        codon=lambda x: x[\"codon\"].str.split(\";\"),\n",
    "        codon_site=lambda x: x[\"codon_site\"].str.split(\";\"),\n",
    "    )\n",
    "    .explode([\"gene\", \"codon\", \"codon_site\"])\n",
    "    .assign(\n",
    "        aa=lambda x: x[\"codon\"].map(codon_table),\n",
    "        codon_site=lambda x: x[\"codon_site\"].astype(int),\n",
    "        clade=lambda x: x[\"clade\"].map(clade_label),\n",
    "    )\n",
    "    .rename(columns={\"codon_site\": \"site\", \"aa\": \"amino acid\"})\n",
    "    .drop(columns=\"codon\")\n",
    "    .query(\"gene != 'ORF1a'\")  # this is just subset of ORF1ab\n",
    ")\n",
    "\n",
    "# now convert ORF1ab numbers to nsp numbers\n",
    "orf1ab_to_nsps_df = pd.concat(\n",
    "    [\n",
    "        pd.DataFrame(\n",
    "            [(i, i - start + 1) for i in range(start, end + 1)],\n",
    "            columns=[\"ORF1ab_site\", \"nsp_site\"],\n",
    "        ).assign(nsp=nsp).drop_duplicates()\n",
    "        for nsp, (start, end) in orf1ab_to_nsps.items()\n",
    "    ],\n",
    "    ignore_index=True,\n",
    ")\n",
    "\n",
    "clade_founder_aas = pd.concat(\n",
    "    [\n",
    "        clade_founder_aas,\n",
    "        (\n",
    "            clade_founder_aas\n",
    "            .query(\"gene == 'ORF1ab'\")\n",
    "            .merge(\n",
    "                orf1ab_to_nsps_df,\n",
    "                left_on=\"site\",\n",
    "                right_on=\"ORF1ab_site\",\n",
    "                validate=\"many_to_one\",\n",
    "            )\n",
    "            .drop(columns=[\"gene\", \"ORF1ab_site\", \"site\"])\n",
    "            .rename(columns={\"nsp_site\": \"site\", \"nsp\": \"gene\"})\n",
    "        ),\n",
    "    ],\n",
    "    ignore_index=False,\n",
    ").drop_duplicates()\n",
    "\n",
    "clade_founder_aas.to_csv(snakemake.output.clade_founder_aas, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f41c4fa-f089-42ea-bb71-335242e6408a",
   "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
}