diff small_rna_map.py @ 3:2e0dc6032a98 draft

planemo upload for repository https://github.com/ARTbio/tools-artbio/tree/master/tools/small_rna_map commit 93f212712d9846c7aaa389de60babb332d38363e

--- a/small_rna_map.py	Wed Jul 12 13:40:36 2017 -0400
+++ b/small_rna_map.py	Tue Jul 18 13:34:36 2017 -0400
@@ -1,84 +1,80 @@
-import os
-from os.path import basename
+import argparse
+from collections import defaultdict
+
+import numpy
+
 import pysam
-import argparse
-import numpy
-import collections
 
 
 def Parser():
     the_parser = argparse.ArgumentParser()
-    the_parser.add_argument('--input', dest="input", required=True, nargs='+', help='input BAM files')
-    the_parser.add_argument('--output', action="store", type=str, help="output tabular file")
+    the_parser.add_argument('--input', dest='input', required=True,
+                            nargs='+', help='input BAM files')
+    the_parser.add_argument('--sample_name', dest='sample_name',
+                            required=True, nargs='+', help='sample name')
+    the_parser.add_argument('--output', action="store",
+                            type=str, help='output tabular file')
     args = the_parser.parse_args()
     return args
 
 
-def Ref_length(bamfile):
-    """ 
-    bamfile.references: list of reference names, and bamfile.lengths is list of their lengths 
-    """
-    ref_length={}
-    for e1, e2 in zip(bamfile.references, bamfile.lengths):
-        ref_length[e1]=e2
-    return ref_length
-
-
-def Ref_coord_pol(bamfile):
-    ref_coord ={}
-    rnames =list(bamfile.references)
-    for read in bamfile:
+def compute_map_dictionary(pysam_object):
+    '''
+    returns
+    - map_dictionary {(read_chromosome,read_position,polarity):
+                     [list_of_read_length]}
+    - max_dictionary {chromosome_name:
+                      max_of_number_of_read_at_any_position}
+    '''
+    unmatched_chromosomes = list(pysam_object.references)
+    map_dictionary = defaultdict(list)
+    max_dictionary = dict.fromkeys(pysam_object.references, 0)
+    for read in pysam_object:
         if not read.is_unmapped:
-            polarity = "F" if (not read.is_reverse) else "R"
-            read_rname = read.reference_name
+            read_chromosome = read.reference_name
+            # to shift from 0-based to 1-based coordinate
             read_position = read.pos + 1
-            the_key = (read_rname, read_position,polarity)
-            read_length = read.qlen
-            try:
-                ref_coord[the_key][0] = ref_coord[the_key][0]+1
-                ref_coord[the_key][1].append(read_length)
-            except:
-                ref_coord[the_key] = [1,[read_length]]
-            if read_rname in rnames:
-                rnames.remove(read_rname)
-    for rname in rnames:
-        ref_coord[(rname, 0, "F")]=[0,[0]]
-    return collections.OrderedDict(sorted(ref_coord.items())) #dictionary {(read_rname,read_position,polarity):[nmbr_read,list_of_alignment_length]}
+            polarity = "F" if (not read.is_reverse) else "R"
+            map_dictionary[(read_chromosome, read_position,
+                            polarity)].append(read.qlen)
+            if read_chromosome in unmatched_chromosomes:
+                unmatched_chromosomes.remove(read_chromosome)
+    for name in unmatched_chromosomes:
+        # to put at least a null value by chromosome
+        map_dictionary[(name, 1, 'F')] = [0]
+    # make a dictionary with max value of reads by chromosome
+    max_dictionary = defaultdict(int)
+    for name in pysam_object.references:
+        for key in map_dictionary.keys():
+            if name in key and len(map_dictionary[key]) > max_dictionary[name]:
+                max_dictionary[name] = len(map_dictionary[key])
+    return map_dictionary, max_dictionary
 
 
-def calcul_stat(dictionary,ref):
-    dico = {}	
-    for keys in dictionary:
-        key = keys[0] #reference name
-        value = dictionary[keys][0] # position
-        try:
-            dico[key].append(value)
-        except: 
-            dico[key] = []
-            dico[key].append(value)
-    the_max = max(dico[ref])
-    return the_max
-
-
-def filename(infile): 
-    filename = basename(infile.filename)
-    the_name = filename.split('.')
-    return the_name[0]
-
-
-def main(infile, output):
-   with open(output,"w") as outfile:
-       outfile.write("Dataset" + "\t" + "Chromosome"+ "\t" + "Chrom_length" + "\t" + "Coordinate" + "\t" + "Nbr_reads" + "\t" + "Polarity" + "\t" + "Max" + "\t" + "Mean" + "\t" + "Median" + "\n")
-       for bamfile in infile:
-            with pysam.AlignmentFile(bamfile,"rb") as bamfile:
-                ref_coord = Ref_coord_pol(bamfile)
-                ref_length = Ref_length(bamfile)
-                tabulation = "\t"
-                for read in ref_coord:
-                    table = (str(filename(bamfile)), str(read[0]), str(ref_length[str(read[0])]), str(read[1]), str(ref_coord[read][0]), str(read[2]), str(calcul_stat(ref_coord,read[0])), str(round(numpy.mean(ref_coord[read][1]),1)), str(numpy.median(ref_coord[read][1])))
-                    outfile.write(tabulation.join(table) + "\n")
+def main(inputs, sample_names, output):
+    with open(output, 'w') as outfile:
+        outfile.write('\t'.join(['Dataset', 'Chromosome', 'Chrom_length',
+                                 'Coordinate', 'Nbr_reads', 'Polarity',
+                                 'Max', 'Mean', 'Median\n']))
+        chr_lengths = dict(zip(pysam.AlignmentFile(inputs[0], 'rb').references,
+                               pysam.AlignmentFile(inputs[0], 'rb').lengths))
+        for file, sample in zip(inputs, sample_names):
+            with pysam.AlignmentFile(file, 'rb') as bamfile:
+                map, maximum = compute_map_dictionary(bamfile)
+                for record in sorted(map):
+                    line = [sample, record[0],
+                            str(chr_lengths[record[0]]),
+                            str(record[1]), str(len(map[record])),
+                            record[2], str(maximum[record[0]]),
+                            str(round(numpy.mean(map[record]), 1)),
+                            str(numpy.median(map[record]))]
+                    outfile.write("\t".join(line) + "\n")
 
 
 if __name__ == "__main__":
-    args= Parser()
-    main(args.input, args.output)
+    args = Parser()
+    # if identical sample names # to be tested
+    if len(set(args.sample_name)) != len(args.sample_name):
+        args.sample_name = [name + '_' + str(i) for
+                            i, name in enumerate(args.sample_name)]
+    main(args.input, args.sample_name, args.output)