--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/aggregate_linelisting.py	Thu Jul 04 19:40:13 2019 -0400
@@ -0,0 +1,279 @@
+#!/usr/bin/env python
+'''Reads in a fasta file of antigenic maps and one with the reference antigenic map as 
+protein SeqRecords. Compares amino acids of sample antigenic maps to corresponding sites
+in the reference and masks identical amino acids with dots. Writes headers (including
+amino acid position numbers read from the respective index array), the reference amino
+acid sequence and column headings required for both non-aggregated and aggregated line lists. 
+Outputs all headers and modified (i.e. dotted) sample sequences to a csv file.'''
+
+'''Author: Diane Eisler, Molecular Microbiology & Genomics, BCCDC Public Health Laboratory, Jan 2018'''
+
+import sys,string,os, time, Bio, re, argparse
+from Bio import Seq, SeqIO, SeqUtils, Alphabet, SeqRecord
+from Bio.SeqRecord import SeqRecord
+from Bio.Alphabet import IUPAC
+from Bio.Seq import Seq
+
+inputAntigenicMaps = sys.argv[1] #batch fasta file with antigenic map sequences
+refAntigenicMap = sys.argv[2] #fasta file of reference antigenic map sequence
+antigenicSiteIndexArray = sys.argv[3] #antigenic site index array csv file
+cladeDefinitionFile = sys.argv[4] #clade definition csv file
+outFileHandle = sys.argv[5] #user-specifed output filename
+
+agg_lineListFile = open(outFileHandle,'w') #open a writable output file
+
+indicesLine = "" #comma-separated antigenic site positions
+cladeList = [] #list of clade names read from clade definition file
+ref_seq = "" #reference antigenic map (protein sequence)
+seqList = [] #list of aa sequences to compare to reference
+
+BC_list = [] #empty list for BC samples
+AB_list = [] #empty list for AB samples
+ON_list = [] #empty list for ON samples
+QC_list = [] #empty list for QC samples
+nonprov_list = [] #empty list for samples not in above 4 provinces
+#dictionary for location-separated sequence lists
+prov_lists = {'1_BC':BC_list,'2_AB':AB_list,'3_ON':ON_list,'4_QC': QC_list, '5_nonprov': nonprov_list}
+
+def replace_matching_aa_with_dot(record):
+    """Compare amino acids in record to reference, mask identical symbols with dots, and return modified record."""
+    orig_seq = str(record.seq) #sequence string from SeqRecord
+    mod_seq = ""
+    #replace only those aa's matching the reference with dots
+    for i in range(0, len(orig_seq)):
+        if (orig_seq[i] == ref_seq[i]):
+            mod_seq = mod_seq  + '.'
+        else:
+            mod_seq  = mod_seq  + orig_seq[i]
+    #assign modified sequence to new SeqRecord and return it
+    rec = SeqRecord(Seq(mod_seq,IUPAC.protein), id = record.id, name = "", description = "")
+    return rec
+
+def extract_clade(record):
+    """Extract clade name (or 'No_Match') from sequence name and return as clade name. """
+    if record.id.endswith('No_Match'):
+        clade_name = 'No_Match'
+    else: #
+        for clade in cladeList:
+            if record.id.endswith(clade):
+                clade_name = clade
+    return clade_name
+    
+def extract_sample_name(record, clade):
+    """Extract sample name from sequence name and return sample name. """
+    end_index = record.id.index(clade)
+    sample_name = record.id[:end_index -1]
+    #return sample name as sequence name minus underscore and clade name
+    return sample_name
+
+def sort_by_location(record):
+    """Search sequence name for province name or 2-letter province code and add SeqRecord to
+    province-specific dictionary."""
+    seq_name = record.id
+    if ('-BC-' in seq_name) or ('/British_Columbia/' in seq_name):
+        BC_list.append(record) #add Sequence record to BC_list
+    elif ('-AB-' in seq_name) or ('/Alberta/' in seq_name):
+        AB_list.append(record) #add Sequence record to AB_list
+    elif ('-ON-' in seq_name) or ('/Ontario/' in seq_name):
+        ON_list.append(record) #add Sequence record to ON_list
+    elif ('-QC-' in seq_name) or ('/Quebec/' in seq_name):
+        QC_list.append(record) #add Sequence record to QC_list
+    else:
+        nonprov_list.append(record) #add Sequence record to nonprov_list
+    return
+
+def extract_province(record):
+    """Search sequence name for province name or 2-letter province code and return province."""
+    seq_name = record.id
+    if ('-BC-' in seq_name) or ('/British_Columbia/' in seq_name):
+        province = 'British Columbia'
+    elif ('-AB-' in seq_name) or ('Alberta' in seq_name):
+        province = '/Alberta/'
+    elif ('-ON-' in seq_name) or ('/Ontario/' in seq_name):
+        province = 'Ontario'
+    elif ('-QC-' in seq_name) or ('/Quebec/' in seq_name):
+        province = 'Quebec'
+    else:
+        province = "other"
+    return province
+
+def get_sequence_length(record):
+    """Return length of sequence in a SeqRecord."""
+    sequenceLength = len(str((record.seq)))
+    return sequenceLength
+
+def get_antigenic_site_substitutions(record):
+    """Count number of non-dotted amino acids in SeqRecord sequence and return as substitutions."""
+    sequenceLength = get_sequence_length(record)
+    seqString = str(record.seq)
+    matches = seqString.count('.')
+    substitutions = sequenceLength - matches
+    return substitutions
+
+def calculate_percent_id(record, substitutions):
+    """Calculate percent sequence identity to reference sequence, based on substitutions
+and sequence length and return percent id as a ratio (i.e. 0.90 no 90%)."""
+    sequenceLength = get_sequence_length(record)
+    percentID = (1.00 - (float(substitutions)/float(sequenceLength)))
+    return percentID
+
+def output_aggregated_linelist(a_list):
+    """Output aggregated line list of SeqRecords in csv format."""
+    sequevars = {} #dict of sequevar: SeqRecord list
+    firstRecordID = None
+    #examine dotted/masked sequences in list and assign unique ones as dict keys
+    for rec in a_list:
+        rec = replace_matching_aa_with_dot(rec)
+        sequence =str(rec.seq)
+        #if the sequence is a key in the dict, add SeqRecord to list
+        if sequence in sequevars:
+            #if sequence already in dict as a key, increment the value
+            sequevars[sequence].append(rec)
+        else:
+            #if sequence not in dict, add is as new key with list of 1 SeqRecord
+            sequevars[sequence] = [rec]
+    #get list of sorted unique sequence keys
+    sorted_unique_seq_keys = sorted(sequevars.keys())
+    #process each list of SeqRecords sharing a unique sequence
+    for u in sorted_unique_seq_keys:
+        #access list of sequences by unique sequence
+        listOfSeqs = sequevars[u]
+        #sort this list of SeqRecords by record.id (i.e. name)
+        listOfSeqs = [f for f in sorted(listOfSeqs, key = lambda x : x.id)]
+        N = len(listOfSeqs)
+        #output details of first SeqRecord to csv
+        firstRecord = listOfSeqs[0]
+        province = extract_province(firstRecord)
+        clade = extract_clade(firstRecord)
+        substitutions = get_antigenic_site_substitutions(firstRecord)
+        percentID = calculate_percent_id(firstRecord,substitutions)
+        name = extract_sample_name(firstRecord, clade)
+        name_part = name.rstrip() + ','
+        N_part = str(N) + ','
+        clade_part = clade + ','
+        substitutions_part = str(substitutions) + ','
+        percID_part = str(percentID) + ','
+        col = " ," #empty column
+        sequence = str(firstRecord.seq).strip()
+        csv_seq = ",".join(sequence) +","
+        comma_sep_output = name_part + N_part + clade_part + col + csv_seq + substitutions_part + percID_part + "\n"
+        #write first member of unique sequence list to csv
+        agg_lineListFile.write(comma_sep_output)
+        #print sequence records in sequevar to console
+        print("\n\t\t%i SeqRecords matching Sequevar: %s" % (len(listOfSeqs), u))
+
+        #to uncollapse sequevar group, print each member of the sequevar list to csv output
+        '''for i in range(1,len(listOfSeqs)):
+            currentRec = listOfSeqs[i]
+            province = extract_province(currentRec)
+            clade = extract_clade(currentRec)
+            substitutions = get_antigenic_site_substitutions(currentRec)
+            percentID = calculate_percent_id(currentRec,substitutions)
+            name_part = (currentRec.id).rstrip() + ','
+            N_part = "n/a" + ','
+            clade_part = clade + ','
+            substitutions_part = str(substitutions) + ','
+            percID_part = str(percentID) + ','
+            col = " ," #empty column
+            sequence = str(currentRec.seq).strip()
+            csv_seq = ",".join(sequence) +","
+            comma_sep_output = name_part + N_part + clade_part + col + csv_seq + substitutions_part + percID_part + "\n"
+            agg_lineListFile.write(comma_sep_output)   '''
+    return
+	
+with open (antigenicSiteIndexArray,'r') as siteIndices:
+    """Read amino acid positions from antigenic site index array and print as header after one empty row."""
+    col = "," #empty column
+    #read amino acid positions and remove trailing whitespace
+    for line in siteIndices:
+        #remove whitespace from the end of each line
+        indicesLine = line.rstrip()
+    row1 = "\n"
+    #add comma-separated AA positions to header line
+    row2 = col + col + col + col + indicesLine + "\n"
+    #write first (empty) and 2nd (amino acid position) lines to output file
+    agg_lineListFile.write(row1)
+    agg_lineListFile.write(row2)
+
+with open (refAntigenicMap,'r') as refMapFile:
+    """Read reference antigenic map from fasta and output amino acids, followed by column headers."""
+    #read sequences from fasta to SeqRecord, uppercase, and store sequence string to ref_seq
+    record = SeqIO.read(refMapFile,"fasta",alphabet=IUPAC.protein)
+    record = record.upper()
+    ref_seq = str(record.seq).strip() #store sequence in variable for comparison to sample seqs
+    col = "," #empty column
+    name_part = (record.id).rstrip() + ','
+    sequence = str(record.seq).strip()
+    csv_seq = ",".join(sequence)
+    #output row with reference sequence displayed above sample sequences
+    row3 = name_part + col + col + col + csv_seq + "\n"
+    agg_lineListFile.write(row3)
+    positions = indicesLine.split(',')
+    numPos = len(positions)
+    empty_indicesLine = ',' * numPos
+    #print column headers for sample sequences
+    row4 = "Sequence Name,N,Clade,Extra Substitutions," + empty_indicesLine + "Number of Amino Acid Substitutions in Antigenic Sites,% Identity of Antigenic Site Residues\n"
+    agg_lineListFile.write(row4)
+    print("\nREFERENCE ANTIGENIC MAP: '%s' (%i amino acids)" % (record.id, len(record)))
+
+with open(cladeDefinitionFile,'r') as cladeFile:
+    """Read clade definition file and store clade names in list."""
+    #remove whitespace from the end of each line and split elements at commas
+    for line in cladeFile:
+        elementList = line.rstrip().split(',')
+        name = elementList[0] #move 1st element to name field
+        cladeList.append(name)
+
+with open(inputAntigenicMaps,'r') as extrAntigMapFile:
+    """Read antigenic maps as protein SeqRecords and add to list."""
+    #read Sequences from fasta file, uppercase and add to seqList
+    for record in SeqIO.parse(extrAntigMapFile, "fasta", alphabet=IUPAC.protein):
+        record = record.upper()
+        seqList.append(record) #add Seq to list of Sequences
+
+#print number of sequences to be processed as user check
+print("\nCOMPARING %i flu antigenic map sequences to the reference..." % len(seqList))
+for record in seqList:
+    #assign SeqRecords to province-specific dictionaries
+    sort_by_location(record)
+
+#access prov segregated lists in order
+sorted_prov_keys = sorted(prov_lists.keys())
+print("\nSequence Lists Sorted by Province: ")
+for prov in sorted_prov_keys:
+    current_list = prov_lists[prov]
+    #mask AA's identical to reference sequence with dot
+    masked_list = [] # empty temporary list to park masked sequences
+    for record in current_list:
+        masked_rec = replace_matching_aa_with_dot(record)
+        masked_list.append(masked_rec)
+    prov_lists[prov] =  masked_list #replace original SeqRecord list with masked list
+
+#group sequences in province-sorted list into clades
+for prov in sorted_prov_keys:
+    prov_list = prov_lists[prov]
+    by_clades_dict = {} #empty dict for clade:seqRecord list groups
+    print("\n'%s' List (Amino Acids identical to Reference are Masked): " % (prov))
+    for rec in prov_list:
+        clade = extract_clade(rec)
+        if clade in by_clades_dict:
+            #if clade already in dict as key, append record to list (value)
+            by_clades_dict[clade].append(rec)
+        else: #add clade as key to dict, value is list of 1 SeqRecord
+            by_clades_dict[clade] = [rec]
+    #get list of alphabetically sorted clade keys
+    sorted_clade_keys = sorted(by_clades_dict.keys())
+    print("\tNumber of clades: ", len(by_clades_dict))
+    #group each list of sequences in clade by sequevars
+    for key in sorted_clade_keys:
+        print("\n\tCLADE: %s Number of Members: %i" % (key, len(by_clades_dict[key])))
+        a_list = by_clades_dict[key]
+        for seqrec in a_list:
+            print("\t %s: %s" %(seqrec.id,str(seqrec.seq)))
+        #output the list to csv as aggregated linelist
+        output_aggregated_linelist(a_list)
+    
+print("Aggregated Linelist written to file: '%s\n'"  % (outFileHandle))
+extrAntigMapFile.close()
+refMapFile.close()
+agg_lineListFile.close()