ENH update 02_Habitat_mapping

Author: cocodyq

General_Scripts/02_Habitat_mapping/01_map_habitat.py

@@ -1,7 +1,7 @@
 '''
 Concept:
 Map habitat for all the smORFs from metaG.
-We split all the smORFs into 8 subfiles because of its large number.
+Split all the smORFs into 8 subfiles because of the large size.
 Map habitat to raw data non-redundant cluster.
 '''
 
@@ -27,23 +27,22 @@ def habitat(infile1,infile2,outpath):
     n = 0
     with open(infile1,'r',encoding = 'utf-8') as f1:
         for line in f1:
-            line = line.strip()
+            linelist = line.strip().split("\t")
             if line.startswith("sample"):
                 continue
             else:
-                linelist = line.split("\t")
                 if len(linelist) > 20:
                     if linelist[20] != "":
                         micro_host[linelist[0]] = linelist[9]+" # "+linelist[20]
                 else:
                     micro_host[linelist[0]] = linelist[9]
+    
     with lzma.open(infile2,'rt') as f2:
         for line in f2:
-            line = line.strip()
+            linelist = line.strip().split("\t")
             if line.startswith("#GMSC"):
                 continue
             else:
-                linelist = line.split("\t")
                 if n < 600000000:
                     out1.write(linelist[0]+"\t"+micro_host[linelist[1]]+"\n")
                 elif n >= 600000000 and n < 1200000000:
@@ -80,33 +79,31 @@ def map_cluster(infile1,infile2,outfile):
 
     with lzma.open(infile1,"rt") as f1:
         for line in f1:
-            line = line.strip()
-            linelist = line.split("\t",1)
+            linelist = line.strip().split("\t",1)
             if len(linelist) == 2:
                 habitat[linelist[0]] = linelist[1]
             else:
                 continue           
 
     with gzip.open(infile2,"rt") as f2:
         for line in f2:
-            line = line.strip()
-            linelist = line.split("\t") 
+            linelist = line.strip().split("\t") 
             number = int(''.join(linelist[1].split(".")[2].split("_")))
             if number < 34617405: #The habitat of smORFs from Progenome is named isolate.
-                out.write(linelist[0]+"\t"+linelist[1]+"\t"+"isolate"+"\n")
+                out.write(f'{linelist[0]}\t{linelist[1]}\tisolate\n')
             else:
                 if linelist[1] in habitat.keys():
-                    out.write(linelist[0]+"\t"+linelist[1]+"\t"+habitat[linelist[1]]+"\n")
+                    out.write(f'{linelist[0]}\t{linelist[1]}\t{habitat[linelist[1]]}\n')
                 else:
-                    out.write(line+"\n")           
+                    out.write(f'{line}')           
     out.close()
 
 
-INPUT_FILE_1 = "./habitat/metadata.tsv"
+INPUT_FILE_1 = "metadata.tsv"
 INPUT_FILE_2 = "GMSC10.metag_smorfs.rename.txt.xz"
 INPUT_FILE_3 = "dedup_cluster.tsv.gz"
-OUT_PATH_1 = "./habitat/metag_habitat"
-OUT_PATH_2 = "./habitat/metag_cluster_habitat" 
+OUT_PATH_1 = "metag_habitat"
+OUT_PATH_2 = "metag_cluster_habitat" 
 
 habitat(INPUT_FILE_1,INPUT_FILE_2,OUT_PATH_1)
 

General_Scripts/02_Habitat_mapping/02_multi_habitat.py

@@ -15,8 +15,7 @@ def multi_habitat(infile,outfile):
 
     with lzma.open(infile,"rt") as f1:
         for line in f1:
-            line = line.strip()
-            cluster,metag,habitat = line.split("\t")
+            cluster,metag,habitat = line.strip().split("\t")
             if cluster_dict:
                 if cluster in cluster_dict:
                     cluster_dict[cluster].append(metag)
@@ -25,7 +24,7 @@ def multi_habitat(infile,outfile):
                     multihabitat = ",".join(sorted(list(habitat_set)))
                     for key,value in cluster_dict.items():
                         for smorf in value:
-                            out.write(key+"\t"+smorf+"\t"+multihabitat+"\n")
+                            out.write(f'{key}\t{smorf}\t{multihabitat}\n')
                     cluster_dict = {}
                     habitat_set = set()      
                     cluster_dict[cluster] = [metag]
@@ -37,60 +36,57 @@ def multi_habitat(infile,outfile):
 
 def extract(infile1,infile2,outfile):
     import lzma
