Use most specific hit NR top hits

Author: lvreynoso

lib/idseq-dag/idseq_dag/util/top_hits.py

@@ -117,23 +117,34 @@ def summary(self):
         return r
 
 
-def _optimal_hit_for_each_query_nr(blast_output_path, max_evalue):
+def _optimal_hit_for_each_query_nr(blast_output_path, lineage_map, accession2taxid_dict, max_evalue):
     contigs_to_best_alignments = defaultdict(list)
     accession_counts = defaultdict(lambda: 0)
 
     with open(blast_output_path) as blastn_6_f:
         # For each contig, get the alignments that have the best total score (may be multiple if there are ties).
+        # Prioritize the specificity of the hit.
+        specificity_to_best_alignments = defaultdict(dict)
         for alignment in BlastnOutput6Reader(blastn_6_f):
             if alignment["evalue"] > max_evalue:
                 continue
             query = alignment["qseqid"]
-            best_alignments = contigs_to_best_alignments[query]
 
-            if len(best_alignments) == 0 or best_alignments[0]["bitscore"] < alignment["bitscore"]:
-                contigs_to_best_alignments[query] = [alignment]
+            lineage_taxids = _get_lineage(alignment["sseqid"], lineage_map, accession2taxid_dict)
+            specificity = next((level for level, taxid_at_level in enumerate(lineage_taxids) if int(taxid_at_level) > 0), float("inf"))
+
+            best_alignments = specificity_to_best_alignments[query]
+
+            if (specificity not in best_alignments) or best_alignments[specificity][0]["bitscore"] < alignment["bitscore"]:
+                specificity_to_best_alignments[query][specificity] = [alignment]
             # Add all ties to best_hits.
-            elif len(best_alignments) > 0 and best_alignments[0]["bitscore"] == alignment["bitscore"]:
-                contigs_to_best_alignments[query].append(alignment)
+            elif len(best_alignments[specificity]) > 0 and best_alignments[specificity][0]["bitscore"] == alignment["bitscore"]:
+                specificity_to_best_alignments[query][specificity].append(alignment)
+
+        # Choose the best alignments with the most specific taxid information.
+        for contig_id, specificity_alignment_dict in specificity_to_best_alignments.items():
+            specific_best_alignments = next(specificity_alignment_dict[specificity] for specificity in sorted(specificity_alignment_dict.keys()))
+            contigs_to_best_alignments[contig_id] = specific_best_alignments
 
         # Create a map of accession to best alignment count.
         for _contig_id, alignments in contigs_to_best_alignments.items():
@@ -199,15 +210,15 @@ def _optimal_hit_for_each_query_nt(blast_output, lineage_map, accession2taxid_di
             # We prioritize the specificity of the hit; hits with species taxids are taken before hits without
             # Specificity is just the index of the tuple returned by _get_lineage(); 0 for species, 1 for genus, etc.
             lineage_taxids = _get_lineage(hit.sseqid, lineage_map, accession2taxid_dict)
-            specificity = next(level for level, taxid_at_level in enumerate(lineage_taxids) if int(taxid_at_level) > 0)
+            specificity = next((level for level, taxid_at_level in enumerate(lineage_taxids) if int(taxid_at_level) > 0), float("inf"))
 
             if (specificity not in best_hits) or best_hits[specificity][0].total_score < hit.total_score:
                 best_hits[specificity] = [hit]
             # Add all ties to best_hits[specificity].
             elif len(best_hits[specificity]) > 0 and best_hits[specificity][0].total_score == hit.total_score:
                 best_hits[specificity].append(hit)
 
-        specific_best_hits = next(hits for specificity, hits in best_hits.items() if len(hits) > 0)
+        specific_best_hits = next(best_hits[specificity] for specificity in sorted(best_hits.keys()))
         contigs_to_blast_candidates[specific_best_hits[0].qseqid] = specific_best_hits
 
     # Create a map of accession to blast candidate count.
@@ -232,13 +243,17 @@ def _optimal_hit_for_each_query_nt(blast_output, lineage_map, accession2taxid_di
 
 def get_top_m8_nr(
     blast_output,
+    lineage_map_path,
+    accession2taxid_dict_path,
     blast_top_blastn_6_path,
     max_evalue=MAX_EVALUE_THRESHOLD,
 ):
     ''' Get top m8 file entry for each contig from blast_output and output to blast_top_m8 '''
-    with open(blast_top_blastn_6_path, "w") as blast_top_blastn_6_f:
+    with open(blast_top_blastn_6_path, "w") as blast_top_blastn_6_f, \
+        open_file_db_by_extension(lineage_map_path, "lll") as lineage_map, \
+        open_file_db_by_extension(accession2taxid_dict_path, "L") as accession2taxid_dict:  # noqa
         BlastnOutput6Writer(blast_top_blastn_6_f).writerows(
-            _optimal_hit_for_each_query_nr(blast_output, max_evalue)
+            _optimal_hit_for_each_query_nr(blast_output, lineage_map, accession2taxid_dict, max_evalue)
         )
 
 

workflows/long-read-mngs/run.wdl

@@ -687,6 +687,8 @@ task FindTopHitsNT {
 task FindTopHitsNR {
     input {
         File nr_m8
+        File lineage_db
+        File accession2taxid_db
         String docker_image_id
     }
 
@@ -695,7 +697,12 @@ task FindTopHitsNR {
         python3 <<CODE
         from idseq_dag.steps.blast_contigs import get_top_m8_nr
 
-        get_top_m8_nr("~{nr_m8}", "rapsearch2.blast.top.m8")
+        get_top_m8_nr(
+            "~{nr_m8}",
+            "~{lineage_db}",
+            "~{accession2taxid_db}",
+            "rapsearch2.blast.top.m8"
+        )
         CODE
 
         python3 - << 'EOF'
@@ -1406,12 +1413,16 @@ workflow czid_long_read_mngs {
     call FindTopHitsNT {
         input:
             nt_m8 = RunNTAlignment.nt_m8,
+            lineage_db = lineage_db,
+            accession2taxid_db = accession2taxid_db,
             docker_image_id = docker_image_id,
     }
 
     call FindTopHitsNR {
         input:
             nr_m8 = RunNRAlignment.nr_m8,
+            lineage_db = lineage_db,
+            accession2taxid_db = accession2taxid_db,
             docker_image_id = docker_image_id,
     }