Replace amplicon_pipeline.py

amplicon_pipeline.py

@@ -71,14 +71,14 @@ def load_amplicons(feature_table_index, argOptions):
     amplicons_list = [[x.split('\n')[0], x.split('\n')[1]] for x in amplicons.split('>')[1:] if x.split('\n')[0] in feature_table_index]
     mibig_amplicons = open(argOptions['MiBIGAmplicons']).read()
     mibig_amplicons_list = [[x.split('\n')[0], x.split('\n')[1]] for x in mibig_amplicons.split('>')[1:]]
-    antismasdb_amplicons = open(argOptions['antismashDBAmplicons']).read()
-    a = [[x.split('\n')[0], x.split('\n')[1]] for x in antismasdb_amplicons.split('>')[1:]]
+    antismashdb_amplicons = open(argOptions['antismashDBAmplicons']).read()
+    antismashdb_amplicons_list = [[x.split('\n')[0], x.split('\n')[1]] for x in antismashdb_amplicons.split('>')[1:]]
     if argOptions['metagenomeAmplicons']:
         metagenome_amplicons = open(argOptions['metagenomeAmplicons']).read()
         metagenome_amplicons_list = [[x.split('\n')[0], x.split('\n')[1]] for x in metagenome_amplicons.split('>')[1:]]
     else:
         metagenome_amplicons_list = []
-    return amplicons_list, mibig_amplicons_list, antismashdb_amplicon_list, metagenome_amplicons_list
+    return amplicons_list, mibig_amplicons_list, antismashdb_amplicons_list, metagenome_amplicons_list
 
 def group_similar_amplicons(feature_table, amplicon_list, argOptions):
     '''
@@ -442,21 +442,24 @@ def parse_taxonomy(amp_taxonomy, outputFolder):
     :param amp_taxonomy: 
     :return: 
     '''
-    NCBI = NCBITaxa()
-    canonical_groups = ['superkingdom','phylum','class','order','family','genus','species']
-    taxa_df = pd.DataFrame(0, index=amp_taxonomy.keys(), columns=canonical_groups)
-    for amp in amplicon_taxonomy:
-        try:
-            ptaxa = amplicon_taxonomy[amp][2:-3]
-            ancestors = NCBI.get_lineage(ptaxa)
-            #names = NCBI.get_taxid_translator(ancestors)
-            ranks = NCBI.get_rank(ancestors)
-            inv_ranks = {v: k for k, v in ranks.items()}
-            _canonical = [x for x in canonical_groups if x in inv_ranks.keys()]
-            taxa_df.loc[amp, _canonical] = [inv_ranks[x] for x in _canonical]
-        except (TypeError, ValueError) as error:
-            continue
-    pickle.dump(taxa_df, open('{}antismashdb2_taxa_df.p'.format(outputFolder),'wb'))
+    if not os.path.isfile('{}antismashdb2_taxa_df.p'.format(outputFolder)):
+        NCBI = NCBITaxa()
+        canonical_groups = ['superkingdom','phylum','class','order','family','genus','species']
+        taxa_df = pd.DataFrame(0, index=amp_taxonomy.keys(), columns=canonical_groups)
+        for amp in amplicon_taxonomy:
+            try:
+                ptaxa = amplicon_taxonomy[amp][2:-3]
+                ancestors = NCBI.get_lineage(ptaxa)
+                #names = NCBI.get_taxid_translator(ancestors)
+                ranks = NCBI.get_rank(ancestors)
+                inv_ranks = {v: k for k, v in ranks.items()}
+                _canonical = [x for x in canonical_groups if x in inv_ranks.keys()]
+                taxa_df.loc[amp, _canonical] = [inv_ranks[x] for x in _canonical]
+            except (TypeError, ValueError) as error:
+                continue
+        pickle.dump(taxa_df, open('{}antismashdb2_taxa_df.p'.format(outputFolder),'wb'))
+    else:
+        taxa_df = pickle.load(open('{}antismashdb2_taxa_df.p'.format(outputFolder),'rb'))
     return taxa_df
 
 def assign_taxonomy_to_amplicons(matches, taxa_df, outputFolder):