Biopython_get_GenBank_genomes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retrieving all GenBank genome sequences for a bacterial genus\n",
    "\n",
    "Some tasks we are interested in need to be repeated for a set of genomes, whenever new isolate sequences become available. This notebook will take us through the process of collecting all genomes from GenBank that belong to a named genus."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing Biopython\n",
    "\n",
    "We will use Biopython's interface to Entrez, the interface to NCBI's online resources. There are some [guidelines](http://biopython.org/DIST/docs/tutorial/Tutorial.html#htoc109) to its use that need to be followed, so you are not banned by NCBI:\n",
    "\n",
    "1. For any series of more than 100 requests, conduct your work outwith US peak times.\n",
    "2. Make no more than three queries each second (enforced by Biopython).\n",
    "3. Set your `Entrez.email` parameter correctly.\n",
    "4. Set your `Entrez.tool` parameter correctly.\n",
    "5. Use session histories where appropriate.\n",
    "\n",
    "We start by importing packages, including Biopython's `Entrez` module, and setting the appropriate parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from Bio import Entrez, SeqIO\n",
    "\n",
    "Entrez.email = \"\"  # Use your own real email\n",
    "Entrez.tool = \"Biopython_get_GenBank_genomes.ipynb\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Searching the database\n",
    "\n",
    "We're using `Bio.Entrez.esearch()` with the `genome` database to look for our search term. In this case, it's *Pectobacterium*, a genus of plant pathogenic bacteria.\n",
    "\n",
    "We know we can look in the `genome` database, because we checked by hand at [http://www.ncbi.nlm.nih.gov/genome](http://www.ncbi.nlm.nih.gov/assembly), and because it's one of the databases named if we look with `Entrez.einfo()`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "genus = \"Pectobacterium\"\n",
    "query_text = \"{0} AND bacteria[Organism]\".format(genus)\n",
    "handle = Entrez.esearch(db='genome', term=query_text)\n",
    "record = Entrez.read(handle)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can get the number of returned records by looking at `record[\"Count\"]`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "record[\"Count\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But what are our records? We can see their GenBank identifiers by looking at `record[\"IdList\"]`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "record[\"IdList\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But this isn't immediately informative. We're going to have to look at the assemblies associated with these identifiers in GenBank. We do this with `Entrez.elink()`, searching for associations between the `genome` database and the `assembly` database, compiling all the resulting `Link` UIDs in a single list.\n",
    "\n",
    "What are the links we're allowed? Well, there's a big list at [http://www.ncbi.nlm.nih.gov/corehtml/query/static/entrezlinks.html](http://www.ncbi.nlm.nih.gov/corehtml/query/static/entrezlinks.html), but we can also inspect NCBI's web interface directly to see that [http://www.ncbi.nlm.nih.gov/assembly](http://www.ncbi.nlm.nih.gov/assembly) is the likely prefix/database we're looking for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "asm_links = []\n",
    "for uid in record[\"IdList\"]:\n",
    "    links = Entrez.read(Entrez.elink(dbfrom=\"genome\", db=\"assembly\", retmode=\"text\", from_uid=uid))\n",
    "    [asm_links.append(d.values()[0]) for d in links[0]['LinkSetDb'][0]['Link']]\n",
    "print(\"We find {0} genome entries: {1}\".format(len(asm_links), asm_links))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can recover links to the `nucleotide` database for each of these UIDs. There may be several such links, but as we are looking for the full assembly, we care only about the `assembly_nuccore_insdc` sequences, which are the contigs.\n",
    "\n",
    "We collect these into a dictionary of contig UIDs, keyed by assembly UID, called `sequid_links`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sequid_links = {}\n",
    "for uid in asm_links:\n",
    "    links = Entrez.read(Entrez.elink(dbfrom=\"assembly\", db=\"nucleotide\", retmode=\"gb\", from_uid=uid))\n",
    "    contigs = [l for l in links[0]['LinkSetDb'] if l['LinkName'] == 'assembly_nuccore_insdc'][0]\n",
    "    sequid_links[uid] = [e['Id'] for e in contigs['Link']]\n",
    "expected_contigs = {}\n",
    "print(\"There are {0} genomes identified for {1}:\".format(len(sequid_links), genus))\n",
    "for k, v in sorted(sequid_links.items()):\n",
    "    print(\"Assembly UID {0}: {1} contigs\".format(k, len(v)))\n",
    "    expected_contigs[k] = len(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have these `nucleotide` database identifiers, we can grab all the sequences and write them out as multi-FASTA files, with `Entrez.efetch()`. The assembly records themselves though, have to be obtained with `Entrez.esummary()`, and then a byzantine set of keywords navigated, to get the information we're interested in.\n",
    "\n",
    "We use the assembly UID without version number as the filename, and write a `labels.txt` file suitable for use with `pyani`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Make sure there's a relevant output directory\n",
    "if not os.path.exists(genus):\n",
    "    os.mkdir(genus)\n",
    "if not os.path.exists(\"failures\"):\n",
    "    os.mkdir(\"failures\")\n",
    "\n",
    "# Write output\n",
    "with open(os.path.join(genus, 'labels.txt'), 'w') as lfh:\n",
    "    with open(os.path.join(genus, 'classes.txt'), 'w') as cfh:\n",
    "        for asm_uid, contigs in sorted(sequid_links.items()):\n",
    "            # Get assembly record information\n",
    "            asm_record = Entrez.read(Entrez.esummary(db='assembly', id=asm_uid, rettype='text'))\n",
    "            asm_organism = asm_record['DocumentSummarySet']['DocumentSummary'][0]['SpeciesName']\n",
    "            try:\n",
    "                asm_strain = asm_record['DocumentSummarySet']['DocumentSummary'][0]['Biosource']['InfraspeciesList'][0]['Sub_value']\n",
    "            except:\n",
    "                asm_strain = \"\"\n",
    "            gname = asm_record['DocumentSummarySet']['DocumentSummary'][0]['AssemblyAccession'].split('.')[0]\n",
    "            filestem = os.path.join(genus, gname)\n",
    "            # Write a labels.txt and a classes.txt file suitable for pyani\n",
    "            glab, species = asm_organism.split(' ', 1)\n",
    "            glab = glab[0]\n",
    "            labelstr = \"{0}\\t{1}. {2} {3}\".format(gname, glab, species, asm_strain)\n",
    "            print >> lfh, labelstr\n",
    "            print >> cfh, \"{0}\\t{1}\".format(gname, asm_organism)\n",
    "            print(labelstr)\n",
    "            # Get FASTA records for each of the contigs (we could do this with GenBank instead,\n",
    "            # but sometimes these are not formatted correctly with sequences)\n",
    "            query_uids = ','.join(contigs)\n",
    "            tries, success = 0, False\n",
    "            while success == False and tries < 20:\n",
    "                # Also check for total sequence length errors?\n",
    "                try:\n",
    "                    print(\"UID:{0} download attempt {1}\".format(asm_uid, tries + 1))\n",
    "                    records = list(SeqIO.parse(Entrez.efetch(db='nucleotide', id=query_uids,\n",
    "                                                             rettype=\"fasta\", retmode='text'),\n",
    "                                               'fasta'))\n",
    "                    if len(records) == expected_contigs[asm_uid]: # No exceptions, num records = expected\n",
    "                        success = True\n",
    "                    else: # No exceptions, but not all contigs\n",
    "                        print('{0} records downloaded, expected {1}'.format(len(records),\n",
    "                                                                           expected_contigs[asm_uid]))\n",
    "                        SeqIO.write(records, os.path.join(\"failures\",\n",
    "                                                          \"{0}_{1}_failed.fasta\".format(asm_uid, tries)),\n",
    "                                                          'fasta')\n",
    "                        tries += 1\n",
    "                except:  # Catch any errors, incl. from SeqIO.parse and Entrez.efetch\n",
    "                    tries += 1\n",
    "            if tries >= 10:\n",
    "                print(\"Download failed for {0}\\n\".format(labelstr))\n",
    "            print(\"UID:{0} has {1} records, total length {2}\\n\".format(asm_uid, len(records),\n",
    "                                                                       sum([len(r) for r in records])))\n",
    "            SeqIO.write(records, \"{0}.fasta\".format(filestem), 'fasta')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}