miga-base 1.2.17.0 → 1.2.17.1

data/utils/FastAAI/FastAAI-legacy/kAAI_v1.0_virus.py

@@ -0,0 +1,1296 @@
+#!/usr/bin/env python
+
+"""
+########################################################################
+# Author:	   Carlos Ruiz
+# Intitution:   Georgia Institute of Technology
+# Version:	  0.8
+# Date:		 March 02, 2020
+
+# Description: Calculates the average amino acid identity using k-mers
+from single copy genes. It is a faster version of the regular AAI (Blast
+or Diamond) and the hAAI implemented in MiGA.
+########################################################################
+"""
+
+################################################################################
+"""---0.0 Import Modules---"""
+import subprocess, argparse, multiprocessing, datetime, shutil
+import textwrap, pickle, gzip
+from random import randint
+from pathlib import Path
+from sys import argv
+from sys import exit
+from functools import partial
+from os.path import realpath
+import numpy
+import tempfile
+
+
+################################################################################
+"""---1.0 Define Functions---"""
+# --- Run prodigal ---
+# ------------------------------------------------------
+def run_prodigal(input_file):
+    """
+    Runs prodigal, compares translation tables and stores faa files
+
+    Arguments:
+       input_file -- Path to genome FastA file
+    
+    Returns:
+        output -- Path to amino acid fasta result
+    """
+    # Predict proteins with translation tables 4 and 11
+    file_path = Path(input_file)
+    filename = file_path.name
+    folder = file_path.parent
+    protein_output = folder / (filename + '.faa')
+    output_11 = folder / (filename + '.faa.11')
+    temp_output = folder / (filename + '.temp')
+    subprocess.call(["prodigal", "-i", str(file_path), "-a", str(output_11), 
+                    "-p", "meta", "-q", "-o", str(temp_output)])
+    output_4 = folder / (filename + '.faa.4')
+    temp_output = folder / (filename + '.temp')
+    subprocess.call(["prodigal", "-i", str(file_path), "-a", str(output_4), 
+                    "-p", "meta", "-g", "4", "-q", "-o", str(temp_output)])
+
+    # Compare translation tables
+    length_4 = 0
+    length_11 = 0
+    with open(output_4, 'r') as table_4:
+        for line in table_4:
+            if line.startswith(">"):
+                continue
+            else:
+                length_4 += len(line.strip())
+
+    with open(output_11, 'r') as table_11:
+        for line in table_11:
+            if line.startswith(">"):
+                continue
+            else:
+                length_11 += len(line.strip())
+    
+    if (length_4 / length_11) >= 1.1:
+        shutil.copy(output_4, protein_output)
+    else:
+        shutil.copy(str(output_11), str(protein_output))
+    
+    # Remove intermediate files
+    output_4.unlink()
+    output_11.unlink()
+    temp_output.unlink()
+
+    # Remove stop '*' codons from protein sequences
+    with open(protein_output, 'r') as final_protein, open(temp_output, 'w') as temporal_file:
+        for line in final_protein:
+            if line.startswith(">"):
+                temporal_file.write("{}".format(line))
+            else:
+                line = line.replace('*', '')
+                temporal_file.write("{}".format(line))
+    shutil.copy(str(temp_output), str(protein_output))
+    temp_output.unlink()
+
+    return str(protein_output)
+# ------------------------------------------------------
+
+# --- Run prodigal for viruses ---
+# ------------------------------------------------------
+def run_prodigal_virus(input_file):
+    """
+    Runs prodigal, compares translation tables and stores faa files
+
+    Arguments:
+       input_file -- Path to genome FastA file
+    
+    Returns:
+        output -- Path to amino acid fasta result
+    """
+    # Predict proteins with translation tables 4 and 11
+    file_path = Path(input_file)
+    filename = file_path.name
+    folder = file_path.parent
+    protein_output = folder / (filename + '.faa')
+    temp_output = folder / (filename + '.temp')
+    subprocess.call(["prodigal", "-i", str(file_path), "-a", str(protein_output), 
+                    "-p", "meta", "-q", "-o", str(temp_output)])
+    
+    # Remove intermediate files
+    temp_output.unlink()
+
+    # Remove stop '*' codons from protein sequences
+    with open(protein_output, 'r') as final_protein, open(temp_output, 'w') as temporal_file:
+        for line in final_protein:
+            if line.startswith(">"):
+                temporal_file.write("{}".format(line))
+            else:
+                line = line.replace('*', '')
+                temporal_file.write("{}".format(line))
+    shutil.copy(str(temp_output), str(protein_output))
+    temp_output.unlink()
+
+    return str(protein_output)
+# ------------------------------------------------------
+
+# --- Run hmmsearch ---
+# ------------------------------------------------------
+def run_hmmsearch(input_file):
+    """
+    Runs hmmsearch on the set of SCGs and select the
+    best Archaea or Bacterial model
+    
+    Arguments:
+        input_file -- Path to protein FastA file
+    
+    Returns:
+        output -- Path to hmmsearch hits table
+    """
+    file_path = Path(input_file)
+    folder = file_path.parent
+    name = file_path.name
+    hmm_output = folder / (name + '.hmm')
+    temp_output = folder / (name + '.temp')
+    script_path = Path(realpath(__file__))
+    script_dir = script_path.parent
+    hmm_complete_model = script_dir / "00.Libraries/01.SCG_HMMs/Complete_SCG_DB.hmm"
+    subprocess.call(["hmmsearch", "--tblout", str(hmm_output), "-o", str(temp_output), "--cut_tc", "--cpu", "1",
+                    str(hmm_complete_model), str(file_path)])
+    temp_output.unlink()
+    return str(hmm_output)
+# ------------------------------------------------------
+
+# --- Filter HMM results for best matches ---
+# ------------------------------------------------------
+def hmm_filter(scg_hmm_file, keep):
+    """
+    Filters HMM results for best hits per protein
+    
+    Arguments:
+        SCG_HMM_file {file path} -- Path to HMM results file
+        keep {bool} -- Keep HMM files
+    
+    Returns:
+        outfile -- Path to filtered files
+    """
+    hmm_path = Path(scg_hmm_file)
+    name = hmm_path.name
+    folder = hmm_path.parent
+    outfile = folder / (name + '.filt')
+    hmm_hit_dict = {}
+    with open(scg_hmm_file, 'r') as hit_file:
+        for line in hit_file:
+            if line.startswith("#"):
+                continue
+            else:
+                hit = line.strip().split()
+                protein_name = hit[0]
+                score = float(hit[8])
+                if protein_name in hmm_hit_dict:
+                    if score > hmm_hit_dict[protein_name][0]:
+                        hmm_hit_dict[protein_name] = [score, line]
+                    elif score < hmm_hit_dict[protein_name][0]:
+                        continue
+                    else:
+                        if randint(2) > 0:
+                            hmm_hit_dict[protein_name] = [score, line]
+                else:
+                    hmm_hit_dict[protein_name] = [score, line]
+    with open(outfile, 'w') as output:
+        for hits in hmm_hit_dict.values():
+            output.write("{}".format(hits[1]))
+    return str(outfile)
+# ------------------------------------------------------
+
+# --- Find Kmers from HMM results ---
+# ------------------------------------------------------
+def kmer_extract(input_files):
+    """
+    Extract kmers from protein files that have hits
+    in the HMM searches.
+    
+    Arguments:
+        SCG_HMM_file {file path} -- Path to filtered HMM results.
+    
+    Returns:
+        [genome_kmers] -- Dictionary of kmers per gene. 
+    """
+    final_filename = input_files[0]
+    protein_file = input_files[1]
+    scg_hmm_file = input_files[2]
+    positive_matches = {}
+    positive_proteins = []
+    with open(scg_hmm_file, 'r') as hmm_input:
+        for line in hmm_input:
+            line = line.strip().split()
+            protein_name = line[0]
+            model_name = line[3]
+            score = line[8]
+            if model_name in positive_matches:
+                if score > positive_matches[model_name][1]:
+                    positive_matches[model_name] = [protein_name, score]
+                else:
+                    continue
+            else:
+                positive_matches[model_name] = [protein_name, score]
+    for proteins in positive_matches.values():
+        positive_proteins.append(proteins[0])
+    scg_kmers = read_kmers_from_file(protein_file, positive_proteins, 4)
+    for accession, protein in positive_matches.items():
+        scg_kmers[accession] = scg_kmers.pop(protein[0])
+    genome_kmers = {final_filename : scg_kmers}
+    return genome_kmers
+# ------------------------------------------------------
+
+# --- Extract kmers from protein sequences ---
+# ------------------------------------------------------
+def read_kmers_from_file(filename, positive_hits, ksize):
+    scg_kmers = {}
+    store_sequence = False
+    protein_name = ""
+    protein_sequence = ""
+    with open(filename) as fasta_in:
+        for line in fasta_in:
+            if line.startswith(">"):
+                if store_sequence == True:
+                    kmers = build_kmers(protein_sequence, ksize)
+                    scg_kmers[protein_name] = kmers
+                protein_sequence = ""
+                store_sequence = False
+                line = line.replace(">", "")
+                protein_name = line.strip().split()[0]
+                if protein_name in positive_hits:
+                    store_sequence = True
+            else:
+                if store_sequence == True:
+                    protein_sequence += line.strip()
+                else:
+                    continue
+            if store_sequence == True:
+                kmers = build_kmers(protein_sequence, ksize)
+                scg_kmers[protein_name] = kmers
+    return scg_kmers
+# ------------------------------------------------------
+
+# --- Extract kmers from viral protein sequences ---
+# ------------------------------------------------------
+def read_viral_kmers_from_file(input_information):
+    final_filename = input_information[0]
+    protein_file = input_information[1]
+    kmer_size = input_information[2]
+    scg_kmers = set()
+    protein_sequence = ""
+    store_sequence = False
+    with open(protein_file) as fasta_in:
+        for line in fasta_in:
+            if line.startswith(">"):
+                if store_sequence == True:
+                    kmers = build_kmers(protein_sequence, kmer_size)
+                    kmers = set(kmers.split(","))
+                    scg_kmers.update(kmers)
+                    protein_sequence = ""
+                else:
+                    protein_sequence = ""
+                    store_sequence = True
+            else:
+                protein_sequence += line.strip()
+    genome_kmers = {final_filename : list(scg_kmers)}
+    return genome_kmers
+# ------------------------------------------------------
+
+# --- Build Kmers ---
+# ------------------------------------------------------
+def build_kmers(sequence, ksize):
+    kmers = []
+    n_kmers = len(sequence) - ksize + 1
+
+    for i in range(n_kmers):
+        kmer = sequence[i:i + ksize]
+        kmers.append(kmer)
+    kmers_set = ','.join(set(kmers))
+    return kmers_set
+# ------------------------------------------------------
+
+# --- Parse kAAI when query == reference ---
+#Carlos, This function is not used with the new changes
+# ------------------------------------------------------
+def single_kaai_parser(query_id):
+    """
+    Calculates Jaccard distances on kmers from proteins shared
+    
+    Arguments:
+        query_id {str} -- Id of the query genome
+    
+    Returns:
+        [Path to output] -- Path to output file
+    """
+    file_path = Path(query_id)
+	
+	#Carlos, tempdir for safety
+    tmp_folder = tempfile.TemporaryDirectory()
+    running_folder = tmp_folder.name
+	
+	
+    temp_output = running_folder / file_path.with_suffix('.aai.temp')
+    # Get number and list of SCG detected in query
+    query_num_scg = len(query_kmer_dictionary[query_id])
+    query_scg_list = query_kmer_dictionary[query_id].keys()
+    # Start comparison with all genomes in the query dictionary
+    with open(temp_output, 'w') as out_file:
+        for target_genome, scg_ids in query_kmer_dictionary.items():
+            jaccard_similarities = []
+            # Get number and list of SCG detected in reference
+            target_num_scg = len(scg_ids)
+            target_scg_list = scg_ids.keys()
+            # Choose the smallest set of proteins
+            if query_num_scg > target_num_scg:
+                final_scg_list = target_scg_list
+            else:
+                final_scg_list = query_scg_list
+            # Compare all the proteins in the final SCG list
+            for accession in final_scg_list:
+                if accession in query_scg_list and accession in target_scg_list:
+                    # Get set and list for each SCG accession
+                    kmers_query = set(query_kmer_dictionary[query_id][accession].split(','))
+                    kmers_target = query_kmer_dictionary[target_genome][accession].split(',')
+                    # Calculate jaccard_similarity
+                    intersection = len(kmers_query.intersection(kmers_target))
+                    union = len(kmers_query.union(kmers_target))
+                    jaccard_similarities.append(intersection / union)
+                else:
+                    continue
+            try:
+                n = len(jaccard_similarities)
+                mean = sum(jaccard_similarities)/n
+                var = sum([ (x - mean)**2 for x in jaccard_similarities ])/(n - 1)
+                out_file.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(query_id, target_genome,
+                           round(mean, 4), round(var**0.5, 4),
+                           len(jaccard_similarities), len(final_scg_list)))
+            except:
+                out_file.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(query_id, target_genome,
+                           "NA", "NA", "NA", "NA"))
+
+    return temp_output
+# ------------------------------------------------------
+
+# --- Parse viral kAAI when query == reference ---
+# ------------------------------------------------------
+def single_virus_kaai_parser(query_id):
+    """
+    Calculates Jaccard distances on kmers from viral proteins
+    
+    Arguments:
+        query_id {str} -- Id of the query genome
+    
+    Returns:
+        [Path to output] -- Path to output file
+    """
+    file_path = Path(query_id)
+	
+	#Carlos, tempdir for safety
+    tmp_folder = tempfile.TemporaryDirectory()
+    running_folder = tmp_folder.name
+	
+	
+    temp_output = running_folder / file_path.with_suffix('.aai.temp')
+    # Start comparison with all genomes in the query dictionary
+    with open(temp_output, 'w') as out_file:
+        for target_genome, kmers_target in query_kmer_dictionary.items():
+            jaccard_index = None
+            kmers_query = set(query_kmer_dictionary[query_id])
+            intersection = len(kmers_query.intersection(kmers_target))
+            union = len(kmers_query.union(kmers_target))
+            try:
+                jaccard_index = intersection / union
+                out_file.write("{}\t{}\t{}\n".format(query_id, target_genome, jaccard_index))
+            except:
+                out_file.write("{}\t{}\tNA\n".format(query_id, target_genome))
+    return temp_output
+# ------------------------------------------------------
+
+# --- Parse kAAI when query != reference ---
+# ------------------------------------------------------
+def double_kaai_parser(query_id):
+    """
+    Calculates Jaccard distances on kmers from proteins shared
+    
+    Arguments:
+        query_id {str} -- Id of the query genome
+    
+    Returns:
+        [Path to output] -- Path to output file
+    """
+    file_path = Path(query_id)
+	
+	#Carlos, tempdir for safety
+    tmp_folder = tempfile.TemporaryDirectory()
+    running_folder = tmp_folder.name
+	
+	
+    temp_output = running_folder / file_path.with_suffix('.aai.temp')
+    # Get number and list of SCG detected in query
+    query_num_scg = len(query_kmer_dictionary[query_id])
+    query_scg_list = query_kmer_dictionary[query_id].keys()
+    # Start comparison with all genomes in the query dictionary
+    with open(temp_output, 'w') as out_file:
+        for target_genome, scg_ids in ref_kmer_dictionary.items():
+            jaccard_similarities = []
+            # Get number and list of SCG detected in reference
+            target_num_scg = len(scg_ids)
+            target_scg_list = scg_ids.keys()
+            # Choose the smallest set of proteins
+            if query_num_scg > target_num_scg:
+                final_scg_list = target_scg_list
+            else:
+                final_scg_list = query_scg_list
+            # Compare all the proteins in the final SCG list
+            for accession in final_scg_list:
+                if accession in query_scg_list and accession in target_scg_list:
+                    # Get set and list for each SCG accession
+                    kmers_query = set(query_kmer_dictionary[query_id][accession].split(','))
+                    kmers_target = ref_kmer_dictionary[target_genome][accession].split(',')
+                    # Calculate jaccard_similarity
+                    intersection = len(kmers_query.intersection(kmers_target))
+                    union = len(kmers_query.union(kmers_target))
+                    jaccard_similarities.append(intersection / union)
+                else:
+                    continue
+            try:
+                n = len(jaccard_similarities)
+                mean = sum(jaccard_similarities)/n
+                var = sum([ (x - mean)**2 for x in jaccard_similarities ])/(n - 1)
+                out_file.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(query_id, target_genome,
+                           round(mean, 4), round(var**0.5, 4),
+                           len(jaccard_similarities), len(final_scg_list)))
+            except:
+                out_file.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(query_id, target_genome,
+                           "NA", "NA", "NA", "NA"))
+    return temp_output
+# ------------------------------------------------------
+
+# --- Parse viral kAAI when query != reference ---
+# ------------------------------------------------------
+def double_viral_kaai_parser(query_id):
+    """
+    Calculates Jaccard distances on kmers from viral proteins
+    
+    Arguments:
+        query_id {str} -- Id of the query genome
+    
+    Returns:
+        [Path to output] -- Path to output file
+    """
+    file_path = Path(query_id)
+	
+	#Carlos, tempdir for safety
+    tmp_folder = tempfile.TemporaryDirectory()
+    running_folder = tmp_folder.name
+	
+	
+    temp_output = running_folder / file_path.with_suffix('.aai.temp')
+    # Start comparison with all genomes in the query dictionary
+    with open(temp_output, 'w') as out_file:
+        for target_genome, kmers_target in ref_kmer_dictionary.items():
+            jaccard_index = None
+            kmers_query = set(query_kmer_dictionary[query_id])
+            intersection = len(kmers_query.intersection(kmers_target))
+            union = len(kmers_query.union(kmers_target))
+            try:
+                jaccard_index = intersection / union
+                out_file.write("{}\t{}\t{}\n".format(query_id, target_genome, jaccard_index))
+            except:
+                out_file.write("{}\t{}\tNA\n".format(query_id, target_genome))
+    return temp_output
+# ------------------------------------------------------
+
+# --- Query == Reference initializer function ---
+# ------------------------------------------------------
+def single_dictionary_initializer(_dictionary):
+    """
+    Make dictionary available for multiprocessing
+    """
+    global query_kmer_dictionary
+    query_kmer_dictionary = _dictionary
+# ------------------------------------------------------
+
+# --- Query != Reference initializer function ---
+# ------------------------------------------------------
+def two_dictionary_initializer(_query_dictionary, _ref_dictionary):
+    """
+    Make dictionary available for multiprocessing
+    """
+    global query_kmer_dictionary
+    global ref_kmer_dictionary
+    query_kmer_dictionary = _query_dictionary
+    ref_kmer_dictionary = _ref_dictionary
+# ------------------------------------------------------
+
+# --- Merge kmer dictionaries ---
+# ------------------------------------------------------
+def merge_dicts(dictionaries):
+    """
+    Given any number of dicts, shallow copy and merge into a new dict,
+    precedence goes to key value pairs in latter dicts.
+    """
+    result = {}
+    for kmer_dictionary in dictionaries:
+        result.update(kmer_dictionary)
+    return result
+# ------------------------------------------------------
+
+
+#My version 1 - numpy-ized
+def single_kaai_parser_all_v_all(args):
+    """
+    Calculates Jaccard distances on kmers from proteins shared
+    
+    Arguments:
+        query_id {str} -- Id of the query genome
+    
+    Returns:
+        [Path to output] -- Path to output file
+    """
+    #Use split as slice if true
+    
+    query_id = args[0]
+    skip_first_n = args[1]
+    
+    file_path = Path(query_id)
+
+    tmp_folder = tempfile.TemporaryDirectory()
+    running_folder = tmp_folder.name
+    
+    #Just for my own testing. Temp dir is definitely the correct choice, here.
+    #running_folder = Path("faster_kaai")
+    
+    temp_output = running_folder / file_path.with_suffix('.aai.temp')
+
+    
+    #The goal is to numpy-ize the following loop in all possible aspects for a (hopeful) speed increase
+    
+    
+    #query_num_scg = len(query_kmer_dictionary[query_id])
+    
+    query_scg_list = numpy.array(list(query_kmer_dictionary[query_id].keys()))
+
+    with open(temp_output, 'w') as out_file:
+    
+        '''
+        Target genomes each control a set of protein family keys
+        
+        The goal is to get the jaccard index for the kmers in all cases 
+        of shared protein families for the two genomes in question, for 
+        each pair of genomes
+        
+        From above, we have the number of proteins in the query dict
+        and a list of the IDs
+        
+        below we get the number of proteins in the target dict
+        and a list of the IDs
+        
+        1 choose the shorter list (each item has to be in both to be used, after all)
+        2 check if each family is in both lists 
+        (kind of an unnecessarily big search cost, yeah? O(n) time with very few n = 1 cases; maybe we can make a dict of dicts of IDs, and check with try: [ID] except: ?)
+        3 get all of the jaccard similarities for kmers in shared protein families
+        
+        4 calculate the mean and variance for each similarity set
+        
+        5 repeat for the remaining genomes.
+            
+        '''
+        
+        #for target_genome, scg_ids in query_kmer_dictionary.items():
+        for target_genome in list(query_kmer_dictionary.keys())[skip_first_n:]:
+            scg_ids = query_kmer_dictionary[target_genome]
+            
+            #If self, 1.0 similarity.
+            if query_id == target_genome:
+                    out_file.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(query_id, target_genome,
+                    1.0, 0.0,
+                    len(query_scg_list), len(query_scg_list)))
+                    continue
+            
+            jaccard_similarities = []
+            # Get number and list of SCG detected in reference
+            #target_num_scg = len(scg_ids)
+            target_scg_list = numpy.array(list(scg_ids.keys()))
+            
+            final_scg_list = numpy.intersect1d(query_scg_list, target_scg_list)
+            
+            #I would like to figure out how to vectorize this.
+            for accession in final_scg_list:        
+                #Because of the prep work, these are already numpy arrays of numbers keying to the kmers they represent from the old kmer dict..
+                kmers_query = query_kmer_dictionary[query_id][accession]
+                kmers_target = query_kmer_dictionary[target_genome][accession]
+                
+                # Calculate jaccard_similarity - intersection is by far the slowest step, so this is by far the best place to optimize.
+                if len(kmers_query) < len(kmers_target):
+                    intersection = len(intersect1d_searchsorted(kmers_query, kmers_target))
+                else:
+                    intersection = len(intersect1d_searchsorted(kmers_target, kmers_query))
+                
+                union = len(numpy.union1d(kmers_query, kmers_target))
+                jaccard_similarities.append(intersection / union)
+            
+            #Allow for numpy in-builts; they're a little faster.
+            jaccard_similarities = numpy.array(jaccard_similarities, dtype=numpy.float_)
+            
+            try:
+                #No longer needed.
+                #n = len(jaccard_similarities)
+                mean = numpy.mean(jaccard_similarities)
+                var = numpy.std(jaccard_similarities)
+                out_file.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(query_id, target_genome,
+                           round(mean, 4), round(var, 4),
+                           len(jaccard_similarities), len(final_scg_list)))
+            except:
+                out_file.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(query_id, target_genome,
+                           "NA", "NA", "NA", "NA"))
+        return temp_output
+
+        
+def initializer_tracker(_dictionary1, _dictionary2):
+    """
+    Make dictionary available for multiprocessing
+    """
+    global kmer_dict
+    global tracker_dict
+    kmer_dict = _dictionary1
+    tracker_dict = _dictionary2
+    
+
+def unique_kmers(kmer_dict):
+
+    tracker_dict = {}
+    
+    counter = 0
+    
+    for file in kmer_dict:
+        for id in kmer_dict[file]:
+            #These are the actual kmers
+            for kmer in kmer_dict[file][id].split(','):
+                #Hash might be fast?
+                try:
+                    tracker_dict[kmer]
+                except:
+                    tracker_dict[kmer] = counter
+                    counter += 1
+    
+    return tracker_dict
+    
+
+def convert_kmers_to_indices(kmer_dict):
+    for genome in kmer_dict:
+        inner_count = 0
+        cur_tup = string_to_tup(genome)
+        for pf in kmer_dict[genome]:
+            kmer_dict[genome][pf] = cur_tup[inner_count]
+            inner_count += 1
+        
+    return kmer_dict
+
+def string_to_tup(genome):
+    sets = []
+    for pf in kmer_dict[genome]:
+        curset = []
+        for kmer in kmer_dict[genome][pf].split(","):
+            curset.append(tracker_dict[kmer])
+        
+        #Do all the overhead here, ONCE.
+        sets.append(numpy.sort(numpy.unique(numpy.array(curset, dtype=numpy.int32))))
+    
+    return(sets)
+    
+def numpyize_kmers(kmer_dict):
+    #make kmer global for tracker
+    single_dictionary_initializer(kmer_dict)
+    #get a list of kmer - index for all unique kmers
+    print("Indexing unique kmers")
+    tracker = unique_kmers(kmer_dict)
+    #Make these global for other functions
+    initializer_tracker(kmer_dict, tracker)
+    #convert comma sep. strings of kmers to ascending sorted lists of unique integers corresponding to the kmers in each protein, for each genome
+    print("Keying kmers")
+    kmer_dict = convert_kmers_to_indices(kmer_dict)
+    
+    #Get skip indices
+    smartargs = []
+    genome_ids = list(kmer_dict.keys())
+    for i in range(0, len(genome_ids)):
+        smartargs.append([genome_ids[i], i])
+		
+    print("Beginning AAI calculations now.")
+        
+    return kmer_dict, smartargs
+    
+#relies on assuming that the values in both of these arrays are unique and sorted, which I do in str_to_tup
+def intersect1d_searchsorted(A,B):
+    idx = numpy.searchsorted(B,A)
+    idx[idx==len(B)] = 0
+    return A[B[idx] == A]    
+	
+
+################################################################################
+"""---2.0 Main Function---"""
+
+def main():
+    # Setup parser for arguments.
+    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
+            description='''This script calculates the average amino acid identity using k-mers\n'''
+                        '''from single copy genes. It is a faster version of the regular AAI '''
+                        '''(Blast or Diamond) and the hAAI implemented in MiGA.'''
+            '''Usage: ''' + argv[0] + ''' -p [Protein Files] -t [Threads] -o [Output]\n'''
+            '''Global mandatory parameters: -g [Genome Files] OR -p [Protein Files] OR -s [SCG HMM Results] -o [AAI Table Output]\n'''
+            '''Optional Database Parameters: See ''' + argv[0] + ' -h')
+    mandatory_options = parser.add_argument_group('Mandatory i/o options. You must select an option for the queries and one for the references.')
+    mandatory_options.add_argument('--qg', dest='query_genomes', action='store', required=False,
+                                    help='File with list of query genomes.')
+    mandatory_options.add_argument('--qp', dest='query_proteins', action='store', required=False,
+                                    help='File with list of query proteins.')
+    mandatory_options.add_argument('--qh', dest='query_hmms', action='store', required=False,
+                                    help=textwrap.dedent('''
+                                    File with list of pre-computed query hmmsearch results.
+                                    If you select this option you must also provide a file with 
+                                    a list of protein files for the queries (with --qp).
+                                    '''))
+    mandatory_options.add_argument('--qd', dest='query_database', action='store', required=False,
+                                    help='File with list of pre-indexed query databases.')
+    mandatory_options.add_argument('--rg', dest='reference_genomes', action='store', required=False,
+                                    help='File with list of reference genomes.')
+    mandatory_options.add_argument('--rp', dest='reference_proteins', action='store', required=False,
+                                    help='File with list of reference proteins.')
+    mandatory_options.add_argument('--rh', dest='reference_hmms', action='store', required=False,
+                                    help=textwrap.dedent('''
+                                    File with list of pre-computed reference hmmsearch results.
+                                    If you select this option you must also provide a file with 
+                                    a list of protein files for the references (with --qp).
+                                    '''))
+    mandatory_options.add_argument('--rd', dest='reference_database', action='store', required=False,
+                                    help='File with list of pre-indexed reference databases.')
+    mandatory_options.add_argument('-o', '--output', dest='output', action='store', required=False, help='Output file. By default kaai_comparisons.txt')
+    additional_input_options = parser.add_argument_group('Behavior modification options.')
+    additional_input_options.add_argument('-e', '--ext', dest='extension', action='store', required=False, 
+                                            help='Extension to remove from original filename, e.g. ".fasta"')
+    additional_input_options.add_argument('-i', '--index', dest='index_db', action='store_true', required=False, 
+                                            help='Only index and store databases, i.e., do not perform comparisons.')
+    misc_options = parser.add_argument_group('Miscellaneous options')
+    misc_options.add_argument('--virus', dest='virus', action='store_true', required=False,
+                                help='Toggle virus-virus comparisons. Use only with viral genomes or proteins.')
+    misc_options.add_argument('-t', '--threads', dest='threads', action='store', default=1, type=int, required=False,
+                                help='Number of threads to use, by default 1')
+    misc_options.add_argument('-k', '--keep', dest='keep', action='store_false', required=False,
+                                help='Keep intermediate files, by default true')
+
+    args = parser.parse_args()
+
+    query_genomes = args.query_genomes
+    reference_genomes = args.reference_genomes
+    query_proteins = args.query_proteins
+    reference_proteins = args.reference_proteins
+    query_hmms = args.query_hmms
+    reference_hmms = args.reference_hmms
+    query_database = args.query_database
+    reference_database = args.reference_database
+    output = args.output
+    if output == None:
+        output == "kaai_comparisons.txt"
+    extension = args.extension
+    index_db = args.index_db
+    threads = args.threads
+    keep = args.keep
+    virus = args.virus
+
+    print("kAAI started on {}".format(datetime.datetime.now()))
+    # Check user input
+    # ------------------------------------------------------
+    # Check if no query was provided
+    if query_genomes == None and query_proteins == None and query_hmms == None and query_database == None:
+        exit('Please prove a file with a list of queries, e.g., --qg, --qp, --qh, or --qd)')
+    # Check query inputs
+    query_input = None
+    if query_hmms != None:
+        if virus == True:
+            exit("If you are comparing viruses, please start from the genome or protein files.")
+        query_input = query_hmms
+        if query_proteins != None:
+            print("Starting from query hmmsearch results.")
+            print("You also provided the list of protein files used for hmmsearch.")
+        elif query_proteins == None:
+            print("You chose to start from pre-computed hmmsearch results for your queries (--qh).")
+            print("However, I also need the location of the query proteins used for hmmsearch.")
+            exit("Please provide them with --qp.")
+    elif query_proteins != None:
+        query_input = query_proteins
+        print("Starting from query proteins.")
+    elif query_genomes != None:
+        query_input = query_genomes
+        print("Starting from query genomes.")
+    elif query_database != None:
+        query_input = query_database
+        print("Starting from the pre-indexed query database.")
+    # Check if no reference was provided
+    if reference_genomes == None and reference_proteins == None and reference_hmms == None and reference_database == None:
+        exit('Please prove a file with a list of references, e.g., --rg, --rp, --rh, or --rd)')
+    # Check reference inputs
+    reference_input = None
+    if reference_hmms != None:
+        if virus == True:
+            exit("If you are comparing viruses, please start from the genome or protein files.")
+        reference_input = reference_hmms
+        if reference_proteins != None:
+            print("Starting from reference hmmsearch results.")
+            print("You also provided the list of protein files used for hmmsearch.")
+        elif reference_proteins == None:
+            print("You chose to start from pre-computed hmmsearch results for your references (--rh).")
+            print("However, I also need the location of the query proteins used for hmmsearch.")
+            exit("Please provide them with --rp.")
+    elif reference_proteins != None:
+        reference_input = reference_proteins
+        print("Starting from reference proteins.")
+    elif reference_genomes != None:
+        reference_input = reference_genomes
+        print("Starting from reference genomes.")
+    elif reference_database != None:
+        reference_input = reference_database
+        print("Starting from the pre-indexed reference database.")
+    # ------------------------------------------------------
+
+    # Check if queries are the same as references (an all-vs-all comparison)
+    # ------------------------------------------------------
+    same_inputs = False
+    if query_input == reference_input:
+        same_inputs = True
+    if same_inputs == True:
+        print('You specified the same query and reference files.')
+        print('I will perform an all vs all comparison :)')
+    # ------------------------------------------------------
+
+    #* Database Parsing is the same regardless of bacterial or viral genomes
+    # If using pre-indexed databases, check if they are valid files.
+    # ------------------------------------------------------
+    # If any of the starting points is from database, then store the
+    # kmer structures in the corresponding dictionaries.
+    # Otherwise read the file list and get the filenames
+    query_kmer_dict = None
+    query_kmer_dict_list = []
+    reference_kmer_dict = None
+    reference_kmer_dict_list = []
+    # If starting from database and query == reference
+    if same_inputs == True:
+        if query_database != None:
+            with open(query_database) as query_database_files:
+                for db_location in query_database_files:
+                    if Path(db_location.strip()).is_file():
+                        with gzip.open(db_location.strip(), 'rb') as database_handle:
+                            temp_dict = pickle.load(database_handle)
+                            if isinstance(temp_dict,dict):
+                                query_kmer_dict_list.append(temp_dict)
+								#Carlos, this line serves no purpose but does take a bunch of time and mem.
+                                #print(query_kmer_dict_list)
+                            else:
+                                exit("One of the database files appear to have the wrong format. Please provide a correctly formated databases.")
+            query_kmer_dict = merge_dicts(query_kmer_dict_list)
+    else:
+    # If the inputs are not the same:
+        # If query and ref are provided
+        if query_database != None and reference_database != None:
+            with open(query_database, 'r') as query_database_files:
+                for db_location in query_database_files:
+                    if Path(db_location.strip()).is_file():
+                        with gzip.open(db_location.strip(), 'rb') as database_handle:
+                            temp_dict = pickle.load(database_handle)
+                            if isinstance(temp_dict,dict):
+                                query_kmer_dict_list.append(temp_dict)
+                            else:
+                                exit("One of the query database files appear to have the wrong format. Please provide a correctly formated databases.")
+            query_kmer_dict = merge_dicts(query_kmer_dict_list)
+            with open(reference_database) as reference_database_files:
+                for db_location in reference_database_files:
+                    if Path(db_location.strip()).is_file():
+                        with gzip.open(db_location.strip(), 'rb') as database_handle:
+                            temp_dict = pickle.load(database_handle)
+                            if isinstance(temp_dict,dict):
+                                reference_kmer_dict_list.append(temp_dict)
+                            else:
+                                exit("One of the reference database files appear to have the wrong format. Please provide a correctly formated databases.")
+            reference_kmer_dict = merge_dicts(reference_kmer_dict_list)
+        # If only the query has a db
+        elif query_database != None and reference_database == None:
+            with open(query_database) as query_database_files:
+                for db_location in query_database_files:
+                    if Path(db_location.strip()).is_file():
+                        with gzip.open(db_location.strip(), 'rb') as database_handle:
+                            temp_dict = pickle.load(database_handle)
+                            if isinstance(temp_dict,dict):
+                                query_kmer_dict_list.append(temp_dict)
+                            else:
+                                exit("One of the query database files appear to have the wrong format. Please provide a correctly formated databases.")
+            query_kmer_dict = merge_dicts(query_kmer_dict_list)
+        # If only the reference has a db
+        elif query_database == None and reference_database != None:
+            with open(reference_database) as reference_database_files:
+                for db_location in reference_database_files:
+                    if Path(db_location.strip()).is_file():
+                        with gzip.open(db_location.strip(), 'rb') as database_handle:
+                            temp_dict = pickle.load(database_handle)
+                            if isinstance(temp_dict,dict):
+                                reference_kmer_dict_list.append(temp_dict)
+                            else:
+                                exit("One of the reference database files appear to have the wrong format. Please provide a correctly formated databases.")
+            reference_kmer_dict = merge_dicts(reference_kmer_dict_list)
+    # ------------------------------------------------------
+
+    # Get files from the query and reference lists and then
+    # create a dictionary with resulting filenames and a list with dictionary keys
+    # The structure of the dictionary is:
+    # original_query, proteins, hmms, filtered_hmms
+    # ------------------------------------------------------
+    # First parse the query:
+    query_list = []
+    query_file_names = {}
+    # For bacterial genomes
+    if virus == False:
+        if query_database != None:
+            pass
+        else:
+            with open(query_input, 'r') as query_input_fh:
+                for line in query_input_fh:
+                    query_list.append(line.strip())
+            for index, query in enumerate(query_list):
+                query_name = str(Path(query).name)
+                if extension != None:
+                    query_name = query_name.replace(extension, "")
+                if query_hmms != None:
+                    query_protein_list = []
+                    with open(query_proteins, 'r') as query_protein_fh:
+                        for line in query_protein_fh:
+                            query_protein_list.append(line.strip())
+                    query_file_names[query_name] = [None, query_protein_list[index], query, query + '.filt']
+                elif query_proteins != None:
+                    query_file_names[query_name] = [None, query, query + '.hmm', query + '.hmm.filt']
+                elif query_genomes != None:
+                    query_file_names[query_name] = [query, query + '.faa', query + '.faa.hmm', query + '.faa.hmm.filt']
+    # For viral genomes
+    else:
+        if query_database != None:
+            pass
+        else:
+            with open(query_input, 'r') as query_input_fh:
+                for line in query_input_fh:
+                    query_list.append(line.strip())
+            for index, query in enumerate(query_list):
+                query_name = str(Path(query).name)
+                if extension != None:
+                    query_name = query_name.replace(extension, "")
+                if query_proteins != None:
+                    query_file_names[query_name] = [None, query]
+                elif query_genomes != None:
+                    query_file_names[query_name] = [query, query + '.faa']
+
+    # Then parse the references:
+    reference_list = []
+    reference_file_names = {}
+    if same_inputs == True:
+            pass
+    else:
+        # For bacterial genomes
+        if virus == False:
+            if reference_database != None:
+                pass
+            else:
+                with open(reference_input, 'r') as reference_input_fh:
+                    for line in reference_input_fh:
+                        reference_list.append(line.strip())
+                for index, reference in enumerate(reference_list):
+                    reference_name = str(Path(reference).name)
+                    if extension != None:
+                        reference_name = reference_name.replace(extension, "")
+                    if reference_hmms != None:
+                        reference_protein_list = []
+                        with open(reference_proteins, 'r') as reference_protein_fh:
+                            for line in reference_protein_fh:
+                                reference_protein_list.append(line.strip())
+                        reference_file_names[reference_name] = [None, reference_protein_list[index], reference, reference + '.filt']
+                    elif reference_proteins != None:
+                        reference_file_names[reference_name] = [None, reference, reference + '.hmm', reference + '.hmm.filt']
+                    elif query_genomes != None:
+                        reference_file_names[reference_name] = [reference, reference + '.faa', reference + '.faa.hmm', reference + '.faa.hmm.filt']
+        # For viral genomes
+        else:
+            if reference_database != None:
+                pass
+            else:
+                with open(reference_input, 'r') as reference_input_fh:
+                    for line in reference_input_fh:
+                        reference_list.append(line.strip())
+                for index, reference in enumerate(reference_list):
+                    reference_name = str(Path(reference).name)
+                    if extension != None:
+                        reference_name = reference_name.replace(extension, "")
+                    if reference_proteins != None:
+                        reference_file_names[reference_name] = [None, reference]
+                    elif query_genomes != None:
+                        reference_file_names[reference_name] = [reference, reference + '.faa']
+    # ------------------------------------------------------
+
+    # Pre-index and store databases
+    # ------------------------------------------------------
+    # Pre-index queries
+    if query_kmer_dict == None:
+        print("Processing queries...")
+        # If using bacterial genomes
+        if virus == False:
+            if query_hmms != None:
+                query_hmm_results = query_list
+            elif query_proteins != None:
+                query_protein_files = query_list
+                print("Searching against HMM models...")
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    query_hmm_results = pool.map(run_hmmsearch, query_protein_files)
+                finally:
+                    pool.close()
+                    pool.join()
+            elif query_genomes != None:
+                print("Predicting proteins...")
+                # Predict query proteins 
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    query_protein_files = pool.map(run_prodigal, query_list)
+                finally:
+                    pool.close()
+                    pool.join()
+                print("Done!")
+                print("Searching against HMM models...")
+                # Run hmmsearch against proteins predicted
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    query_hmm_results = pool.map(run_hmmsearch, query_protein_files)
+                finally:
+                    pool.close()
+                    pool.join()
+                print("Done!")
+            print("Filtering query hmmsearch results...")
+            # Filter query HMM search results
+            try:
+                pool = multiprocessing.Pool(threads)
+                pool.map(partial(hmm_filter, keep=keep), query_hmm_results)
+            finally:
+                pool.close()
+                pool.join()
+            print("Extracting kmers from query proteins...")
+            # Finding kmers for all queries
+            query_information = []
+            for name, values in query_file_names.items():
+                query_information.append((name, values[1], values[3]))
+            try:
+                pool = multiprocessing.Pool(threads)
+                kmer_results = pool.map(kmer_extract, query_information)
+            finally:
+                pool.close()
+                pool.join()
+            query_kmer_dict = merge_dicts(kmer_results)
+            del kmer_results
+        # If using viral genomes
+        else:
+            if query_genomes != None:
+                print("Predicting proteins...")
+                # Predict query proteins 
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    query_protein_files = pool.map(run_prodigal_virus, query_list)
+                finally:
+                    pool.close()
+                    pool.join()
+                print("Done!")
+            elif query_proteins != None:
+                query_protein_files = query_list
+            print("Extracting kmers from query proteins...")
+            query_information = []
+            for name, values in query_file_names.items():
+                query_information.append((name, values[1], 4))
+            try:
+                pool = multiprocessing.Pool(threads)
+                kmer_results = pool.map(read_viral_kmers_from_file, query_information)
+            finally:
+                pool.close()
+                pool.join()
+            query_kmer_dict = merge_dicts(kmer_results)
+            del kmer_results
+
+    # Pre-index references (if different from queries)
+    if same_inputs == False and reference_kmer_dict == None:
+        print("Processing references...")
+        # If using bacterial genomes
+        if virus == False:
+            if reference_hmms != None:
+                reference_hmm_results = reference_list
+            elif reference_proteins != None:
+                reference_protein_files = reference_list
+                print("Searching against HMM models...   ")
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    reference_hmm_results = pool.map(run_hmmsearch, reference_protein_files)
+                finally:
+                    pool.close()
+                    pool.join()
+            if reference_genomes != None:
+                print("Predicting proteins...")
+                # Predict reference proteins 
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    reference_protein_files = pool.map(run_prodigal, reference_list)
+                finally:
+                    pool.close()
+                    pool.join()
+                print("Done!")
+                print("Searching against HMM models...")
+                # Run hmmsearch against proteins predicted
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    reference_hmm_results = pool.map(run_hmmsearch, reference_protein_files)
+                finally:
+                    pool.close()
+                    pool.join()
+                print("Done!")
+            print("Filtering reference hmmsearch results...")
+            # Filter reference HMM search results
+            try:
+                pool = multiprocessing.Pool(threads)
+                pool.map(partial(hmm_filter, keep=keep), reference_hmm_results)
+            finally:
+                pool.close()
+                pool.join()
+            print("Extracting kmers from reference proteins...") 
+            # Finding kmers for all queries
+            reference_information = []
+            for name, values in reference_file_names.items():
+                reference_information.append((name, values[1], values[3]))
+            try:
+                pool = multiprocessing.Pool(threads)
+                kmer_results = pool.map(kmer_extract, reference_information)
+            finally:
+                pool.close()
+                pool.join()
+            reference_kmer_dict = merge_dicts(kmer_results)
+            del kmer_results
+        # If using viral genomes
+        else:
+            if query_genomes != None:
+                print("Predicting proteins...")
+                # Predict query proteins 
+                try:
+                    pool = multiprocessing.Pool(threads)
+                    query_protein_files = pool.map(run_prodigal, query_list)
+                finally:
+                    pool.close()
+                    pool.join()
+                print("Done!")
+            elif query_proteins != None:
+                query_protein_files = query_list
+            print("Extracting kmers from query proteins...")
+            reference_information = []
+            for name, values in reference_file_names.items():
+                reference_information.append((name, values[1], 4))
+            try:
+                pool = multiprocessing.Pool(threads)
+                kmer_results = pool.map(read_viral_kmers_from_file, reference_information)
+            finally:
+                pool.close()
+                pool.join()
+            query_kmer_dict = merge_dicts(kmer_results)
+            del kmer_results
+    # ------------------------------------------------------
+    
+    # Create or database(s) and compress it(them)
+    # ------------------------------------------------------
+    if same_inputs == True and query_database == None:
+        print("Saving pre-indexed database...")
+        query_database_name = query_input + '.db.gz'
+        with gzip.open(query_database_name, 'wb') as database_handle:
+            pickle.dump(query_kmer_dict, database_handle, protocol=4)
+    if same_inputs == False and query_database == None and reference_database == None:
+        print("Saving pre-indexed databases...")
+        query_database_name = query_input + '.db.gz'
+        reference_database_name = reference_input + '.db.gz'
+        with gzip.open(query_database_name, 'wb') as database_handle:
+            pickle.dump(query_kmer_dict, database_handle, protocol=4)
+        with gzip.open(reference_database_name, 'wb') as database_handle:
+            pickle.dump(reference_kmer_dict, database_handle, protocol=4)
+    elif same_inputs == False and query_database == None:
+        print("Saving pre-indexed query database...")
+        query_database_name = query_input + '.db.gz'
+        with gzip.open(query_database_name, 'wb') as database_handle:
+            pickle.dump(query_kmer_dict, database_handle, protocol=4)
+    elif same_inputs == False and reference_database == None:
+        print("Saving pre-indexed reference database...")
+        reference_database_name = reference_input + '.db.gz'
+        with gzip.open(reference_database_name, 'wb') as database_handle:
+            pickle.dump(reference_kmer_dict, database_handle, protocol=4)
+    # ------------------------------------------------------
+    # Calculate Jaccard distances
+    # ------------------------------------------------------
+    if index_db == True:
+        print("Finished pre-indexing databases.")
+        print("Next time you can run the program using only these files with --qd and(or) --rd.")
+    else:
+        print("Calculating shared Kmer fraction...")
+        if virus == False:
+            if same_inputs == True:
+                query_id_list = query_kmer_dict.keys()
+                try:
+					
+                    fixed_dict, smart_args = numpyize_kmers(query_kmer_dict)
+                    #single_dictionary_initializer(fixed_dict)
+                
+                    pool = multiprocessing.Pool(threads, initializer = single_dictionary_initializer, initargs = (fixed_dict,))
+                    Fraction_Results = pool.map(single_kaai_parser_all_v_all, smart_args)
+                finally:
+                    pool.close()
+                    pool.join()
+            else:
+                query_id_list = query_kmer_dict.keys()
+                try:
+                    pool = multiprocessing.Pool(threads, initializer = two_dictionary_initializer, initargs = (query_kmer_dict, reference_kmer_dict))
+                    Fraction_Results = pool.map(double_kaai_parser, query_id_list)
+                finally:
+                    pool.close()
+                    pool.join()
+        else:
+            if same_inputs == True:
+                query_id_list = query_kmer_dict.keys()
+                try:
+                    pool = multiprocessing.Pool(threads, initializer = single_dictionary_initializer, initargs = (query_kmer_dict,))
+                    Fraction_Results = pool.map(single_virus_kaai_parser, query_id_list)
+                finally:
+                    pool.close()
+                    pool.join()
+            else:
+                query_id_list = query_kmer_dict.keys()
+                try:
+                    pool = multiprocessing.Pool(threads, initializer = two_dictionary_initializer, initargs = (query_kmer_dict, reference_kmer_dict))
+                    Fraction_Results = pool.map(double_viral_kaai_parser, query_id_list)
+                finally:
+                    pool.close()
+                    pool.join()
+    # ------------------------------------------------------
+    
+    # Merge results into a single output
+    # ------------------------------------------------------
+        print("Merging results...")
+        with open(output, 'w') as outfile:
+            for file in Fraction_Results:
+                with open(file) as Temp:
+                    shutil.copyfileobj(Temp, outfile)
+                file.unlink()
+        print("kAAI finishied correctly on {}".format(datetime.datetime.now()))
+    # ------------------------------------------------------
+    # If comparing viral genomes
+
+
+	
+	
+	
+if __name__ == "__main__":
+    main()