v2.3.0 PR

bin/aggregate_lineages_bracken.py

@@ -6,30 +6,17 @@
 import json
 import sys
 
+from ranks import RANKS, UNCLASSIFIED_RANK
+
 UNCLASSIFIED = "Unclassified"
 UNKNOWN = "Unknown"
 
-RANKS = [
-    "superkingdom",
-    "clade",
-    "kingdom",
-    "phylum",
-    "subphylum",
-    "class",
-    "order",
-    "family",
-    "genus",
-    "species",
-    "subspecies",
-    "serotype",
-]
-
 
 def update_or_create_unclassified(entries, unclassified_count):
     """Handle unclassified entries."""
     entries[UNCLASSIFIED] = {
         "taxid": 0,
-        "rank": RANKS[0],
+        "rank": UNCLASSIFIED_RANK,
         "count": int(unclassified_count),
         "children": {
             UNKNOWN: {
@@ -65,10 +52,12 @@ def update_or_create_count(entry, entries, bracken_counts):
             new_entry = {"taxid": taxid, "rank": rank, "count": count, "children": {}}
             previous[name] = new_entry
             previous = new_entry["children"]
-            continue
-
-        current["count"] += count
-        previous = current["children"]
+        else:
+            current["count"] += count
+            previous = current["children"]
+        # Assigned on both branches: skipping it on the new-entry branch meant the
+        # species -> subspecies fallback above only fired when the parent species
+        # node happened to already exist in the tree.
         previous_rank = rank
 
     return entries

bin/assignment.py

@@ -161,7 +161,7 @@ def __eq__(self, other):
         else:
             return False
 
-    def get_read_map(self, taxon_id_map, parents={}):
+    def get_read_map(self, taxon_id_map, parents={}, paired=True):
         """
         Parses the kraken assignment file and collects the read_ids associated with each of the
         required taxon ids. If paired reads are provided, will consider the common ancestor of
@@ -170,17 +170,34 @@ def get_read_map(self, taxon_id_map, parents={}):
         Parameters:
             taxon_id_map (iter): Iterable of taxon ids to identify reads for.
             parents (dict): A dict mapping taxon id to parent taxon id from NCBI Taxonomy
+            paired (bool): If False, the input is known to be single-end, so each read_id can only
+                           appear once in the assignment file. `extended_map` exists purely to
+                           detect a read_id's second (mate) assignment, so for single-end data it
+                           can never be used and is skipped entirely - it would otherwise be a
+                           second whole-assignment-file-sized dict alongside `read_map`.
 
         Returns:
             read_map (dict): A dict from read_id to a taxon_id in the input iterable.
         """
         read_map = defaultdict(str)
-        extended_map = defaultdict(str)
+        extended_map = defaultdict(str) if paired else None
         comments = set()
         with open(self.file_name, "r") as kfile:
             for line in kfile:
-                assignment = KrakenAssignmentEntry(line)
-                taxon_id, read_id = assignment.taxon_id, assignment.read_id
+                # Inlined rather than going through KrakenAssignmentEntry: this loop runs once per
+                # read in the whole assignment file, and only the taxon_id/read_id fields are
+                # actually used here, so there's no need to build an object, validate the raw
+                # (unstripped) line, and then split a second time to unpack.
+                fields = line.strip().split("\t")
+                if len(fields) != 5:
+                    sys.stderr.write(
+                        f"Kraken assignment line {line} badly formatted - must have 5 fields"
+                    )
+                    sys.exit(11)
+                read_id = trim_read_id(fields[1])
+                taxon_id = fields[2]
+                if taxon_id == "A":
+                    taxon_id = "81077"
 
                 corrected_taxon_id = taxon_id
                 if parents:
@@ -198,7 +215,7 @@ def get_read_map(self, taxon_id_map, parents={}):
                                 f"Assign {taxon_id} to {corrected_taxon_id} list"
                             )
 
-                if read_id in extended_map:
+                if paired and read_id in extended_map:
                     mrca_taxon_id = get_mrca(
                         corrected_taxon_id, extended_map[read_id], parents
                     )
@@ -218,11 +235,91 @@ def get_read_map(self, taxon_id_map, parents={}):
                     if corrected_taxon_id != taxon_id:
                         comments.add(f"Assign {taxon_id} to {corrected_taxon_id} list")
                     read_map[read_id] = corrected_taxon_id
-                extended_map[read_id] = taxon_id
+                if paired:
+                    extended_map[read_id] = taxon_id
         for c in comments:
             print(c)
         return read_map
 
+    def get_read_maps(self, taxon_id_maps, parents={}, paired=True):
+        """
+        Like get_read_map, but builds several independent read_maps in a single pass over the
+        assignment file, instead of one pass per taxon_id_map. Each read_map in the result is
+        computed exactly as if get_read_map had been called separately for the corresponding
+        taxon_id_map - the maps do not interact with each other in any way, this purely shares
+        the (expensive, once-per-read) file parsing work across all of them.
+
+        Parameters:
+            taxon_id_maps (list): A list of taxon_id_map iterables, one per desired read_map.
+            parents (dict): A dict mapping taxon id to parent taxon id from NCBI Taxonomy
+            paired (bool): See get_read_map.
+
+        Returns:
+            read_maps (list): A list of read_map dicts, in the same order as taxon_id_maps.
+        """
+        read_maps = [defaultdict(str) for _ in taxon_id_maps]
+        extended_maps = [defaultdict(str) if paired else None for _ in taxon_id_maps]
+        comments = set()
+        with open(self.file_name, "r") as kfile:
+            for line in kfile:
+                fields = line.strip().split("\t")
+                if len(fields) != 5:
+                    sys.stderr.write(
+                        f"Kraken assignment line {line} badly formatted - must have 5 fields"
+                    )
+                    sys.exit(11)
+                read_id = trim_read_id(fields[1])
+                taxon_id = fields[2]
+                if taxon_id == "A":
+                    taxon_id = "81077"
+
+                for read_map, extended_map, taxon_id_map in zip(
+                    read_maps, extended_maps, taxon_id_maps
+                ):
+                    corrected_taxon_id = taxon_id
+                    if parents:
+                        while (
+                            corrected_taxon_id in parents
+                            and corrected_taxon_id not in taxon_id_map
+                            and corrected_taxon_id != "1"
+                        ):
+                            corrected_taxon_id = parents[corrected_taxon_id]
+                            if (
+                                corrected_taxon_id in taxon_id_map
+                                and corrected_taxon_id != taxon_id
+                            ):
+                                comments.add(
+                                    f"Assign {taxon_id} to {corrected_taxon_id} list"
+                                )
+
+                    if paired and read_id in extended_map:
+                        mrca_taxon_id = get_mrca(
+                            corrected_taxon_id, extended_map[read_id], parents
+                        )
+                        if mrca_taxon_id in taxon_id_map:
+                            if mrca_taxon_id != read_map[read_id]:
+                                comments.add(
+                                    f"Reassign {extended_map[read_id]} (and {corrected_taxon_id}) to mrca {mrca_taxon_id} list"
+                                )
+                                read_map[read_id] = mrca_taxon_id
+                        elif read_id in read_map:
+                            comments.add(
+                                f"MRCA {mrca_taxon_id} of {extended_map[read_id]} and {corrected_taxon_id} not in taxon_id_map"
+                            )
+                            del read_map[read_id]
+
+                    elif corrected_taxon_id in taxon_id_map:
+                        if corrected_taxon_id != taxon_id:
+                            comments.add(
+                                f"Assign {taxon_id} to {corrected_taxon_id} list"
+                            )
+                        read_map[read_id] = corrected_taxon_id
+                    if paired:
+                        extended_map[read_id] = taxon_id
+        for c in comments:
+            print(c)
+        return read_maps
+
     def load_file(self, taxon_ids=None):
         """
         Loads all entries in the kraken assignment file. If this is a paired file and there is a clash, result is unclassified

bin/check_hcid.py

@@ -74,10 +74,18 @@ def check_report_for_hcid(hcid_dict, taxonomy_dir, kreport_file):
             hcid_dict[taxid]["classified_parent_found"] = True
 
 
-def map_to_refs(query, ref_sam):
+_COMPLEMENT_TABLE = str.maketrans("ACGTNacgtn", "TGCANtgcan")
+
+
+def reverse_complement(seq):
+    return seq.translate(_COMPLEMENT_TABLE)[::-1]
+
+
+def map_to_refs(ref_sam):
     counts = defaultdict(int)
     ranges = defaultdict(list)
     read_ids = defaultdict(list)
+    read_seqs = {}
 
     in_file = open(ref_sam, "r")
     in_sam = Reader(in_file)
@@ -91,8 +99,20 @@ def map_to_refs(query, ref_sam):
         counts[sam_record.rname] += 1
         ranges[sam_record.rname].append(sam_record.coords)
         read_ids[sam_record.rname].append(sam_record.qname)
+        # SAM always reports SEQ/QUAL in forward-reference-strand orientation, so a
+        # reverse-strand (flag 16) alignment's SEQ/QUAL is the reverse-complement of what
+        # was actually in the original FASTQ - undo that here, so the reads exported below
+        # match the original file exactly, the same as reading them from the FASTQ (as this
+        # used to do, via a second full decompression pass) would have given.
+        if sam_record.flag == 16:
+            read_seqs[sam_record.qname] = (
+                reverse_complement(sam_record.seq),
+                sam_record.qual[::-1],
+            )
+        else:
+            read_seqs[sam_record.qname] = (sam_record.seq, sam_record.qual)
 
-    return counts, ranges, read_ids
+    return counts, ranges, read_ids, read_seqs
 
 
 def check_pileup(ref, ref_ranges, reference_file, min_coverage=0):
@@ -122,8 +142,8 @@ def coverage_hist(coverages):
     return hist
 
 
-def check_ref_coverage(hcid_dict, query, reference, ref_sam):
-    counts, ranges, read_ids = map_to_refs(query, ref_sam)
+def check_ref_coverage(hcid_dict, reference, ref_sam):
+    counts, ranges, read_ids, read_seqs = map_to_refs(ref_sam)
 
     for taxon in hcid_dict:
         taxon_found = True
@@ -169,8 +189,10 @@ def check_ref_coverage(hcid_dict, query, reference, ref_sam):
                 hcid_dict[taxon]["mapped_additional"]
             ) / len(hcid_dict[taxon]["additional_refs"])
 
+    return read_seqs
+
 
-def report_findings(hcid_dict, read_file, prefix):
+def report_findings(hcid_dict, read_seqs, prefix):
     keys = [
         "name",
         "taxon_id",
@@ -189,31 +211,17 @@ def report_findings(hcid_dict, read_file, prefix):
             w.writerow({key: hcid_dict[taxid][key] for key in keys})
 
     found = []
-    needed = {
-        name
-        for taxid in hcid_dict
-        if hcid_dict[taxid]["mapped_found"]
-        for name in hcid_dict[taxid]["mapped_read_ids"]
-    }
-
-    read_records = {}
-    for name, seq, qual in pyfastx.Fastq(read_file, full_name=False, build_index=False):
-        if name in needed:
-            read_records[name] = (seq, qual)
-
-        if needed - set(read_records.keys()) == set():
-            break
 
     for taxid in hcid_dict:
         if hcid_dict[taxid]["mapped_found"]:
             quals = []
             lens = []
             with open("%s.reads.fq" % taxid, "w") as f_reads:
                 for mapped_name in hcid_dict[taxid]["mapped_read_ids"]:
-                    record = read_records.get(mapped_name)
+                    record = read_seqs.get(mapped_name)
                     if record is None:
                         print(
-                            f"WARNING: read {mapped_name} not found in FASTQ, skipping",
+                            f"WARNING: read {mapped_name} not found in SAM, skipping",
                             file=sys.stderr,
                         )
                         continue
@@ -252,12 +260,6 @@ def main():
         required=False,
         help="Kraken or Bracken file of taxon relationships and quantities",
     )
-    parser.add_argument(
-        "-r",
-        dest="reads",
-        required=True,
-        help="FASTQ of reads",
-    )
     parser.add_argument(
         "-t",
         dest="taxonomy",
@@ -304,9 +306,9 @@ def main():
     sys.stderr.write("Check kraken report for counts\n")
     check_report_for_hcid(hcid_dict, args.taxonomy, args.kreport_file)
     sys.stderr.write("Check mapped coverage\n")
-    check_ref_coverage(hcid_dict, args.reads, args.ref_fasta, args.ref_sam)
+    read_seqs = check_ref_coverage(hcid_dict, args.ref_fasta, args.ref_sam)
     sys.stderr.write("Report findings\n")
-    report_findings(hcid_dict, args.reads, args.prefix)
+    report_findings(hcid_dict, read_seqs, args.prefix)
 
     now = datetime.now()
     time = now.strftime("%m/%d/%Y, %H:%M:%S")