anvio_pangenome.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "25fd90ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "import os\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0da4a9f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_anvio_contig_db(taxa, anvioBasedb, genomeFilePath, genomeID, sh_cmd_db):\n",
    "    \"\"\"Generate contig databases for all genomes\"\"\"\n",
    "    # Make the directory if it does not already exist\n",
    "    if not os.path.isdir('{}{}'.format(anvioBasedb, taxa)):\n",
    "        os.mkdir('{}{}'.format(anvioBasedb, taxa))\n",
    "        os.mkdir('{}{}/fna'.format(anvioBasedb, taxa))\n",
    "        os.mkdir('{}{}/contig_db'.format(anvioBasedb, taxa))\n",
    "    # Create the contig database if it does not exist\n",
    "    if not os.path.isfile('{}{}/contig_db/{}_contigs.db'.format(anvioBasedb, taxa, genomeID)):\n",
    "        # Format the FASTA file\n",
    "        anvio_reformat_fasta_cmd = 'anvi-script-reformat-fasta {} --simplify-names --output-file {}{}/fna/{}.fna --seq-type NT'.format(genomeFilePath, anvioBasedb, taxa, genomeID)\n",
    "        # Create the contig database\n",
    "        anvio_contigdb_cmd = 'anvi-gen-contigs-database -o {}{}/contig_db/{}_contigs.db -f {}{}/fna/{}.fna'.format(anvioBasedb, taxa, genomeID, anvioBasedb, taxa, genomeID)\n",
    "        # Add the commands to the shell script\n",
    "        sh_cmd_db += anvio_reformat_fasta_cmd + '\\n'\n",
    "        sh_cmd_db += anvio_contigdb_cmd + '\\n'\n",
    "        print (anvio_contigdb_cmd)\n",
    "    return sh_cmd_db\n",
    "\n",
    "def add_to_storage(external_genome_storage_file, taxa, genome, anvioBasedb):\n",
    "    # Add the genome and its location to the file\n",
    "    external_genome_storage_file += '{}\\t{}{}/contig_db/{}_contigs.db\\n'.format(genome, anvioBasedb, taxa, genome)\n",
    "    # Return the updated file\n",
    "    return external_genome_storage_file\n",
    "\n",
    "def run_kegg_annot(sh_genome_kegg, taxa, genome, anvioBasedb):\n",
    "    # First, add the command to the shell script\n",
    "    sh_genome_kegg += 'anvi-run-kegg-kofams -T 64 -c {}{}/contig_db/{}_contigs.db\\n'.format(anvioBasedb, taxa, genome)\n",
    "    # Then return the shell script\n",
    "    return sh_genome_kegg\n",
    "\n",
    "def run_metabolism_reconstruction(sh_cmd_metabolism, taxa, genome, anvioBasedb):\n",
    "    # First, add the command to the shell script\n",
    "    sh_cmd_metabolism += 'anvi-estimate-metabolism -c {}{}/contig_db/{}_contigs.db -O {}{}/metabolism/{}-METABOLISM.txt\\n'.format(anvioBasedb, taxa, genome, anvioBasedb, taxa, genome)\n",
    "    # Then return the shell script\n",
    "    return sh_cmd_metabolism"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bf62f5b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sphingomonadaceae\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "# This code takes in a dataframe of genomes and their taxonomic annotations, and creates a set of commands for running anvio to generate the contig databases, the genomes storage, and the pangenome for each genus.\n",
    "# The output is a set of shell scripts that can be run on a cluster, with one script for each step.\n",
    "# taxa_df is a dataframe with the genome names as the index and the first column containing the genus names and the second column containing the species names.\n",
    "# fastaDir is the directory in which the fasta files are located.\n",
    "# anvioBasedb is the directory in which the anvio databases are stored.\n",
    "\n",
    "taxa_df = pd.read_csv('/DATA_RAID2/vtracann/shared/db/isolates/family_gtdbtk_annot.tsv', sep='\\t', index_col=0, header=None)\n",
    "anvioBasedb = '/DATA_RAID2/vtracann/shared/db/isolates/anvio/'\n",
    "fastaDir = '/DATA_RAID2/vtracann/shared/db/isolates/global/unzipped/'\n",
    "#print (taxa_df)\n",
    "sh_cmd_db = \"\"\n",
    "sh_cmd_storage = \"\"\n",
    "sh_cmd_pangenome = \"\"\n",
    "sh_genome_kegg = \"\"\n",
    "sh_cmd_metabolism = \"\"\n",
    "\n",
    "for taxa in set(taxa_df.loc[:, 1].values):\n",
    "    if taxa != 'unclassified':\n",
    "        genomes = taxa_df[taxa_df.loc[:, 1] == taxa]\n",
    "        if taxa == 'Sphingomonadaceae':\n",
    "            print (taxa)\n",
    "            print (len(genomes))\n",
    "\n",
    "        if len(genomes) > 10:\n",
    "            external_genome_storage_file = \"name\\tcontigs_db_path\\n\"\n",
    "            #print (taxa)\n",
    "            #print (len(genomes))\n",
    "            for genome in genomes.index:\n",
    "                genomesFileName = [x for x in os.listdir(fastaDir) if x.startswith(genome)][0]\n",
    "                genomeFilePath = fastaDir+genomesFileName\n",
    "                sh_cmd_db = make_anvio_contig_db(taxa, anvioBasedb, genomeFilePath, genome, sh_cmd_db)\n",
    "                external_genome_storage_file = add_to_storage(external_genome_storage_file, taxa, genome, anvioBasedb)\n",
    "                if not os.path.isfile('{}{}/{}_annot-GENOMES.db'.format(anvioBasedb, taxa, taxa)):\n",
    "                    sh_genome_kegg = run_kegg_annot(sh_genome_kegg, taxa, genome, anvioBasedb)\n",
    "                if not os.path.isfile('{}{}/metabolism/{}_metabolism'.format(anvioBasedb, taxa, genome)):\n",
    "                    sh_cmd_metabolism = run_metabolism_reconstruction(sh_cmd_metabolism, taxa, genome, anvioBasedb)\n",
    "                    \n",
    "            #print (len(os.listdir('{}{}/contig_db/'.format(anvioBasedb, taxa))))\n",
    "                \n",
    "            externaldbFile = open('{}{}/{}_external_db.tsv'.format(anvioBasedb, taxa, taxa), 'w')\n",
    "            externaldbFile.write(external_genome_storage_file)\n",
    "            externaldbFile.close()\n",
    "            \n",
    "            sh_cmd_storage+= \"anvi-gen-genomes-storage -e {}{}/{}_external_db.tsv -o {}{}/{}-GENOMES.db\\n\".format(anvioBasedb, taxa, taxa, anvioBasedb, taxa, taxa)\n",
    "            \n",
    "            if not os.path.isdir('{}{}/pangenome/'.format(anvioBasedb, taxa)):\n",
    "                os.mkdir('{}{}/pangenome/'.format(anvioBasedb, taxa))\n",
    "            sh_cmd_pangenome+= \"anvi-pan-genome -g {}{}/{}-GENOMES.db --mcl-inflation 1.4 --min-occurrence 2 -n {} -O {}{}/pangenome/ -T 64\\n\".format(anvioBasedb, taxa, taxa, taxa, anvioBasedb, taxa)\n",
    "            \n",
    "            if not os.path.isdir('{}{}/metabolism/'.format(anvioBasedb, taxa)):\n",
    "                os.mkdir('{}{}/metabolism/'.format(anvioBasedb, taxa))\n",
    "            \n",
    "            \n",
    "\n",
    "outFile = open('anvio_make_contigdb.sh', 'w')\n",
    "outFile.write(sh_cmd_db)\n",
    "outFile.close()\n",
    "\n",
    "outFile = open('anvio_make_genome_storage.sh', 'w')\n",
    "outFile.write(sh_cmd_storage)\n",
    "outFile.close()\n",
    "\n",
    "outFile = open('anvio_make_pangenome.sh', 'w')\n",
    "outFile.write(sh_cmd_pangenome)\n",
    "outFile.close()\n",
    "\n",
    "outFile = open('anvio_run_kegg_annot.sh', 'w')\n",
    "outFile.write(sh_genome_kegg)\n",
    "outFile.close()\n",
    "\n",
    "outFile = open('anvio_run_metabolism.sh', 'w')\n",
    "outFile.write(sh_cmd_metabolism)\n",
    "outFile.close()\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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}