--- a/mismatch_frequencies.py	Wed Mar 23 09:59:33 2016 -0400
+++ b/mismatch_frequencies.py	Sat Dec 22 04:15:47 2018 -0500
@@ -1,28 +1,33 @@
-import pysam, re, string
-import matplotlib.pyplot as plt
+import re
+import string
+import pysam
+import matplotlib
 import pandas as pd
-import json
 from collections import defaultdict
 from collections import OrderedDict
 import argparse
 import itertools
 
+matplotlib.use('pdf')
+import matplotlib.pyplot as plt  # noqa: E402
+
+
 class MismatchFrequencies:
     '''Iterate over a SAM/BAM alignment file, collecting reads with mismatches. One
     class instance per alignment file. The result_dict attribute will contain a
     nested dictionary with name, readlength and mismatch count.'''
-    def __init__(self, result_dict={}, alignment_file=None, name="name", minimal_readlength=21, 
+    def __init__(self, result_dict={}, alignment_file=None, name="name", minimal_readlength=21,
                  maximal_readlength=21,
-                 number_of_allowed_mismatches=1, 
-                 ignore_5p_nucleotides=0, 
+                 number_of_allowed_mismatches=1,
+                 ignore_5p_nucleotides=0,
                  ignore_3p_nucleotides=0,
-                 possible_mismatches = [
+                 possible_mismatches=[
                         'AC', 'AG', 'AT',
                         'CA', 'CG', 'CT',
                         'GA', 'GC', 'GT',
                         'TA', 'TC', 'TG'
                 ]):
-    
+
         self.result_dict = result_dict
         self.name = name
         self.minimal_readlength = minimal_readlength
@@ -31,20 +36,19 @@
         self.ignore_5p_nucleotides = ignore_5p_nucleotides
         self.ignore_3p_nucleotides = ignore_3p_nucleotides
         self.possible_mismatches = possible_mismatches
-        
+
         if alignment_file:
             self.pysam_alignment = pysam.Samfile(alignment_file)
-            self.references = self.pysam_alignment.references #names of fasta reference sequences
-            result_dict[name]=self.get_mismatches(
-                self.pysam_alignment, 
-                minimal_readlength, 
+            self.references = self.pysam_alignment.references  # names of fasta reference sequences
+            result_dict[name] = self.get_mismatches(
+                self.pysam_alignment,
+                minimal_readlength,
                 maximal_readlength,
                 possible_mismatches
             )
-    
-    def get_mismatches(self, pysam_alignment, minimal_readlength, 
+
+    def get_mismatches(self, pysam_alignment, minimal_readlength,
                        maximal_readlength, possible_mismatches):
-        mismatch_dict = defaultdict(int)
         rec_dd = lambda: defaultdict(rec_dd)
         len_dict = rec_dd()
         for alignedread in pysam_alignment:
@@ -56,8 +60,8 @@
                     len_dict[int(alignedread.rlen)][chromosome]['total valid reads'] = 1
                 MD = alignedread.opt('MD')
                 if self.read_has_mismatch(alignedread, self.number_of_allowed_mismatches):
-                    (ref_base, mismatch_base)=self.read_to_reference_mismatch(MD, alignedread.seq, alignedread.is_reverse)
-                    if ref_base == None:
+                    (ref_base, mismatch_base) = self.read_to_reference_mismatch(MD, alignedread.seq, alignedread.is_reverse)
+                    if not ref_base:
                             continue
                     else:
                         for i, base in enumerate(ref_base):
@@ -68,7 +72,7 @@
                             except TypeError:
                                 len_dict[int(alignedread.rlen)][chromosome][ref_base[i]+mismatch_base[i]] = 1
         return len_dict
-    
+
     def read_is_valid(self, read, min_readlength, max_readlength):
         '''Filter out reads that are unmatched, too short or
         too long or that contian insertions'''
@@ -80,17 +84,17 @@
             return False
         else:
             return True
-    
+
     def read_has_mismatch(self, read, number_of_allowed_mismatches=1):
         '''keep only reads with one mismatch. Could be simplified'''
-        NM=read.opt('NM')
-        if NM <1: #filter out reads with no mismatch
+        NM = read.opt('NM')
+        if NM < 1:  # filter out reads with no mismatch
             return False
-        if NM >number_of_allowed_mismatches: #filter out reads with more than 1 mismtach
+        if NM > number_of_allowed_mismatches:  # filter out reads with more than 1 mismtach
             return False
         else:
             return True
-        
+
     def mismatch_in_allowed_region(self, readseq, mismatch_position):
         '''
         >>> M = MismatchFrequencies()
@@ -108,29 +112,29 @@
         >>> M.mismatch_in_allowed_region(readseq, mismatch_position)
         False
         '''
-        mismatch_position+=1 # To compensate for starting the count at 0
+        mismatch_position += 1  # To compensate for starting the count at 0
         five_p = self.ignore_5p_nucleotides
         three_p = self.ignore_3p_nucleotides
         if any([five_p > 0, three_p > 0]):
-            if any([mismatch_position <= five_p, 
-                    mismatch_position >= (len(readseq)+1-three_p)]): #Again compensate for starting the count at 0
+            if any([mismatch_position <= five_p,
+                    mismatch_position >= (len(readseq) + 1 - three_p)]):  # Again compensate for starting the count at 0
                 return False
             else:
                 return True
         else:
             return True
-            
+
     def read_to_reference_mismatch(self, MD, readseq, is_reverse):
         '''
         This is where the magic happens. The MD tag contains SNP and indel information,
         without looking to the genome sequence. This is a typical MD tag: 3C0G2A6.
         3 bases of the read align to the reference, followed by a mismatch, where the
-        reference base is C, followed by 10 bases aligned to the reference. 
+        reference base is C, followed by 10 bases aligned to the reference.
         suppose a reference 'CTTCGATAATCCTT'
                              |||  || ||||||
-                 and a read 'CTTATATTATCCTT'. 
-        This situation is represented by the above MD tag. 
-        Given MD tag and read sequence this function returns the reference base C, G and A, 
+                 and a read 'CTTATATTATCCTT'.
+        This situation is represented by the above MD tag.
+        Given MD tag and read sequence this function returns the reference base C, G and A,
         and the mismatched base A, T, T.
         >>> M = MismatchFrequencies()
         >>> MD='3C0G2A7'
@@ -140,82 +144,87 @@
         True
         >>> result[1]=="ATT"
         True
-        >>> 
+        >>>
         '''
-        search=re.finditer('[ATGC]',MD)
+        search = re.finditer('[ATGC]', MD)
         if '^' in MD:
             print 'WARNING insertion detected, mismatch calling skipped for this read!!!'
             return (None, None)
-        start_index=0 # refers to the leading integer of the MD string before an edited base
-        current_position=0 # position of the mismatched nucleotide in the MD tag string
-        mismatch_position=0 # position of edited base in current read 
-        reference_base=""
-        mismatched_base=""
+        start_index = 0  # refers to the leading integer of the MD string before an edited base
+        current_position = 0  # position of the mismatched nucleotide in the MD tag string
+        mismatch_position = 0  # position of edited base in current read
+        reference_base = ""
+        mismatched_base = ""
         for result in search:
-            current_position=result.start()
-            mismatch_position=mismatch_position+1+int(MD[start_index:current_position]) #converts the leading characters before an edited base into integers
-            start_index=result.end()
-            reference_base+=MD[result.end()-1]
-            mismatched_base+=readseq[mismatch_position-1]
+            current_position = result.start()
+            mismatch_position = mismatch_position + 1 + int(MD[start_index:current_position])  # converts the leading characters before an edited base into integers
+            start_index = result.end()
+            reference_base += MD[result.end() - 1]
+            mismatched_base += readseq[mismatch_position - 1]
         if is_reverse:
-            reference_base=reverseComplement(reference_base)
-            mismatched_base=reverseComplement(mismatched_base)
-            mismatch_position=len(readseq)-mismatch_position-1
-        if mismatched_base=='N':
+            reference_base = reverseComplement(reference_base)
+            mismatched_base = reverseComplement(mismatched_base)
+            mismatch_position = len(readseq)-mismatch_position-1
+        if mismatched_base == 'N':
             return (None, None)
         if self.mismatch_in_allowed_region(readseq, mismatch_position):
             return (reference_base, mismatched_base)
         else:
             return (None, None)
 
+
 def reverseComplement(sequence):
     '''do a reverse complement of DNA base.
     >>> reverseComplement('ATGC')=='GCAT'
     True
-    >>> 
+    >>>
     '''
-    sequence=sequence.upper()
+    sequence = sequence.upper()
     complement = string.maketrans('ATCGN', 'TAGCN')
     return sequence.upper().translate(complement)[::-1]
 
+
 def barplot(df, library, axes):
-    df.plot(kind='bar', ax=axes, subplots=False,\
-            stacked=False, legend='test',\
+    df.plot(kind='bar', ax=axes, subplots=False,
+            stacked=False, legend='test',
             title='Mismatch frequencies for {0}'.format(library))
-    
+
+
 def df_to_tab(df, output):
     df.to_csv(output, sep='\t')
 
+
 def reduce_result(df, possible_mismatches):
     '''takes a pandas dataframe with full mismatch details and
     summarises the results for plotting.'''
     alignments = df['Alignment_file'].unique()
     readlengths = df['Readlength'].unique()
-    combinations = itertools.product(*[alignments, readlengths]) #generate all possible combinations of readlength and alignment files
+    combinations = itertools.product(*[alignments, readlengths])  # generate all possible combinations of readlength and alignment files
     reduced_dict = {}
-    frames = []
-    last_column = 3+len(possible_mismatches)
+    last_column = 3 + len(possible_mismatches)
     for combination in combinations:
         library_subset = df[df['Alignment_file'] == combination[0]]
         library_readlength_subset = library_subset[library_subset['Readlength'] == combination[1]]
-        sum_of_library_and_readlength = library_readlength_subset.iloc[:,3:last_column+1].sum()
-        if not reduced_dict.has_key(combination[0]):
+        sum_of_library_and_readlength = library_readlength_subset.iloc[:, 3:last_column+1].sum()
+        if combination[0] not in reduced_dict:
             reduced_dict[combination[0]] = {}
         reduced_dict[combination[0]][combination[1]] = sum_of_library_and_readlength.to_dict()
     return reduced_dict
 
+
 def plot_result(reduced_dict, args):
-    names=reduced_dict.keys()
-    nrows=len(names)/2+1
-    fig = plt.figure(figsize=(16,32))
-    for i,library in enumerate (names):
-        axes=fig.add_subplot(nrows,2,i+1)
-        library_dict=reduced_dict[library]
-        df=pd.DataFrame(library_dict)
+    names = reduced_dict.keys()
+    nrows = len(names) / 2 + 1
+    fig = plt.figure(figsize=(16, 32))
+    for i, library in enumerate(names):
+        axes = fig.add_subplot(nrows, 2, i+1)
+        library_dict = reduced_dict[library]
+        df = pd.DataFrame(library_dict)
         df.drop(['total aligned reads'], inplace=True)
         barplot(df, library, axes),
         axes.set_ylabel('Mismatch count / all valid reads * readlength')
-    fig.savefig(args.output_pdf, format='pdf')    
+    fig.savefig(args.output_pdf, format='pdf')
+
 
 def format_result_dict(result_dict, chromosomes, possible_mismatches):
     '''Turn nested dictionary into preformatted tab seperated lines'''
@@ -237,35 +246,39 @@
                     line.extend([0])
                 result.append(line)
     df = pd.DataFrame(result, columns=header.split('\t'))
-    last_column=3+len(possible_mismatches)
-    df['mismatches/per aligned nucleotides'] = df.iloc[:,3:last_column].sum(1)/(df.iloc[:,last_column]*df['Readlength'])
+    last_column = 3 + len(possible_mismatches)
+    df['mismatches/per aligned nucleotides'] = df.iloc[:, 3:last_column].sum(1)/(df.iloc[:, last_column] * df['Readlength'])
     return df
-  
+
+
 def setup_MismatchFrequencies(args):
-    resultDict=OrderedDict()
-    kw_list=[{'result_dict' : resultDict, 
-             'alignment_file' :alignment_file, 
-             'name' : name, 
-             'minimal_readlength' : args.min, 
-             'maximal_readlength' : args.max,
-             'number_of_allowed_mismatches' : args.n_mm,
-             'ignore_5p_nucleotides' : args.five_p, 
-             'ignore_3p_nucleotides' : args.three_p,
-             'possible_mismatches' : args.possible_mismatches }
-             for alignment_file, name in zip(args.input, args.name)]
+    resultDict = OrderedDict()
+    kw_list = [{'result_dict': resultDict,
+                'alignment_file': alignment_file,
+                'name': name,
+                'minimal_readlength': args.min,
+                'maximal_readlength': args.max,
+                'number_of_allowed_mismatches': args.n_mm,
+                'ignore_5p_nucleotides': args.five_p,
+                'ignore_3p_nucleotides': args.three_p,
+                'possible_mismatches': args.possible_mismatches}
+               for alignment_file, name in zip(args.input, args.name)]
     return (kw_list, resultDict)
 
+
 def nested_dict_to_df(dictionary):
     dictionary = {(outerKey, innerKey): values for outerKey, innerDict in dictionary.iteritems() for innerKey, values in innerDict.iteritems()}
-    df=pd.DataFrame.from_dict(dictionary).transpose()
+    df = pd.DataFrame.from_dict(dictionary).transpose()
     df.index.names = ['Library', 'Readlength']
     return df
 
+
 def run_MismatchFrequencies(args):
-    kw_list, resultDict=setup_MismatchFrequencies(args)
+    kw_list, resultDict = setup_MismatchFrequencies(args)
     references = [MismatchFrequencies(**kw_dict).references for kw_dict in kw_list]
     return (resultDict, references[0])
 
+
 def main():
     result_dict, references = run_MismatchFrequencies(args)
     df = format_result_dict(result_dict, references, args.possible_mismatches)
@@ -273,12 +286,12 @@
     plot_result(reduced_dict, args)
     reduced_df = nested_dict_to_df(reduced_dict)
     df_to_tab(reduced_df, args.output_tab)
-    if not args.expanded_output_tab == None:
+    if args.expanded_output_tab:
         df_to_tab(df, args.expanded_output_tab)
     return reduced_dict
 
 if __name__ == "__main__":
-    
+
     parser = argparse.ArgumentParser(description='Produce mismatch statistics for BAM/SAM alignment files.')
     parser.add_argument('--input', nargs='*', help='Input files in SAM/BAM format')
     parser.add_argument('--name', nargs='*', help='Name for input file to display in output file. Should have same length as the number of inputs')
@@ -286,15 +299,13 @@
     parser.add_argument('--output_tab', help='Output filename for table')
     parser.add_argument('--expanded_output_tab', default=None, help='Output filename for table')
     parser.add_argument('--possible_mismatches', default=[
-            'AC', 'AG', 'AT','CA', 'CG', 'CT', 'GA', 'GC', 'GT', 'TA', 'TC', 'TG'
+            'AC', 'AG', 'AT', 'CA', 'CG', 'CT', 'GA', 'GC', 'GT', 'TA', 'TC', 'TG'
         ], nargs='+', help='specify mismatches that should be counted for the mismatch frequency. The format is Reference base -> observed base, eg AG for A to G mismatches.')
     parser.add_argument('--min', '--minimal_readlength', type=int, help='minimum readlength')
     parser.add_argument('--max', '--maximal_readlength', type=int, help='maximum readlength')
     parser.add_argument('--n_mm', '--number_allowed_mismatches', type=int, default=1, help='discard reads with more than n mismatches')
     parser.add_argument('--five_p', '--ignore_5p_nucleotides', type=int, default=0, help='when calculating nucleotide mismatch frequencies ignore the first N nucleotides of the read')
     parser.add_argument('--three_p', '--ignore_3p_nucleotides', type=int, default=1, help='when calculating nucleotide mismatch frequencies ignore the last N nucleotides of the read')
-    #args = parser.parse_args(['--input', '3mismatches_ago2ip_s2.bam', '3mismatches_ago2ip_ovary.bam','--possible_mismatches','AC','AG', 'CG', 'TG', 'CT','--name', 'Siomi1', 'Siomi2' , '--five_p', '3','--three_p','3','--output_pdf', 'out.pdf', '--output_tab', 'out.tab', '--expanded_output_tab', 'expanded.tab', '--min', '20', '--max', '22'])
+    # args = parser.parse_args(['--input', '3mismatches_ago2ip_s2.bam', '3mismatches_ago2ip_ovary.bam','--possible_mismatches','AC','AG', 'CG', 'TG', 'CT','--name', 'Siomi1', 'Siomi2' , '--five_p', '3','--three_p','3','--output_pdf', 'out.pdf', '--output_tab', 'out.tab', '--expanded_output_tab', 'expanded.tab', '--min', '20', '--max', '22'])
     args = parser.parse_args()
     reduced_dict = main()
-
-