kb_iMAG-viz_NOT-EXECTUABLE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "\n",
    "import sys\n",
    "import time\n",
    "import pandas\n",
    "import numpy\n",
    "import os\n",
    "import re\n",
    "import argparse\n",
    "from argparse import RawTextHelpFormatter\n",
    "from pathlib import Path\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### example quick command: \n",
    "\n",
    "```kb_iMAG-viz-workflow.py -i /Applications/ResearchSoftware/kb_iMAG-viz/test/query-genomes/TARA-MAGs_Delmont-Archaea-only-2017.RAST.txt```\n",
    "\n",
    "#### example long command: \n",
    "\n",
    "```/Applications/ResearchSoftware/kb_iMAG-viz/kb_iMAG-viz-workflow.py -i /Applications/ResearchSoftware/kb_iMAG-viz/test/query-genomes/TARA-MAGs_Delmont-Archaea-only-2017.RAST.txt --taxa_level Phylum --save_master_table Yes --path_to_kb_imagviz /Applications/ResearchSoftware/kb_iMAG-viz --generate_count_tables n --dimensional_reduction_method pca --plotting_method ggplot```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check to see if python3 installed and being used; should be automatic because of the shebang, but useful if someone calls explicitly with python 2\n",
    "def check_for_py3():\n",
    "    print(\"\\n\"+\"Running check_for_py3\")\n",
    "    time_start = time.time()\n",
    "    try:\n",
    "        if sys.version_info.major != 3:\n",
    "            sys.stderr.write('\\nError: python version 3 is required - you have python version %d.\\n\\n' % sys.version_info.major)\n",
    "            sys.exit(-1)\n",
    "    except Exception:\n",
    "        sys.stderr.write(\"Unable to detect python version - assuming python 3 installed.\\n\\n\")\n",
    "    print('check_for_py3 done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#KBase-specific function to return object ID based on object name\n",
    "def get_obj_id(obj_name):\n",
    "    \"\"\"\n",
    "    Example: get_obj_id(\"MAG-QC_Archaea.SAGs.Prokka\")\n",
    "    \"\"\"\n",
    "    from biokbase.workspace.client import Workspace\n",
    "    import os\n",
    "    ws = Workspace('https://kbase.us/services/ws')\n",
    "    ws_name = os.environ['KB_WORKSPACE_ID']\n",
    "    try:\n",
    "        obj_id = ws.get_object_info3({'objects': [{'workspace': ws_name, 'name': obj_name}]})['paths'][0][0]\n",
    "        return obj_id\n",
    "    except:\n",
    "        return False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#KBase-specific function to return object data based on object ID\n",
    "def get_object_data(object_id):\n",
    "    \"\"\"\n",
    "    Fetch data from the workspace. Example1: get_object_data(get_obj_id(\"iMAG-viz_Archaea.SAGs.RAST\")); Example2: get_object_data(u'43402/2132/1')\n",
    "    \"\"\"\n",
    "    from biokbase.workspace.client import Workspace\n",
    "    ws = Workspace('https://kbase.us/services/ws')\n",
    "    return ws.get_objects([{\"ref\": object_id}])[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to extract genome annotation information from a KBase GenomeSets Object\n",
    "def extract_annotations_from_genomeset(genomeset_name, export_filename):\n",
    "    print(\"\\n\"+\"Running extract_annotations_from_genomeset\")\n",
    "    time_start = time.time()\n",
    "    from biokbase.narrative.jobs.appmanager import AppManager\n",
    "    ws = biokbase.narrative.clients.get(\"workspace\")\n",
    "    for genome in get_object_data(get_obj_id(genomeset_name)).values()[3]:\n",
    "        functionList = ws.get_objects([{'ref': genome}])[0]['data']['cdss']\n",
    "        df = []\n",
    "        for function in range(len(functionList)):\n",
    "            df.append(functionList[function]['functions'])\n",
    "        with open(export_filename, 'a') as f:\n",
    "            if get_object_data(genome).values()[0][1] != \"2528311097___MAG-QC_Archaea__Isolates.RAST\":  # this genome has something weird going on in the annotation; result is it breaks the JSON file structure, suspected ({,},') symbols, need to troubleshoot\n",
    "                f.write(get_object_data(genome).values()[0][1] + '\\t' + str(df) + '\\n')\n",
    "        f.close()\n",
    "    print('extract_annotations_from_genomeset done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to clean up annotation data and parse as a tsv-like file\n",
    "def read_and_parse_rast_annotations(query_annotation_table, output_query_annotation):\n",
    "    print(\"\\n\"+\"Running read_and_parse_rast_annotations\")\n",
    "    time_start = time.time()\n",
    "    f1 = open(query_annotation_table)\n",
    "    if os.path.exists(output_query_annotation):\n",
    "        os.remove(output_query_annotation)\n",
    "    f2 = open(output_query_annotation, \"w+\")\n",
    "    f2 = open(output_query_annotation, \"a\")\n",
    "    ftrim = f1.readlines()[1:]\n",
    "    for line in ftrim:\n",
    "        rep = {\"], [u\": \"|\", \"', u'\": \"|\", '\"': \"\", \"[[u\": \"\"}\n",
    "        rep = dict((re.escape(k), v) for k, v in rep.items())\n",
    "        pattern = re.compile(\"|\".join(rep.keys()))\n",
    "        line = pattern.sub(lambda m: rep[re.escape(m.group(0))], line)\n",
    "        for j in range(len(line.split('\\t')[1].split(\"|\")) - 1):  # last entry needs to be removed\n",
    "            if \"''\" != str(line.split('\\t')[1].split(\"|\")[j]): # write only non-empty lines\n",
    "                if \"hypothetical protein\" not in line.split('\\t')[1].split(\"|\")[j]: # not working\n",
    "                    if \"unknown\" not in line.split('\\t')[1].split(\"|\")[j]: # not working\n",
    "                        if \"conserved protein\" not in line.split('\\t')[1].split(\"|\")[j]: # not working\n",
    "                            if \"Conserved protein\" not in line.split('\\t')[1].split(\"|\")[j]: # not working\n",
    "                                if \"predicted protein\" not in line.split('\\t')[1].split(\"|\")[j]: # not working\n",
    "                                    if (str(line.split('\\t')[1].split(\"|\")[j])[0] == \"'\") and (str(line.split('\\t')[1].split(\"|\")[j])[-1] == \"'\"):\n",
    "                                        f2.write(line.split('\\t')[0].split('___')[0]+'\\t'+line.split('\\t')[1].split(\"|\")[j][1:-1]+'\\n')\n",
    "                                    elif (str(line.split('\\t')[1].split(\"|\")[j])[0] == \"'\") and (str(line.split('\\t')[1].split(\"|\")[j])[-1] != \"'\"):\n",
    "                                        f2.write(line.split('\\t')[0].split('___')[0]+'\\t'+line.split('\\t')[1].split(\"|\")[j][1:]+'\\n')\n",
    "                                    elif (str(line.split('\\t')[1].split(\"|\")[j])[0] != \"'\") and (str(line.split('\\t')[1].split(\"|\")[j])[-1] == \"'\"):\n",
    "                                        f2.write(line.split('\\t')[0].split('___')[0]+'\\t'+line.split('\\t')[1].split(\"|\")[j][:-1]+'\\n')\n",
    "                                    elif (str(line.split('\\t')[1].split(\"|\")[j])[0] != '\"') and (str(line.split('\\t')[1].split(\"|\")[j])[-1] == '\"'):\n",
    "                                        f2.write(line.split('\\t')[0].split('___')[0]+'\\t'+line.split('\\t')[1].split(\"|\")[j][1:-1]+'\\n')\n",
    "                                    else:\n",
    "                                        f2.write(line.split('\\t')[0].split('___')[0]+'\\t'+line.split('\\t')[1].split(\"|\")[j]+'\\n')\n",
    "    f1.close()\n",
    "    f2.close()\n",
    "    print('read_and_parse_rast_annotations done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to combine multiple external required tables into a single table\n",
    "def combine_external_checkm_and_taxonomy_info(query_genome_data, reference_IMG_genome_data, reference_Other_genome_data):\n",
    "    print(\"\\n\" + \"Running combine_external_checkm_and_taxonomy_info\")\n",
    "    time_start = time.time()\n",
    "    filenames = [query_genome_data, reference_IMG_genome_data, reference_Other_genome_data]\n",
    "    with open(output_combined_genome_data, 'w') as outfile:\n",
    "        for fname in filenames:\n",
    "            with open(fname) as infile:\n",
    "                for line in infile:\n",
    "                    outfile.write(line)\n",
    "    print(\"\\nTotal number of genomes in metadata table: \" + str(len(open(output_combined_genome_data).readlines())))\n",
    "    print('combine_external_checkm_and_taxonomy_info done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to merge all tables together (query and reference) into a large parseable master table used for key opertations\n",
    "def import_and_merge_tables(save_master_table):\n",
    "    print(\"\\n\" + \"Running import_and_merge_tables\")\n",
    "    time_start = time.time()\n",
    "    f1 = pandas.read_csv(query_genome_data, header=None, sep=\"\\t\", index_col=False, dtype=str)\n",
    "    global f2\n",
    "    f2 = pandas.read_csv(reference_IMG_genome_data, header=None, sep=\"\\t\", index_col=False, dtype=str)\n",
    "    f3 = pandas.read_csv(reference_Other_genome_data, header=None, sep=\"\\t\", index_col=False, dtype=str)\n",
    "    d1 = pandas.read_csv(query_isolate_annotation_data, header=None, sep=\"\\t\", index_col=False, dtype=str)\n",
    "    d2 = pandas.read_csv(query_MAG_annotation_data, header=None, sep=\"\\t\", index_col=False, dtype=str)\n",
    "    d3 = pandas.read_csv(query_SAG_annotation_data, header=None, sep=\"\\t\", index_col=False, dtype=str)\n",
    "    q1 = pandas.read_csv(query_annotation_data, header=None, sep=\"\\t\", index_col=False, dtype=str)\n",
    "    f123 = pandas.concat([f1, f2, f3], axis=0)\n",
    "    d123q1 = pandas.concat([d1, d2, d3, q1], axis=0)\n",
    "    global merge\n",
    "    merge = f123.merge(d123q1, how='outer', left_on=0, right_on=0)\n",
    "    merge.columns = ['genomeID', 'genomeSet', 'V1', 'V2', 'V3', 'V4', 'GenomeType', 'V6', 'Completeness', 'Contamination', 'Domain', 'Phylum', 'Class', 'Order', 'Family', 'Genus', 'Species', 'RAST_Annotation']\n",
    "    if save_master_table == \"Yes\":\n",
    "        merge.to_csv(\"iMAG-viz-output_ALL_genomeQC-and-annotation-data.csv\")\n",
    "    print('import_and_merge_tables done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to assist with logical selection of categorical data\n",
    "def extract_lineages_for_selected_level(taxa_level, query_genome_data):\n",
    "    print(\"\\n\"+\"Running extract_lineages_for_selected_level\")\n",
    "    time_start = time.time()\n",
    "    if taxa_level == 'Domain':\n",
    "        taxa_level_list = merge.Domain.unique()\n",
    "        taxa_level_down1 = 'Phylum'\n",
    "        taxa_level_down2 = 'Class'\n",
    "        #query_unique = f2.Domain.unique()\n",
    "    elif taxa_level == 'Phylum':\n",
    "        taxa_level_list = merge.Phylum.unique()\n",
    "        taxa_level_down1 = 'Class'\n",
    "        taxa_level_down2 = 'Order'\n",
    "        #print(query_genome_data)\n",
    "        #query_unique = query_annotation_data[11].unique()\n",
    "        #print(\"query_unique\"+query_unique)\n",
    "    elif taxa_level == 'Class':\n",
    "        taxa_level_list = merge.Class.unique()\n",
    "        taxa_level_down1 = 'Order'\n",
    "        taxa_level_down2 = 'Family'\n",
    "        #query_unique = f2.Class.unique()\n",
    "    elif taxa_level == 'Order':\n",
    "        taxa_level_list = merge.Order.unique()\n",
    "        taxa_level_down1 = 'Family'\n",
    "        taxa_level_down2 = 'Genus'\n",
    "        #query_unique = f2.Order.unique()\n",
    "    elif taxa_level == 'Family':\n",
    "        taxa_level_list = merge.Family.unique()\n",
    "        taxa_level_down1 = 'Genus'\n",
    "        taxa_level_down2 = 'Species'\n",
    "        #query_unique = f2.Family.unique()\n",
    "    else:\n",
    "        taxa_level_list = merge.Genus.unique()\n",
    "        taxa_level_down1 = 'Species'\n",
    "        taxa_level_down2 = ''\n",
    "        #query_unique = f2.Genus.unique()\n",
    "    return taxa_level_list, taxa_level_down1, taxa_level_down2\n",
    "    print('extract_lineages_for_selected_level done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Data table example (variables: merge, merge_reduced, merge_reduced_trim)\n",
    "\n",
    "# merge example\n",
    "# genomeID    genomeSet   V1  V2  V3  ... Order   Family  Genus   Species RAST_Annotation\n",
    "# 1861359 GCA_002509575.1_ASM250957v1_genomic PRJNA348753-8000Genomes y   y   y   ... o__Methanomicrobiales   f__Methanocullaceae g__Methanoculleus   s__GCA_002508705.1  Acyl-CoA    synthetase  (NDP    forming)\n",
    "# 1861360 GCA_002509575.1_ASM250957v1_genomic PRJNA348753-8000Genomes y   y   y   ... o__Methanomicrobiales   f__Methanocullaceae g__Methanoculleus   s__GCA_002508705.1  Iron(III)   dicitrate-binding   protein\n",
    "# 1861361 GCA_002509575.1_ASM250957v1_genomic PRJNA348753-8000Genomes y   y   y   ... o__Methanomicrobiales   f__Methanocullaceae g__Methanoculleus   s__GCA_002508705.1  Iron(III)   dicitrate   transport   system  permease    ...\n",
    "# 1861362 GCA_002509575.1_ASM250957v1_genomic PRJNA348753-8000Genomes y   y   y   ... o__Methanomicrobiales   f__Methanocullaceae g__Methanoculleus   s__GCA_002508705.1  Iron(III)   dicitrate   transport   ATP-binding protein\n",
    "# 1861363 GCA_002509575.1_ASM250957v1_genomic PRJNA348753-8000Genomes y   y   y   ... o__Methanomicrobiales   f__Methanocullaceae g__Methanoculleus   s__GCA_002508705.1  Diadenosine 5'5'''-P1,P4-tetraphosphate pyroph..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to subset data by lineage (aka: taxonomic level) (e.g. Phylum, Order, Family, Genus, etc.)\n",
    "\n",
    "def subset_data_by_lineage(lineage, taxa_level):\n",
    "    print(\"\\n\\t\"+\"Running subset_data_by_lineage\")\n",
    "    time_start = time.time()\n",
    "    global merge_reduced\n",
    "    global merge_reduced_trim\n",
    "    merge_reduced = merge.loc[merge[taxa_level] == lineage]\n",
    "    merge_reduced_trim = merge_reduced.drop(columns=\"RAST_Annotation\").drop_duplicates() # per genomeID, extra information.\n",
    "    print('\\tsubset_data_by_lineage done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to count annotation data for the subset lineage\n",
    "\n",
    "def count_annotation_data_for_level(merge_reduced, output_annotation_count, lineage):\n",
    "    print(\"\\n\\t\"+\"Running count_annotation_data_for_level\")\n",
    "    time_start = time.time()\n",
    "    if Path(output_annotation_count+\"_\"+lineage+\".tsv\").is_file():\n",
    "        os.remove(output_annotation_count+\"_\"+lineage+\".tsv\")\n",
    "\n",
    "    f1 = open(output_annotation_count+\"_\"+lineage+\".tsv\", \"w+\")\n",
    "    f1 = open(output_annotation_count+\"_\"+lineage+\".tsv\", \"a\")\n",
    "    genomelist = merge_reduced.genomeID.unique()\n",
    "    annotationlist = merge_reduced.RAST_Annotation.unique()\n",
    "    f1.write(\"Annotation\" + \"\\t\")  # write header line\n",
    "    for genome in range(len(genomelist)):\n",
    "        f1.write(str(genomelist[genome]))\n",
    "        if genome != (len(genomelist) - 1):\n",
    "            f1.write(\"\\t\")\n",
    "        else:\n",
    "            f1.write(\"\\n\")\n",
    "    for annotation in range(len(annotationlist)): # takes long to run (up to hours)\n",
    "    #for annotation in range(0, 20):\n",
    "        f1.write(str(annotationlist[annotation]) + \"\\t\")\n",
    "        temp_merge = merge_reduced.loc[merge_reduced['RAST_Annotation'] == annotationlist[annotation]]\n",
    "        #dat = temp_merge['genomeID'].value_counts().rename_axis('genomeID').reset_index(name='counts')\n",
    "        for genome in range(len(genomelist)):\n",
    "            dat = temp_merge['genomeID'].value_counts().rename_axis('genomeID').reset_index(name='counts')\n",
    "            f1.write(str(dat['genomeID'].astype(str).str.contains(str(genomelist[genome])).sum()))\n",
    "            if genome != (len(genomelist) - 1):\n",
    "                f1.write(\"\\t\")\n",
    "            else:\n",
    "                f1.write(\"\\n\")\n",
    "    print('\\tcount_annotation_data_for_level done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to combine annotation count table with genome metadata\n",
    "\n",
    "def import_count_and_combine_with_genome_metadata(lineage):\n",
    "    global dat\n",
    "    global dat_combined\n",
    "    global genome_number\n",
    "    with open(output_annotation_count+\"_\"+lineage+\".tsv\") as f:\n",
    "        ncols = len(f.readline().split('\\t'))\n",
    "    dat = numpy.loadtxt(output_annotation_count+\"_\"+lineage+\".tsv\", delimiter=\"\\t\", skiprows=1, usecols=range(1, ncols))\n",
    "    dat = numpy.transpose(dat)  # transpose data\n",
    "    genome_number = len(merge_reduced_trim)\n",
    "    #print(dat)\n",
    "    #print(merge_reduced_trim)\n",
    "    print(\"\\tNumber of genomes: \" + str(len(merge_reduced_trim)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to run TSNE (in dev, not really working to produce a meaningful visual)\n",
    "\n",
    "def run_tsne_dimensional_reduction():\n",
    "    print(\"\\n\\t\"+\"Running run_tsne_dimensional_reduction\")\n",
    "    time_start = time.time()\n",
    "    from sklearn.manifold import TSNE\n",
    "    tsne = TSNE(n_components=2, perplexity=40.0, early_exaggeration=12.0, learning_rate=200.0, n_iter=300, n_iter_without_progress=300, min_grad_norm=1e-07, metric=\"euclidean\", init=\"random\", verbose=1, random_state=None, method=\"barnes_hut\", angle=0.5)\n",
    "    tsne_results = tsne.fit_transform(dat)  # run tsne, fit into an embedded space\n",
    "    global dat1\n",
    "    dat1 = pandas.DataFrame(data=tsne_results, columns=['tsne-2d-one', 'tsne-2d-two'])\n",
    "    dat1_combined = pandas.DataFrame(numpy.concatenate((dat1, merge_reduced_trim),axis = 1)) # combine count table with genome metadata\n",
    "    dat1_combined.columns = ['tsne-2d-one', 'tsne-2d-two', 'genomeID', 'genomeSet', 'V1', 'V2', 'V3', 'V4', 'GenomeType', 'V6', 'Completeness', 'Contamination', 'Domain', 'Phylum', 'Class', 'Order', 'Family', 'Genus', 'Species']\n",
    "    print('\\trun_tsne_dimensional_reduction done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to run PCA (working fairly well)\n",
    "\n",
    "def run_principal_component_analysis():\n",
    "    print(\"\\n\\t\"+\"Running run_principal_component_analysis\")\n",
    "    time_start = time.time()\n",
    "    from sklearn.decomposition import PCA\n",
    "    pca = PCA(n_components=3)\n",
    "    pca_result = pca.fit_transform(dat)\n",
    "    global dat1_combined\n",
    "    dat1 = pandas.DataFrame(data=pca_result, columns=['PC1', 'PC2', 'PC3'])\n",
    "    dat1_combined = pandas.DataFrame(numpy.concatenate((dat1, merge_reduced_trim),axis = 1)) # combine count table with genome metadata\n",
    "    dat1_combined.columns = ['PC1', 'PC2', 'PC3', 'genomeID', 'genomeSet', 'V1', 'V2', 'V3', 'V4', 'GenomeType', 'V6', 'Completeness', 'Contamination', 'Domain', 'Phylum', 'Class', 'Order', 'Family', 'Genus', 'Species']\n",
    "    #print(dat1_combined)\n",
    "    #print(dat1_combined.shape)\n",
    "    dat1_combined.to_csv(\"iMAG-viz-output_PCA-input-data_\"+str(lineage)+\".tsv\", sep='\\t')\n",
    "    print('\\trun_principal_component_analysis done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to plot with Seaborn (not used at the moment)\n",
    "\n",
    "def plot_dimensional_reduction_results_seaborn(lineage, mode):\n",
    "    print(\"\\n\\t\" + \"Running plot_dimensional_reduction_results_seaborn\")\n",
    "    time_start = time.time()\n",
    "    import seaborn\n",
    "    # list_genomeid = merge_reduced[\"genomeID\"]\n",
    "    if mode == \"tsne\":\n",
    "        sns_plot = seaborn.scatterplot(x=\"tsne-2d-one\", y=\"tsne-2d-two\", hue=\"Class\", style=None, size=None, data=dat1_combined, palette=seaborn.color_palette(\"hls\", len(dat1_combined[\"Class\"].unique())), hue_order=None, hue_norm=None, sizes=None, size_order=None, size_norm=None, markers=dat1_combined[\"Class\"], style_order=None, x_bins=None, y_bins=None, units=None, estimator=None, ci=95, n_boot=1000, alpha=0.3, x_jitter=None, y_jitter=None, legend='full', ax=None)\n",
    "    else:\n",
    "        # print(type(merge_reduced[\"Class\"]))\n",
    "        df = pandas.DataFrame(merge_reduced_trim)\n",
    "        #sns_plot = seaborn.scatterplot(x=\"PC1\", y=\"PC2\", hue=df[\"Class\"], style=None, size=None, data=dat1, palette=seaborn.color_palette(\"hls\", len(df[\"Class\"].unique())), hue_order=None, hue_norm=None, sizes=None, size_order=None, size_norm=None, markers=True, style_order=None, x_bins=None, y_bins=None, units=None, estimator=None, ci=95, n_boot=1000, alpha=0.3, x_jitter=None, y_jitter=None, legend='full', ax=None)\n",
    "        sns_plot = seaborn.scatterplot(x=\"PC1\", y=\"PC2\", hue=\"Class\", style=None, size=None, data=dat1_combined, palette=seaborn.color_palette(\"hls\", len(dat1_combined[\"Class\"].unique())), hue_order=None, hue_norm=None, sizes=None, size_order=None, size_norm=None, markers=dat1_combined[\"Class\"], style_order=None, x_bins=None, y_bins=None, units=None, estimator=None, ci=95, n_boot=1000, alpha=0.3, x_jitter=None, y_jitter=None, legend='full', ax=None)\n",
    "    fig = sns_plot.get_figure()\n",
    "    fig.savefig(\"iMAG-viz-output_PCA_\"+str(lineage)+\".png\")\n",
    "    print('\\tplot_dimensional_reduction_results_seaborn done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to plot with ggplot by calling R externally\n",
    "\n",
    "def plot_dimensional_reduction_results_rggplot(lineage, mode):\n",
    "    print(\"\\n\\t\"+\"Running plot_dimensional_reduction_results_rggplot\")\n",
    "    import shutil\n",
    "    import os\n",
    "    shutil.copy(\"iMAG-viz-output_PCA-input-data_\"+str(lineage)+\".tsv\", \"R-input-data.tsv\")\n",
    "    time_start = time.time()\n",
    "    import subprocess\n",
    "    command = \"Rscript /Applications/ResearchSoftware/kb_iMAG-viz/make_ggplot_scatterplot.R\"\n",
    "    process=subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    process.wait()\n",
    "    out, err = process.communicate()\n",
    "    if process.returncode != 0: sys.exit(\"*Error generating figure with ggplot2\")\n",
    "    shutil.move(\"R-output-plot.pdf\",\"iMAG-viz-output_PCA_\"+str(lineage)+\".pdf\")\n",
    "    os.remove(\"R-input-data.tsv\")\n",
    "    print('\\tplot_dimensional_reduction_results_rggplot done! Time elapsed: ' + '{}'.format(time.time() - time_start)[:7] + ' seconds\\n')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Declare variables\n",
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(prog='kb_iMAG-viz-workflow',usage='%(prog)s.py -i [query_annotation_table] --version', description=\"\"\"\n",
    "    kb_iMAG-viz is software to evaluate microbial genome quality against reference genomes.\n",
    "    ----------------------------------------------------------------------------------------------------------------\n",
    "    kb_iMAG-viz-workflow.py performs the following steps.\n",
    "    The workflow goes as follows:\n",
    "    STEP 1. Annotations data is read from a source (e.g. RAST), and converted into a tab-delimited text\n",
    "    STEP 1a. Raw data is cleaned to remove extraneous characters\n",
    "    STEP 2. Data are subsetted based on the taxonomic level of interest (e.g. domain, phylum, class, order, ...)\n",
    "    STEP 3. Annotation data is converted to presence/absence data (for that taxonomic level)\n",
    "    STEP 4: Dimensional reduction on annotation count data to identify major trends and outliers\n",
    "    STEP 5: Plotting dimensional reduction results and associated genome quality information.\"\"\"\n",
    "    ,formatter_class=RawTextHelpFormatter)\n",
    "    #requiredNamed = parser.add_argument_group('required arguments')\n",
    "    #requiredNamed.add_argument(\"-i\", dest=\"query_annotation_table\", help=\"\"\"Indicate a tab-delimited table where genome ID is the first column and subset columns correspond to gene names.\".\"\"\", required=False)\n",
    "    parser.add_argument(\"-i\", dest=\"query_annotation_table\", help=\"\"\"Indicate a tab-delimited table where genome ID is the first column and subset columns correspond to gene names (default: /Applications/ResearchSoftware/kb_iMAG-viz/test/query-genomes/TARA-MAGs_Delmont-Archaea-only-2017.RAST.txt)\"\"\", default=\"/Applications/ResearchSoftware/kb_iMAG-viz/test/query-genomes/TARA-MAGs_Delmont-Archaea-only-2017.RAST.txt\")\n",
    "    parser.add_argument(\"--taxa_level\", dest=\"taxa_level\", help=\"\"\"Indicate the taxonomic lineage. (default: Phylum)\"\"\", default=\"Phylum\")\n",
    "    parser.add_argument(\"--save_master_table\", dest=\"save_master_table\", help=\"\"\"Save the master merged table; warning it can be large as it contains an unflattened version of all the input data. (default: Yes)\"\"\", default=\"Yes\")\n",
    "    parser.add_argument(\"--generate_count_tables\", dest=\"generate_count_tables\", help=\"\"\"Regenerate annotation count tables (time-consuming step). (default: Yes)\"\"\", default=\"Yes\")\n",
    "    parser.add_argument(\"--dimensional_reduction_method\", dest=\"dimensional_reduction_method\", help=\"\"\"Pick a dimensional reduction method. (default: pca)\"\"\", default=\"pca\")\n",
    "    parser.add_argument(\"--plotting_method\", dest=\"plotting_method\", help=\"\"\"Pick a ploting method. (default: ggplot)\"\"\", default=\"ggplot\")\n",
    "    parser.add_argument(\"--path_to_kb_imagviz\", dest=\"path_to_kb_imagviz\", help=\"\"\"Indicate the path to kb_imagviz software\"\"\", default=\"/Applications/ResearchSoftware/kb_iMAG-viz\")\n",
    "    parser.add_argument('--version', action='version', version='%(prog)s v0.1')\n",
    "    return parser.parse_args()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def banner():\n",
    "    print(\"\"\"\n",
    "__________________________________________________________\n",
    "|  _    _        _ __  __    _    ____            _      |\n",
    "| | | _| |__    (_)  \\/  |  / \\  / ___|    __   _(_)____ |\n",
    "| | |/ / '_ \\   | | |\\/| | / _ \\| |  _ ____\\ \\ / / |_  / |\n",
    "| |   <| |_) |  | | |  | |/ ___ \\ |_| |_____\\ V /| |/ /  |\n",
    "| |_|\\_\\_.__/___|_|_|  |_/_/   \\_\\____|      \\_/ |_/___| |\n",
    "|          |_____|                                       |\n",
    "|                                           version %s  |\n",
    "|________________________________________________________|\\n\\n\"\"\" % (version))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'version' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-20-46d3ea65ae6b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;31m# run main kb_iMAG-viz-workflow.py script\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"__main__\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m     \u001b[0mbanner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m     \u001b[0margs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mparse_args\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margv\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-19-3e4f146681a2>\u001b[0m in \u001b[0;36mbanner\u001b[0;34m()\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;34m|\u001b[0m          \u001b[0;34m|\u001b[0m\u001b[0m_____\u001b[0m\u001b[0;34m|\u001b[0m                                       \u001b[0;34m|\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m|\u001b[0m                                           \u001b[0mversion\u001b[0m \u001b[0;34m%\u001b[0m\u001b[0ms\u001b[0m  \u001b[0;34m|\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m |________________________________________________________|\\n\\n\"\"\" % (version))\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'version' is not defined"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "# run main kb_iMAG-viz-workflow.py script\n",
    "if __name__ == \"__main__\":\n",
    "    banner()\n",
    "    args = parse_args()\n",
    "    if len(sys.argv) is None:\n",
    "        parser.print_help()\n",
    "    elif args.query_annotation_table is None:\n",
    "        parser.print_help()\n",
    "    elif args.path_to_kb_imagviz is None:\n",
    "        parser.print_help()\n",
    "    else:\n",
    "        output_query_annotation = \"iMAG-viz-output_query-flattened-annotation-data.tsv\"\n",
    "        output_annotation_count = \"iMAG-viz-output_annotation-count-data\"\n",
    "        # files needed for function combine_external_checkm_and_taxonomy_info AND import_and_merge_tables\n",
    "        query_genome_data = os.path.realpath(args.path_to_kb_imagviz + \"/test/query-genomes/Delmont_genomeQC-data.tsv\")\n",
    "        reference_IMG_genome_data = os.path.realpath(args.path_to_kb_imagviz + \"/test/reference-genomes/IMG_genomeQC-data.tsv\")\n",
    "        reference_Other_genome_data = os.path.realpath(args.path_to_kb_imagviz + \"/test/reference-genomes/Other_genomeQC-data.tsv\")\n",
    "        # this output file contains all of the genome data used in the analysis\n",
    "        output_combined_genome_data = os.path.realpath(\"iMAG-viz-output_ALL_genomeQC-data.tsv\")\n",
    "        # files needed for function import_and_merge_tables\n",
    "        query_annotation_data = os.path.realpath(args.path_to_kb_imagviz + \"/test/query-genomes/TARA-MAGs_Delmont-Archaea-only-2017_clean.RAST.txt\")\n",
    "        query_isolate_annotation_data = os.path.realpath(args.path_to_kb_imagviz + \"/test/reference-genomes/Archaea.Isolates_clean.RAST.txt\")\n",
    "        query_MAG_annotation_data = os.path.realpath(args.path_to_kb_imagviz + \"/test/reference-genomes/Archaea.MAGs_clean.RAST.txt\")\n",
    "        query_SAG_annotation_data = os.path.realpath(args.path_to_kb_imagviz + \"/test/reference-genomes/Archaea.SAGs_clean.RAST.txt\")\n",
    "        check_for_py3()\n",
    "        read_and_parse_rast_annotations(str(args.query_annotation_table), output_query_annotation)\n",
    "        combine_external_checkm_and_taxonomy_info(query_genome_data, reference_IMG_genome_data, reference_Other_genome_data)\n",
    "        import_and_merge_tables(args.save_master_table)\n",
    "        taxalist = extract_lineages_for_selected_level(args.taxa_level, query_genome_data)[0]  # get list of lineages to iterate over\n",
    "        for lineage_number in range(len(taxalist)):\n",
    "            lineage = taxalist[lineage_number]\n",
    "            # if (str(lineage) == \"p__Nanoarchaeota\") or (str(lineage) == \"p__Micrarchaeota\") or (str(lineage) == \"p__Euryarchaeota\") or (str(lineage) == \"p__Asgardarchaeota\"):\n",
    "            # if (str(lineage) == \"p__Thermoplasmatota\"):\n",
    "            if (str(lineage) != \"nan\"): # some groups don't have phyla? need to work this out.\n",
    "                print(\"Starting with lineage: \"+str(lineage))\n",
    "                subset_data_by_lineage(lineage, args.taxa_level)\n",
    "                if (args.generate_count_tables == \"Yes\"):\n",
    "                    count_annotation_data_for_level(merge_reduced, output_annotation_count, lineage)\n",
    "                else:\n",
    "                    print(\"Skipping the count table generation (time-consuming step; assuming this has been done already or next steps will fail).\\n\")\n",
    "                import_count_and_combine_with_genome_metadata(lineage)\n",
    "                if genome_number > 3:\n",
    "                    if (args.dimensional_reduction_method == \"pca\"):\n",
    "                        run_principal_component_analysis()\n",
    "                    else:\n",
    "                        run_tsne_dimensional_reduction()\n",
    "                    if (args.plotting_method == \"ggplot\"):\n",
    "                        plot_dimensional_reduction_results_rggplot(lineage, mode=\"pca\")\n",
    "                    else:\n",
    "                        plot_dimensional_reduction_results_seaborn(lineage, mode=\"pca\")\n",
    "                else:\n",
    "                    print(\"\\nWarning: not enough data (genomes) to run a meaningful dimensional reduction. Select a different group.\")\n",
    "                print(\"Finished with lineage: \"+str(lineage)+'\\n')\n",
    "        print(\"Success! Done running kb_iMAG-viz\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}