+    import gzip
     smorf = set()
     out = lzma.open(outfile, "wt")
 
-    with lzma.open(infile1,"rt") as f1:
+    with gzip.open(infile1,"rt") as f1:
         for line in f1:
-            line = line.strip()
-            linelist = line.split("\t")
-            smorf.add(linelist[0])
+            member,cluster = line.strip().split("\t")
+            smorf.add(member)
 
     with lzma.open(infile2,"rt") as f2:
         for line in f2:
-            line = line.strip()
-            linelist = line.split("\t")
-            if linelist[1] in smorf:
-                out.write(linelist[1]+"\t"+linelist[2]+"\n")
+            seq,habitat = line.strip().split("\t")
+            if seq in smorf:
+                out.write(f'{seq}\t{habitat}\n')
             else:
                 continue
-
     out.close()
 
 def general(infile1,infile2,outfile):
     import lzma
     out = lzma.open(outfile,"wt")
     env = {}
+
     with open(infile1,"rt") as f1:
         for line in f1:
-            line = line.strip()
-            sample,amp,microontology,name,host,habitat = line.split("\t")
+            sample,amp,microontology,name,host,habitat = line.strip().split("\t")
             if host != "":
                 fullhabitat = microontology + " # " + host
             else:
                 fullhabitat = microontology
             env[fullhabitat] = habitat
         env["isolate"] = "isolate"
+
     with lzma.open(infile2,"rt") as f2:
         for line in f2:
-            line = line.strip()
-            smorf,multihabitat = line.split("\t")
+            smorf,multihabitat = line.strip().split("\t")
             multilist = multihabitat.split(",")
             change = set()
             for i in multilist:
                 i = i.replace("-"," ")
                 change.add(env[i])
             generalhabitat = ",".join(sorted(list(change)))
-            out.write(smorf+"\t"+generalhabitat+"\n")
-            
+            out.write(f'{smorf}\t{generalhabitat}\n')
     out.close()
 
-INPUT_FILE_1 = "./habitat/metag_cluster_habitat.tsv.xz"
-INPUT_FILE_2 = "./frozen/100AA_rename.tsv.xz"
-INPUT_FILE_3 = "./habitat/habitat_general.txt"
-OUTPUT_FILE_1 = "./habitat/all_cluster_multi_habitat.tsv.xz"
-OUTPUT_FILE_2 = "./habitat/id100/100AA_multi_habitat.tsv.xz" 
-OUTPUT_FILE_3 = "./habitat/id100/100AA_multi_general_habitat.tsv.xz"
+INPUT_FILE_1 = "metag_cluster_habitat.tsv.xz"
+INPUT_FILE_2 = "GMSC.cluster.tsv.gz"
+INPUT_FILE_3 = "habitat_general.txt"
+OUTPUT_FILE_1 = "all_cluster_multi_habitat.tsv.xz"
+OUTPUT_FILE_2 = "100AA_multi_habitat.tsv.xz" 
+OUTPUT_FILE_3 = "100AA_multi_general_habitat.tsv.xz"
 
 multi_habitat(INPUT_FILE_1,OUTPUT_FILE_1)
 extract(INPUT_FILE_2,OUTPUT_FILE_1,OUTPUT_FILE_2)

General_Scripts/02_Habitat_mapping/03_map_cluster_habitat.py

@@ -1,57 +1,96 @@
 '''
 Concept:
 Map habitat to 90% identity clusters.
+Combine multiple habitats for each smORF from the same cluster.
+Change habitats to general name.
 '''
 
 import gzip
 import lzma
 
-'''
-Change format of 90% identity clusters including two columns(90AA clusters and 100AA name) 
-'''
-def change_format_90(infile1,outfile):
-    cluster = {}
+def mapcluster(infile1,infile2,outfile):
+    habitat = {}
     out = lzma.open(outfile, "wt")
-    with gzip.open(infile1,"rt") as f1:
+    with lzma.open(infile1,"rt") as f1:
+        for line in f1:
+            smorf,habitat = line.strip().split("\t",1)
+            habitat[smorf] = habitat
+
+    with gzip.open(infile2,"rt") as f2:
+        for line in f2:
+            member,cluster = line.strip().split("\t") 
+            out.write(f'{cluster}\t{member}\t{habitat[member]}\n')
+    out.close()
+
+def multi_habitat(infile,outfile):
+    import lzma
+
+    cluster_dict = {}
+    habitat_set = set()
+    out  = lzma.open(outfile, "wt")
+    
+    with lzma.open(infile,"rt") as f1:
         for line in f1:
