b0-parse-evidence.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ac08accd-0dbd-4673-8a20-a986c5ebe743",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%bash\n",
    "# diamond blastp \\\n",
    "#     -q sarg_ref.fa \\\n",
    "#     -d tmp/refseq_protein.dmnd \\\n",
    "#     --out tmp/refseq_protein.txt \\\n",
    "#     --outfmt 6 qseqid sseqid stitle nident qlen slen pident qcovhsp scovhsp bitscore evalue \\\n",
    "#     --id 50 --subject-cover 75 --query-cover 75 \\\n",
    "#     -k 0 --threads 48 --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9d278fbe-e8d3-4cf5-9f98-d7a959b27872",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import time\n",
    "import glob\n",
    "import regex as re\n",
    "import json\n",
    "import gzip\n",
    "\n",
    "from collections import defaultdict\n",
    "from Bio import SeqIO\n",
    "from tqdm.contrib.concurrent import process_map\n",
    "\n",
    "names = set()\n",
    "with open('tmp/refseq_protein.txt') as f:\n",
    "    for line in f:\n",
    "        ls = line.rstrip().split('\\t')\n",
    "        name = ls[2].split(' >')[0].split(' ',1)[1].rsplit(' [')[0].split('MULTISPECIES: ')[-1]\n",
    "        if 'hypothetical' not in name and ', partial' not in name:\n",
    "            names.add(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "08945252-afc1-4b2a-b0ea-a0c9482e3e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "rep = set()\n",
    "ref = set()\n",
    "rev = defaultdict(set)\n",
    "tset = set()\n",
    "with open('sarg_ref.fa') as handle:\n",
    "    for record in SeqIO.parse(handle, 'fasta'):\n",
    "        name = record.description.split(' >')[0].split(' ',1)[1].rsplit(' [')[0].split('MULTISPECIES: ')[-1]\n",
    "        rev[name].add('@'.join(record.id.split('|')[1:3]))\n",
    "        rep.add(record.seq)\n",
    "        ref.add(record.id.split('|')[-1])\n",
    "        tset.add('@'.join(record.id.split('|')[1:3]))\n",
    "\n",
    "rev = {x: ' | '.join(sorted(y)) for x, y in rev.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7faf3cb8-cc6f-4324-98c3-934f44da06e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "98293e5c7c7b4dd78df0017ae335dd93",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/593"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def reader(file):\n",
    "    df = pd.read_table(file, header=None, low_memory=False, names= ['accession','name','evidence','symbol','identifier']).fillna('NA')\n",
    "    df = df[df.name.isin(names)]\n",
    "    df = df[~df.accession.isin(ref)]\n",
    "    df = df[~df.accession.isin(df[df.evidence.str.contains('NA|Domain')].accession)]\n",
    "    df = df[~df.accession.isin(df[(df.evidence.str.contains('TIGR') | (df.identifier.str.contains('PF'))) & (df.symbol == 'NA')].accession)]\n",
    "    return df\n",
    "\n",
    "evidence = pd.concat(process_map(reader, glob.glob('tmp/refseq/*.tsv'), bar_format = '{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt}', max_workers=48, chunksize=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f0e6ac0-e48d-4c79-9e65-d2241dcb60c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "## bugfix\n",
    "## https://www.ncbi.nlm.nih.gov/genome/annotation_prok/evidence/NBR012193/\n",
    "evidence.loc[evidence.evidence == 'NBR012193', 'name'] = 'quaternary ammonium compound efflux SMR transporter EmrC'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6f70cb62-ab40-4109-80a8-b25b85d853b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "evidence['cnt'] = evidence.groupby(['name', 'evidence']).transform('size')\n",
    "rmk = evidence[['name', 'evidence', 'symbol', 'identifier', 'cnt']].drop_duplicates()\n",
    "rmk['mapped'] = rmk.name.map(rev).fillna('')\n",
    "rmk['unique'] = (rmk.mapped.str.count('\\\\|') == 0) & (rmk.mapped != '')\n",
    "rmk['valid'] = rmk.apply(lambda x: re.sub('^bla|\\\\(|\\\\)', '', x.mapped.split('@')[-1].lower()) in re.sub('\\\\(|\\\\)', '', x['name'].lower()) if len(x.mapped.split('@')[-1]) > 4 else x.mapped.lower().split('@')[-1] in x['name'].lower(), axis=1)\n",
    "rmk.loc[rmk.name.str.contains('/'), 'valid'] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f3b3a8c4-90d8-4c49-99ba-16d185b0b81a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unreviewed: 329\n"
     ]
    }
   ],
   "source": [
    "with open('misc/whitelist.json') as f:\n",
    "    wl = json.load(f)\n",
    "\n",
    "rmkk = rmk[~((rmk.unique) & (rmk.valid))].copy()\n",
    "rmkk['reviewed'] = rmkk.name.apply(lambda x: wl.get(x, {}).get('reviewed', ''))\n",
    "rmkk = rmkk.sort_values(['name', 'evidence', 'symbol'])\n",
    "\n",
    "## easy cases\n",
    "bla = rmkk[(rmkk.symbol.str.contains('^bla', regex=True)) & (rmkk.symbol.str.replace('^bla', '', regex=True) == rmkk.name.str.split(' |-').str.get(0)) & (~rmkk.symbol.str.contains('CMA|CSA|CFE|LAT|LCR|NPS', regex=True))].copy()\n",
    "bla['gene'] = 'beta-lactam@' + bla.symbol\n",
    "mcr = rmk[rmk.name.str.contains('^MCR', regex=True)].copy()\n",
    "mcr['gene'] = 'colistin@' + mcr.name.str.split(' |-related').str.get(0).str.lower()\n",
    "\n",
    "rmkk = rmkk[~(rmkk.name.isin(bla.name) | (rmkk.name.isin(mcr.name)))]\n",
    "wl = defaultdict(dict)\n",
    "for _, i in rmkk.iterrows():\n",
    "    wl[i['name']]['reviewed'] = i.reviewed\n",
    "    wl[i['name']]['gene'] = i.mapped\n",
    "    if 'remark' not in wl[i['name']]:\n",
    "        wl[i['name']]['remark'] = []\n",
    "    wl[i['name']]['remark'].append(' | '.join([i.evidence, i.symbol, i.identifier, str(i['cnt'])]))\n",
    "\n",
    "wl = {x:y for x,y in wl.items() if x in set(rmkk.name)}\n",
    "with open('misc/whitelist.json', 'w') as f:\n",
    "    json.dump(wl, f, sort_keys=True, indent=4)\n",
    "\n",
    "print('unreviewed:', rmkk.name.nunique())\n",
    "assert sum(rmkk.reviewed == '') == 0, 'Make sure all uncertain cases are reviewed in <misc/whitelist.json>.'\n",
    "\n",
    "rev.update(rmkk.set_index('name').reviewed.to_dict())\n",
    "rev.update(bla.set_index('name').gene.to_dict())\n",
    "rev.update(mcr.set_index('name').gene.to_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d19bc585-a585-4598-96bd-87aad6bff896",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for i,j in rmkk[(rmkk.reviewed == '') | (rmkk.reviewed.str.contains(' '))].groupby(['name', 'mapped', 'reviewed']):\n",
    "    print(f'{i[0]}')\n",
    "    print(f'gene: {i[1]}')\n",
    "    print(f'reviewed: {i[2]}')\n",
    "\n",
    "    for _, k in j.iterrows():\n",
    "        print(k.evidence, k.symbol, k.identifier, k['cnt'])\n",
    "\n",
    "    print('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "de612794-3d33-4796-8eb5-853e9b6caf98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a1e6bc9119e24b67b97c6396363b4be2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/593"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def parser(file):\n",
    "    records = []\n",
    "    with gzip.open(file, 'rt') as f:\n",
    "        for record in SeqIO.parse(f, 'fasta'):\n",
    "            if record.seq not in rep:\n",
    "                name = record.description.split(' >')[0].split(' ',1)[1].rsplit(' [')[0].split('MULTISPECIES: ')[-1]\n",
    "                gene = rev.get(name)\n",
    "                if gene is not None and gene != '' and gene != '*':\n",
    "                    if record.id in aset:\n",
    "                        record.id = 'SARG|' + gene.replace('@', '|') + '|' + record.id\n",
    "                        record.description = record.description.split(' ', 1)[-1].split(' >')[0]\n",
    "                        records.append(record)\n",
    "    return records\n",
    "\n",
    "aset = set(evidence.accession)\n",
    "r = process_map(parser, glob.glob('tmp/refseq/*.faa.gz'), bar_format = '{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt}', max_workers=48, chunksize=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "69c4723c-9074-42bb-80af-1017a499f38d",
   "metadata": {},
   "outputs": [],
   "source": [
    "records = []\n",
    "ids = set()\n",
    "sset = set()\n",
    "for i in r:\n",
    "    records.extend(i)\n",
    "    ids.update(set(record.id.split('|')[-1] for record in i))\n",
    "    sset.update(set('@'.join(record.id.split('|')[1:3]) for record in i))\n",
    "\n",
    "with open('tmp/sarg_tmp.fa', 'w') as output_handle:\n",
    "    SeqIO.write(records, output_handle, 'fasta')\n",
    "\n",
    "evidence['type'] = evidence.name.map(rev).str.split('@').str.get(0).fillna('')\n",
    "evidence['subtype'] = evidence.name.map(rev).str.split('@').str.get(1).fillna('')\n",
    "evidence[evidence.accession.isin(ids)].groupby(['type', 'subtype', 'name', 'evidence', 'symbol', 'identifier'], as_index=False).size().rename({'size': 'cnt'}, axis=1).to_csv('misc/evidence.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "abe7cae5-d3a2-4bc2-b645-3c835ee06318",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\"aminoglycoside@aph(3'')*\",\n",
       " 'aminoglycoside@rmt*',\n",
       " 'beta-lactam@blaASDC',\n",
       " 'beta-lactam@blaCSR',\n",
       " 'beta-lactam@mec*',\n",
       " 'macrolide-lincosamide-streptogramin@erm*',\n",
       " 'macrolide-lincosamide-streptogramin@lsa*',\n",
       " 'macrolide-lincosamide-streptogramin@mph*',\n",
       " 'macrolide-lincosamide-streptogramin@msr*',\n",
       " 'macrolide-lincosamide-streptogramin@sal*',\n",
       " 'macrolide-lincosamide-streptogramin@vat*',\n",
       " 'macrolide-lincosamide-streptogramin@vga*',\n",
       " 'nitroimidazole@nim*',\n",
       " 'rifamycin@rph*',\n",
       " 'tetracycline@tet*',\n",
       " 'trimethoprim@dfr*'}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sset - tset"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}