-            line = line.strip()
-            linelist = line.split("\t")
-            if len(linelist) == 3:
-                if linelist[2] in cluster.keys():
-                    cluster[linelist[2]].append(linelist[0])
+            cluster,metag,habitat = line.strip().split("\t")
+            if cluster_dict:
+                if cluster in cluster_dict:
+                    cluster_dict[cluster].append(metag)
+                    habitatlist = habitat.split(",")
+                    for h in habitatlist:
+                        habitat_set.add(h)
                 else:
-                    cluster[linelist[2]] = [linelist[0]]
+                    multihabitat = ",".join(sorted(list(habitat_set)))
+                    for key,value in cluster_dict.items():
+                        for smorf in value:
+                            out.write(f'{key}\t{smorf}\t{multihabitat}\n')
+                    cluster_dict = {}
+                    habitat_set = set()      
+                    cluster_dict[cluster] = [metag]
+                    habitatlist = habitat.split(",")
+                    for h in habitatlist:
+                        habitat_set.add(h)
             else:
-                cluster[linelist[0]] = [linelist[0]]
-    for key,value in cluster.items():
-        for i in range(len(value)):
-            out.write(key+"\t"+value[i]+"\n")
-    out.close()
+                cluster_dict[cluster] = [metag]
+                habitatlist = habitat.split(",")
+                for h in habitatlist:
+                    habitat_set.add(h)
+    out.close() 
 
-'''
-Map habitat to 90% identity clusters.
-'''
-def mapcluster(infile1,infile2,outfile):
-    habitat = {}
-    out = lzma.open(outfile, "wt")
-    with lzma.open(infile1,"rt") as f1:
+def general(infile1,infile2,outfile):
+    import lzma
+    out = lzma.open(outfile,"wt")
+    env = {}
+
+    with open(infile1,"rt") as f1:
         for line in f1:
-            line = line.strip()
-            linelist = line.split("\t",1)
-            habitat[linelist[0]] = linelist[1]
+            sample,amp,microontology,name,host,habitat = line.strip().split("\t")
+            if host != "":
+                fullhabitat = microontology + " # " + host
+            else:
+                fullhabitat = microontology
+            env[fullhabitat] = habitat
+        env["isolate"] = "isolate"
 
     with lzma.open(infile2,"rt") as f2:
         for line in f2:
-            line = line.strip()
-            linelist = line.split("\t") 
-            out.write(linelist[0]+"\t"+linelist[1]+"\t"+habitat[linelist[1]]+"\n")
-
+            smorf_cluster,smorf,multihabitat = line.strip().split("\t")
+            multilist = multihabitat.split(",")
+            change = set()
+            for i in multilist:
+                i = i.replace("-"," ")
+                change.add(env[i])
+            generalhabitat = ",".join(sorted(list(change)))
+            out.write(f'{smorf_cluster}\t{smorf}\t{generalhabitat}\n')
+            
     out.close()
 
-INPUT_FILE_1 = "./clust_result/result/all_0.5_0.9.tsv.gz" 
-INPUT_FILE_2 = "./habitat/id100/100AA_multi_general_habitat.tsv.xz"   
-OUTPUT_FILE_1 = "all_cluster_0.9.tsv.xz"  
-OUTPUT_FILE_2 = "./habitat/id90/cluster_multi_habitat_90.tsv.xz" 
+INPUT_FILE_1 = "100AA_multi_general_habitat.tsv.xz"   
+INPUT_FILE_2 = "GMSC.cluster.tsv.gz"
+INPUT_FILE_3 = "habitat_general.txt"
+OUTPUT_FILE_1 = "cluster_multi_habitat_90.tsv.xz" 
+OUTPUT_FILE_2 = "90AA_multi_habitat.tsv.xz"
+OUTPUT_FILE_3 = "90AA_multi_general_habitat.tsv.xz"
 
-change_format_90(INPUT_FILE_1,OUTPUT_FILE_1)
-mapcluster(INPUT_FILE_2,OUTPUT_FILE_1,OUTPUT_FILE_2)
+mapcluster(INPUT_FILE_1,INPUT_FILE_2,OUTPUT_FILE_1)
+multi_habitat(OUTPUT_FILE_1,OUTPUT_FILE_2)
+general(INPUT_FILE_3,OUTPUT_FILE_2,OUTPUT_FILE_3)