diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml deleted file mode 100644 index e424839c..00000000 --- a/.github/workflows/docker-build-push.yml +++ /dev/null @@ -1,40 +0,0 @@ -name: quay_build_push - -on: - release: - types: [published] - -jobs: - docker: - runs-on: ubuntu-latest - steps: - - name: Docker meta - id: meta - uses: docker/metadata-action@v5 - with: - images: quay.io/biowilko/scylla - flavor: | - latest=true - tags: | - type=semver,pattern={{version}} - type=sha - - name: Checkout - uses: actions/checkout@v3 - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - name: Login to Docker Hub - uses: docker/login-action@v3 - with: - registry: quay.io - username: ${{ secrets.QUAY_USERNAME }} - password: ${{ secrets.QUAY_PASSWORD }} - - name: Build and push - uses: docker/build-push-action@v5 - with: - context: ./dockerfiles/scylla/ - platforms: linux/amd64 - push: true - tags: ${{ steps.meta.outputs.tags }} - labels: ${{ steps.meta.outputs.labels }} diff --git a/bin/aggregate_lineages_bracken.py b/bin/aggregate_lineages_bracken.py index 9f1ebe29..9c6daa25 100755 --- a/bin/aggregate_lineages_bracken.py +++ b/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 diff --git a/bin/assignment.py b/bin/assignment.py index 705cb847..a06ea9e4 100644 --- a/bin/assignment.py +++ b/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 diff --git a/bin/check_hcid.py b/bin/check_hcid.py index 0f0c2b7b..5468047e 100755 --- a/bin/check_hcid.py +++ b/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,20 +211,6 @@ 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"]: @@ -210,10 +218,10 @@ def report_findings(hcid_dict, read_file, prefix): 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") diff --git a/bin/check_reads.py b/bin/check_reads.py index 25674aa6..3ca86c05 100755 --- a/bin/check_reads.py +++ b/bin/check_reads.py @@ -8,35 +8,91 @@ from assignment import trim_read_id -def check_fastq(read_file): +def check_fastq(read_file, sample_size=1000, max_illumina_read_length=1000): is_duplicates = True is_interleaved = False is_concat = False position = 0 - seen_names = set() - seen_trimmed = set() - positions = defaultdict(int) + # Reference structures are only ever populated from the first + # `sample_size` reads, so memory is bounded regardless of file size. + # Reads beyond that window are only *probed* against this reference, + # never added to it - any duplicate whose first occurrence falls + # outside the window is missed, but for both the interleaved case + # (pairs are adjacent) and the concatenated case (the whole first + # half is duplicated verbatim in the second half) every duplicate's + # first occurrence is guaranteed to fall within an early window. + ref_names = set() + ref_trimmed = set() + ref_positions = {} + ref_checks = {} + differences = defaultdict(int) - checks = {} + dup_positions = [] + early_exit = False sys.stderr.write(f"Reading in {read_file}\n") for record in pyfastx.Fastq(read_file, build_index=False): name, seq, qual = record + + # Illumina chemistry has a hard per-read length ceiling well under + # 1kb (even long-read kits top out around 600bp), so a single read + # longer than this proves the file cannot be interleaved or + # concatenated Illumina paired-end data - the entire pairing check + # is moot and we can stop scanning immediately, which matters a lot + # for very large long-read (e.g. ONT) FASTQs. Checked on every read + # for the whole scan (not just the reference window) since it's a + # free len() on data already parsed, and a degraded/fragmented run + # could have a run of short reads before the long ones appear. + # + # Gated on `not differences` so it never fires mid-detection once + # duplicate evidence has actually been seen. In practice this rarely + # protects the "whole ONT file accidentally duplicated onto itself" + # case though: that duplicate evidence only appears once the scan + # reaches the original file's length, and a long-read file almost + # always has a read over the threshold well before that - so this + # optimisation trades away that (rarer, non-Illumina-specific) + # detection in favour of the much more common speed win. + if not differences and len(seq) > max_illumina_read_length: + sys.stderr.write( + f"Read {name} is {len(seq)}bp (> {max_illumina_read_length}bp): " + "file cannot be Illumina paired-end data, skipping pairing check\n" + ) + return 0 + trimmed_name = trim_read_id(name) - if name in seen_names: - differences[trimmed_name] = position - positions[trimmed_name] - elif trimmed_name in seen_trimmed: - differences[trimmed_name] = position - positions[trimmed_name] - if trimmed_name in checks and checks[trimmed_name] != seq: - is_duplicates = False - positions[trimmed_name] = position + if position < sample_size: + if name in ref_names or trimmed_name in ref_trimmed: + differences[trimmed_name] = position - ref_positions[trimmed_name] + dup_positions.append(position) + if trimmed_name in ref_checks and ref_checks[trimmed_name] != seq: + is_duplicates = False + ref_positions[trimmed_name] = position + ref_names.add(name) + ref_trimmed.add(trimmed_name) + if str(position + 1).startswith("1") or str(position + 1).startswith("5"): + ref_checks[trimmed_name] = seq + elif trimmed_name in ref_trimmed: + differences[trimmed_name] = position - ref_positions[trimmed_name] + dup_positions.append(position) + if trimmed_name in ref_checks and ref_checks[trimmed_name] != seq: + is_duplicates = False + ref_positions[trimmed_name] = position + position += 1 - seen_names.add(name) - seen_trimmed.add(trimmed_name) - if str(position).startswith("1") or str(position).startswith("5"): - checks[trimmed_name] = seq + + # Once we're well past the reference window and every duplicate + # found so far is an adjacent pair, this is confidently an + # interleaved file - stop scanning early rather than reading the + # rest of a potentially very large FASTQ just to confirm it. + if ( + position >= 2 * sample_size + and differences + and set(differences.values()) == {1} + ): + early_exit = True + break num_seqs = position @@ -45,8 +101,7 @@ def check_fastq(read_file): return 0 difference_set = set([v for v in differences.values()]) - position_set = set([positions[k] for k in differences]) - min_duplicate = min(position_set) + min_duplicate = min(dup_positions) if difference_set == {1}: # if all pairs are next to each other, have interleaved file is_interleaved = True @@ -109,13 +164,14 @@ def check_fastq(read_file): out_handle = r2 key = "r2" + num_seqs_desc = f">= {num_seqs}" if early_exit else str(num_seqs) if is_duplicates: sys.stderr.write( - f"Input {num_seqs} sequences have resulted in out file with the following read counts: {out_prefix}.fixed.fastq : {counts['r']}\n" + f"Input {num_seqs_desc} sequences have resulted in out file with the following read counts: {out_prefix}.fixed.fastq : {counts['r']}\n" ) else: sys.stderr.write( - f"Input {num_seqs} sequences have resulted in out files with the following read counts: {out_prefix}.R1.fastq : {counts['r1']}, {out_prefix}.R2.fastq : {counts['r2']}\n" + f"Input {num_seqs_desc} sequences have resulted in out files with the following read counts: {out_prefix}.R1.fastq : {counts['r1']}, {out_prefix}.R2.fastq : {counts['r2']}\n" ) return 11 @@ -128,8 +184,33 @@ def check_fastq(read_file): ) ) parser.add_argument("--fastq", help="Input FASTQ.") + parser.add_argument( + "--sample-size", + dest="sample_size", + type=int, + default=1000, + help=( + "Number of leading reads to use to build the reference set for " + "duplicate/interleave/concatenation detection (default: 1000)." + ), + ) + parser.add_argument( + "--max-illumina-read-length", + dest="max_illumina_read_length", + type=int, + default=1000, + help=( + "A single read longer than this (bp) proves the file cannot be " + "Illumina paired-end data, so the pairing check is skipped " + "immediately (default: 1000)." + ), + ) args = parser.parse_args() - exit_code = check_fastq(args.fastq) + exit_code = check_fastq( + args.fastq, + sample_size=args.sample_size, + max_illumina_read_length=args.max_illumina_read_length, + ) sys.exit(exit_code) diff --git a/bin/concatenate_reads.py b/bin/concatenate_reads.py index 6ccb8ac3..dc4172fa 100755 --- a/bin/concatenate_reads.py +++ b/bin/concatenate_reads.py @@ -22,74 +22,91 @@ def enforce_headers(f1_header, f2_header): sys.exit(8) +FLUSH_BYTES = 1 << 20 + + +def _close_out(f_out): + if f_out is not sys.stdout: + f_out.close() + + def join_w_separator(f1, f2, f_out, strict=False, sep="|"): - readlines = True ix = 0 # Use 'G', which is in most Phred scores. # http://en.wikipedia.org/wiki/FASTQ_format phred_sep = "|" * len(sep) - while readlines: + buf = [] + buf_len = 0 + + while True: f1_line = f1.readline() f2_line = f2.readline() if f1_line == "": - readlines = False + if buf: + f_out.write("".join(buf)) f1.close() f2.close() - f_out.close() + _close_out(f_out) if f2_line != "": raise Exception("Input FASTQ files do not match in length.") - continue + return if ix % 4 == 0: # Header if strict: # Fail if they don't match up to the first whitespace enforce_headers(f1_line, f2_line) # Write the header out - f_out.write(f1_line) + line = f1_line elif (ix - 1) % 4 == 0: # Sequence - f_out.write(f1_line.strip() + sep + f2_line) + line = f1_line.strip() + sep + f2_line elif (ix - 2) % 4 == 0: # Separator - f_out.write("+\n") + line = "+\n" else: - f_out.write(f1_line.strip() + phred_sep + f2_line) + line = f1_line.strip() + phred_sep + f2_line + + buf.append(line) + buf_len += len(line) + if buf_len >= FLUSH_BYTES: + f_out.write("".join(buf)) + buf = [] + buf_len = 0 ix += 1 def join_interleaved(f1, f2, f_out, strict=False): - readlines = True ix = 0 f1_lines = [] f2_lines = [] + buf = [] + buf_len = 0 - def flush_buffer(f_out, f1_lines, f2_lines): + def flush_record(f1_lines, f2_lines): + nonlocal buf, buf_len if f1_lines and f2_lines: assert len(f1_lines) == 4 assert len(f2_lines) == 4 - f_out.write(f1_lines[0]) - f_out.write(f1_lines[1]) - f_out.write(f1_lines[2]) - f_out.write(f1_lines[3]) - f_out.write(f2_lines[0]) - f_out.write(f2_lines[1]) - f_out.write(f2_lines[2]) - f_out.write(f2_lines[3]) + for line in f1_lines: + buf.append(line) + buf_len += len(line) + for line in f2_lines: + buf.append(line) + buf_len += len(line) f1_lines = [] f2_lines = [] return f1_lines, f2_lines - while readlines: + while True: f1_line = f1.readline() f2_line = f2.readline() if f1_line == "": - readlines = False if f2_line != "": raise Exception("Input FASTQ files do not match in length.") break @@ -98,17 +115,23 @@ def flush_buffer(f_out, f1_lines, f2_lines): if strict: enforce_headers(f1_line, f2_line) - f1_lines, f2_lines = flush_buffer(f_out, f1_lines, f2_lines) + f1_lines, f2_lines = flush_record(f1_lines, f2_lines) + if buf_len >= FLUSH_BYTES: + f_out.write("".join(buf)) + buf = [] + buf_len = 0 # Fill buffer up to 4 lines f1_lines.append(f1_line) f2_lines.append(f2_line) ix += 1 - _, _ = flush_buffer(f_out, f1_lines, f2_lines) + flush_record(f1_lines, f2_lines) + if buf: + f_out.write("".join(buf)) f1.close() f2.close() - f_out.close() + _close_out(f_out) if __name__ == "__main__": @@ -151,18 +174,18 @@ def flush_buffer(f_out, f1_lines, f2_lines): if ".gz" in args.fastq1 or ".gzip" in args.fastq1: f1 = gzip.open(args.fastq1, mode="rt") else: - f1 = open(args.fastq1, mode="rt") + f1 = open(args.fastq1, mode="rt", buffering=FLUSH_BYTES) if ".gz" in args.fastq2 or ".gzip" in args.fastq2: f2 = gzip.open(args.fastq2, mode="rt") else: - f2 = open(args.fastq2, mode="rt") + f2 = open(args.fastq2, mode="rt", buffering=FLUSH_BYTES) if args.output_fastq is None: f_out = sys.stdout elif args.gzip: - f_out = gzip.open(args.output_fastq, mode="w") + f_out = gzip.open(args.output_fastq, mode="wt") else: - f_out = open(args.output_fastq, mode="w") + f_out = open(args.output_fastq, mode="w", buffering=FLUSH_BYTES) if not args.no_interleave: join_interleaved(f1, f2, f_out, strict=args.strict) diff --git a/bin/extract_fraction_from_reads.py b/bin/extract_fraction_from_reads.py index 9d0775be..c17255d7 100755 --- a/bin/extract_fraction_from_reads.py +++ b/bin/extract_fraction_from_reads.py @@ -2,6 +2,7 @@ import sys import argparse +import json as json_module from datetime import datetime from extract_utils import * @@ -26,52 +27,107 @@ def setup_prefixes( return outprefix -def extract_reads( - read_map, - taxon_id_map, - entries, - reads1, - reads2, - prefix, - taxon_ids, - exclude, - include_unclassified, -): - # check read files - filetype, zipped = check_read_files(reads1) - - prefixes = setup_prefixes( - taxon_ids, - entries, - prefix, - inverse=exclude, - include_unclassified=include_unclassified, +def build_fraction_spec(cfg, loaded_taxonomy, taxonomy_dir): + """ + Build the (taxon_id_map, prefixes) pieces for one fraction config entry - the + taxonomy-dependent setup that's identical whether this fraction is extracted alone or as part + of a shared multi-fraction pass. + + Args: + cfg (dict): A fraction config with keys 'prefix', 'taxid' (list), 'exclude' (bool) and + 'include_unclassified' (bool). + loaded_taxonomy (Taxonomy): A Taxonomy object with parent/child relationships already + loaded. `entries` is populated (accumulated across calls) as + needed by this function. + taxonomy_dir (str): The unzipped directory downloaded from NCBI taxonomy. + + Returns: + taxon_id_map (dict): Taxon ID (and descendants) to set of taxon_ids in cfg["taxid"]. + prefixes (dict): Key (usually taxon_id) to output filename prefix for this fraction. + """ + taxon_id_map = loaded_taxonomy.get_taxon_id_map( + cfg["taxid"], cfg["include_unclassified"] ) - out_counts, quals, lens, filenames, total_length = process_read_files( - prefixes, - filetype, - read_map, - taxon_id_map, - reads1, - reads2, - inverse=exclude, - get_handles=True, + # setup_prefixes only ever looks up entries for the originally-requested taxid (and not at + # all when exclude is set, since it returns early), so loading taxonomy info for the whole + # descendant closure in taxon_id_map - which can be the entire nodes.dmp file - is wasted work. + if not cfg["exclude"]: + loaded_taxonomy.load_entries(taxonomy_dir, cfg["taxid"]) + prefixes = setup_prefixes( + cfg["taxid"], + loaded_taxonomy.entries, + cfg["prefix"], + inverse=cfg["exclude"], + include_unclassified=cfg["include_unclassified"], ) - - generate_summary( - taxon_ids, - entries, - prefix, - out_counts, - quals, - lens, - filenames, - total_length, - include_unclassified=(include_unclassified != exclude), - short=True, + return taxon_id_map, prefixes + + +def extract_fractions(fraction_configs, kraken_assignment, reads1, reads2, taxonomy_dir): + """ + Extract one or more independent fractions in a single pass over the kraken assignment file + and read file(s) - N separate fraction configs are not a partition of the read set (a read + can legitimately belong to more than one fraction, e.g. a viral read is also "not human"), so + this keeps each fraction's read_map/taxon_id_map/output files completely independent; only + the (expensive, once-per-read) assignment file parse and FASTQ decompress+parse are shared. + + Args: + fraction_configs (list): A list of fraction config dicts, each with keys 'prefix', + 'taxid' (list), 'exclude' (bool) and 'include_unclassified' (bool). + kraken_assignment (KrakenAssignments): The loaded kraken assignment file. + reads1 (str): FASTA/FASTQ file containing the raw reads. + reads2 (str): 2nd FASTA/FASTQ file containing the raw reads (paired), or "" if unpaired. + taxonomy_dir (str): The unzipped directory downloaded from NCBI taxonomy. + + Returns: + dict: fraction prefix -> out_counts (taxon_id -> read count) for that fraction. + """ + filetype, zipped = check_read_files(reads1) + loaded_taxonomy = Taxonomy(taxonomy_dir) + + taxon_id_maps = [] + prefixes_list = [] + for cfg in fraction_configs: + taxon_id_map, prefixes = build_fraction_spec(cfg, loaded_taxonomy, taxonomy_dir) + taxon_id_maps.append(taxon_id_map) + prefixes_list.append(prefixes) + + read_maps = kraken_assignment.get_read_maps(taxon_id_maps, paired=bool(reads2)) + + fraction_specs = [ + { + "prefixes": prefixes_list[i], + "read_map": read_maps[i], + "subtaxa_map": taxon_id_maps[i], + "inverse": fraction_configs[i]["exclude"], + } + for i in range(len(fraction_configs)) + ] + + results, total_length = process_read_files_multi( + fraction_specs, filetype, reads1, reads2 ) - return out_counts + out_counts_by_fraction = {} + for i, cfg in enumerate(fraction_configs): + out_counts, read_stats, filenames = results[i] + generate_summary( + cfg["taxid"], + loaded_taxonomy.entries, + cfg["prefix"], + out_counts, + read_stats, + filenames, + total_length, + include_unclassified=(cfg["include_unclassified"] != cfg["exclude"]), + short=True, + # total_length is the length of the whole input file - the same value for every + # fraction here, so only the first fraction's call needs to write it out. + write_total_length=(i == 0), + ) + out_counts_by_fraction[cfg["prefix"]] = out_counts + + return out_counts_by_fraction # Main method @@ -110,8 +166,8 @@ def main(): "-p", "--prefix", dest="prefix", - required=True, - help="Prefix for output files", + required=False, + help="Prefix for output files. Mutually exclusive with --fraction_config.", ) parser.add_argument( @@ -137,35 +193,48 @@ def main(): default=False, help="List of taxonomy ID[s] or names to exclude (space-delimited) from outputs", ) + parser.add_argument( + "--fraction_config", + dest="fraction_config", + required=False, + help=( + "Path to a JSON file listing multiple fractions to extract in a single pass over the " + "read file(s), instead of a single -p/--taxid/--exclude/--include_unclassified " + "fraction. Each entry is an object with keys 'prefix', 'taxid' (list of taxonomy IDs " + "or names), 'exclude' (bool) and 'include_unclassified' (bool). Mutually exclusive " + "with -p/--prefix." + ), + ) parser.set_defaults(append=False) args = parser.parse_args() + if bool(args.prefix) == bool(args.fraction_config): + parser.error( + "Exactly one of -p/--prefix or --fraction_config must be provided." + ) + # Start Program now = datetime.now() time = now.strftime("%m/%d/%Y, %H:%M:%S") sys.stderr.write("PROGRAM START TIME: " + time + "\n") - loaded_taxonomy = Taxonomy(args.taxonomy) - taxon_id_map = loaded_taxonomy.get_taxon_id_map( - args.taxid, args.include_unclassified - ) - loaded_taxonomy.load_entries(args.taxonomy, taxon_id_map.keys()) + if args.fraction_config: + with open(args.fraction_config) as f: + fraction_configs = json_module.load(f) + else: + fraction_configs = [ + { + "prefix": args.prefix, + "taxid": args.taxid, + "exclude": args.exclude, + "include_unclassified": args.include_unclassified, + } + ] - # Initialize kraken assignment file kraken_assignment = KrakenAssignments(args.kraken_assignment_file) - read_map = kraken_assignment.get_read_map(taxon_id_map) - - out_counts = extract_reads( - read_map, - taxon_id_map, - loaded_taxonomy.entries, - args.reads1, - args.reads2, - args.prefix, - args.taxid, - args.exclude, - args.include_unclassified, + out_counts_by_fraction = extract_fractions( + fraction_configs, kraken_assignment, args.reads1, args.reads2, args.taxonomy ) now = datetime.now() @@ -174,8 +243,9 @@ def main(): sys.stderr.write("READ COUNTS: \n") - for taxon in out_counts: - sys.stderr.write("%s: %i\n" % (taxon, out_counts[taxon])) + for prefix, out_counts in out_counts_by_fraction.items(): + for taxon in out_counts: + sys.stderr.write("%s/%s: %i\n" % (prefix, taxon, out_counts[taxon])) sys.exit(0) diff --git a/bin/extract_taxa_from_reads.py b/bin/extract_taxa_from_reads.py index 9951ed61..4ee236d3 100755 --- a/bin/extract_taxa_from_reads.py +++ b/bin/extract_taxa_from_reads.py @@ -45,7 +45,7 @@ def get_taxon_id_lists( entry = kraken_report.entries[taxon] pass_count_thresh = True pass_perc_thresh = True - if len(target_ranks) > 0 and entry.rank not in target_ranks: + if len(target_ranks) > 0 and entry.simple_rank not in target_ranks: continue if min_count and entry.ucount < min_count: pass_count_thresh = False @@ -59,6 +59,15 @@ def get_taxon_id_lists( pass_perc_thresh = False if len(names) > 0 and entry.name not in names and taxon not in names: continue + if entry.rank != entry.simple_rank: + parent = entry.parent + while ( + kraken_report.entries[parent].rank + != kraken_report.entries[parent].simple_rank + ): + parent = kraken_report.entries[parent].parent + if parent in lists_to_extract: + continue if not pass_count_thresh and not pass_perc_thresh: continue @@ -68,7 +77,7 @@ def get_taxon_id_lists( while len(lookup) > 0: parent = lookup.pop() if parent in loaded_taxonomy.parents and parent != "1": - lookup.append(loaded_taxonomy.parents[lookup]) + lookup.append(loaded_taxonomy.parents[parent]) if parent in kraken_report.entries and parent != "1": lookup.append(kraken_report.entries[parent].parent) if parent != "1": @@ -108,9 +117,7 @@ def setup_prefixes(lists_to_extract, prefix=None): return outprefix -def extract_taxa( - kraken_report, lists_to_extract, kraken_assignment, reads1, reads2, prefix -): +def extract_taxa(entries, lists_to_extract, kraken_assignment, reads1, reads2, prefix): # open read files filetype, zipped = check_read_files(reads1) @@ -123,10 +130,10 @@ def extract_taxa( # "INCLUDING PARENTS/CHILDREN, HAVE %i TAXA TO INCLUDE IN READ FILES for %s\n" # % (len(lists_to_extract[taxon]), taxon) # ) - read_map = kraken_assignment.get_read_map(subtaxa_map) + read_map = kraken_assignment.get_read_map(subtaxa_map, paired=bool(reads2)) prefixes = setup_prefixes(lists_to_extract, prefix) - out_counts, quals, lens, filenames, total_length = process_read_files( + out_counts, read_stats, filenames, total_length = process_read_files( prefixes, filetype, read_map, @@ -134,16 +141,14 @@ def extract_taxa( reads1, reads2, inverse=False, - get_handles=False, ) generate_summary( lists_to_extract, - kraken_report.entries, + entries, prefix, out_counts, - quals, - lens, + read_stats, filenames, total_length, short=False, @@ -151,6 +156,58 @@ def extract_taxa( return out_counts +def load_report_and_extract( + report_file, + loaded_taxonomy, + taxid, + target_ranks, + min_count, + min_count_descendants, + min_percent, + top_n, + include_parents, + include_children, + max_human, + merged_lists_to_extract, + merged_entries, +): + """ + Load a single kraken report, apply the human-count guard, identify the taxon ID lists to extract from it, and + merge the result into the (possibly already-populated) merged_lists_to_extract/merged_entries structures. Used + to combine several kreport splits (e.g. Bacteria/Viruses/Metazoa/default) into a single extraction pass. + + A taxon can legitimately appear in more than one split report (split_kraken_report.py copies ancestor lines into + each split that needs them), so results are merged via set union rather than overwritten. Entries are merged + first-report-wins, since the same taxon_id should describe the same name/rank regardless of which split it was + read from. + """ + sys.stderr.write(f"Loading kraken report {report_file}\n") + kraken_report = KrakenReport(report_file) + if max_human: + # Exits the whole process immediately (sys.exit(2)) if this report's + # human read count exceeds the threshold - this intentionally fails + # the entire combined extraction, not just this one report's share. + kraken_report.check_host({"9606": max_human}) + + lists_to_extract = get_taxon_id_lists( + kraken_report, + loaded_taxonomy, + names=taxid, + target_ranks=target_ranks, + min_count=min_count, + min_count_descendants=min_count_descendants, + min_percent=min_percent, + top_n=top_n, + include_parents=include_parents, + include_children=include_children, + ) + + for taxon, subtaxa in lists_to_extract.items(): + merged_lists_to_extract[taxon] |= subtaxa + for taxon_id, entry in kraken_report.entries.items(): + merged_entries.setdefault(taxon_id, entry) + + def check_out_counts(out_counts, kraken_report): for taxon_id in out_counts: if out_counts[taxon_id] != kraken_report.entries[taxon_id].count: @@ -173,8 +230,20 @@ def main(): parser.add_argument( "-r", dest="report_file", - required=True, - help="Kraken or Bracken file of taxon relationships and quantities", + required=False, + help="Kraken or Bracken file of taxon relationships and quantities. Mutually exclusive with --report_config.", + ) + parser.add_argument( + "--report_config", + dest="report_config", + required=False, + help=( + "Path to a JSON file listing multiple kreport splits to extract from in a single pass over the read " + "file(s), instead of a single -r report. Each entry is an object with keys 'report' (path), 'rank' " + "(a rank code or list of rank codes, same as --rank), 'min_count_descendants' and 'min_percent'. " + "Results across all reports are merged (taxa unioned, first-report-wins on name/rank metadata). " + "Mutually exclusive with -r." + ), ) parser.add_argument( "-t", @@ -295,6 +364,11 @@ def main(): "G": "G", "S": "S", } + if bool(args.report_file) == bool(args.report_config): + parser.error( + "Exactly one of -r/--report_file or --report_config must be provided." + ) + if args.rank: target_ranks = [rank_dict[r] for r in args.rank] else: @@ -305,34 +379,61 @@ def main(): sys.stderr.write("Loading taxonomy\n") loaded_taxonomy = Taxonomy(args.taxonomy) - # Load kraken report entries - sys.stderr.write("Loading kraken report\n") - kraken_report = KrakenReport(args.report_file) - if args.max_human: - kraken_report.check_host({"9606": args.max_human}) - # Initialize kraken assignment file sys.stderr.write("Loading kraken assignments\n") kraken_assignment = KrakenAssignments(args.kraken_assignment_file) + merged_lists_to_extract = defaultdict(set) + merged_entries = {} + sys.stderr.write("Identifying lists to extract\n") - lists_to_extract = get_taxon_id_lists( - kraken_report, - loaded_taxonomy, - names=args.taxid, - target_ranks=target_ranks, - min_count=args.min_count, - min_count_descendants=args.min_count_descendants, - min_percent=args.min_percent, - top_n=args.top_n, - include_parents=args.include_parents, - include_children=args.include_children, - ) + if args.report_config: + with open(args.report_config) as f: + report_config = json.load(f) + for report_entry in report_config: + entry_rank = report_entry.get("rank") + if entry_rank is None: + entry_target_ranks = target_ranks + elif isinstance(entry_rank, list): + entry_target_ranks = [rank_dict[r] for r in entry_rank] + else: + entry_target_ranks = [rank_dict[entry_rank]] + load_report_and_extract( + report_entry["report"], + loaded_taxonomy, + args.taxid, + entry_target_ranks, + args.min_count, + report_entry.get("min_count_descendants"), + report_entry.get("min_percent"), + args.top_n, + args.include_parents, + args.include_children, + args.max_human, + merged_lists_to_extract, + merged_entries, + ) + else: + load_report_and_extract( + args.report_file, + loaded_taxonomy, + args.taxid, + target_ranks, + args.min_count, + args.min_count_descendants, + args.min_percent, + args.top_n, + args.include_parents, + args.include_children, + args.max_human, + merged_lists_to_extract, + merged_entries, + ) sys.stderr.write("Extracting reads from file\n") out_counts = extract_taxa( - kraken_report, - lists_to_extract, + merged_entries, + merged_lists_to_extract, kraken_assignment, args.reads1, args.reads2, diff --git a/bin/extract_utils.py b/bin/extract_utils.py index 7a979327..d434c89c 100755 --- a/bin/extract_utils.py +++ b/bin/extract_utils.py @@ -4,7 +4,7 @@ import gzip import pyfastx import json -from collections import defaultdict +from collections import defaultdict, OrderedDict from pathlib import Path import statistics as stats @@ -21,6 +21,19 @@ def ascii_to_phred(q: str, offset: int = 33) -> np.ndarray: return arr - offset +def mean_phred(q: str, offset: int = 33) -> float: + """ + Mean Phred quality for a single FASTQ quality string. Unlike np.mean(ascii_to_phred(q)), + this never materialises a full subtracted array - summing the raw ASCII byte values + (numpy upcasts the accumulator, so this doesn't overflow) and subtracting the offset once + from the mean is the same value, computed in about a third of the time. + """ + if not q: + return 0.0 + arr = np.frombuffer(q.encode("ascii"), dtype=np.uint8) + return float(arr.sum()) / arr.size - offset + + def record_to_fastq_entry(record): """ Converts a record to a FASTQ entry string. @@ -63,145 +76,157 @@ def check_read_files(reads): return filetype, zipped -def setup_outfiles(fastq_2, prefixes, filetype, get_handles=True): +def setup_outfiles(fastq_2, prefixes, filetype): """ - Sets up the output filenames, optionally opening them and returning handles. + Sets up the output filenames for each key (usually a taxon ID). Args: fastq_2 (bool): True or a string if fastq_2 exists, False or None if reads were not paired. prefixes (dict): A dictionary from the key (usually taxon_id) to the prefix of the output file for that key. filetype (str): "fasta" or "fastq" depending on input filetype. - get_handles (bool): Whether to open file handles for the output files. + + Every returned file is also created empty, so that a key with zero extracted + reads still leaves a file behind: the Nextflow processes declare globbed + outputs (e.g. `*.f*q.gz`), which a missing file - unlike an empty one - fails + to match, taking the task down the retry path instead of publishing an empty + result. Returns: filenames (dict): Dictionary with keys from the prefixes, to a list of output filenames. - out_handles_1 (dict): Key (matching filenames) to out_handle (if opened) - out_handles_2 (dict): Key (matching filenames) to out_handle (if opened and paired reads) """ - - out_handles_1 = {} - out_handles_2 = {} filenames = defaultdict(list) for key, prefix in prefixes.items(): if fastq_2: filenames[key].append(f"{prefix}_1.{filetype}") filenames[key].append(f"{prefix}_2.{filetype}") - if get_handles: - if key in out_handles_1: - print("already open") - else: - out_handles_1[key] = open(f"{prefix}_1.{filetype}", "w") - out_handles_2[key] = open(f"{prefix}_2.{filetype}", "w") - else: filenames[key].append(f"{prefix}.{filetype}") - if get_handles: - if key in out_handles_1: - print("already open") - else: - out_handles_1[key] = open(f"{prefix}.{filetype}", "w") - return filenames, out_handles_1, out_handles_2 + for paths in filenames.values(): + for path in paths: + open(path, "w").close() + return filenames -def close_outfiles(out_handles_1, out_handles_2): - """ - Checks each handle in the dictionary and closes if it is open. - Args: - out_handles_1 (dict): Key to out_handle - out_handles_1 (dict): Key to out_handle (if paired reads) - """ - for handle in out_handles_1: - out_handles_1[handle].close() - for handle in out_handles_2: - out_handles_2[handle].close() +def _default_max_open_files(): + try: + import resource + soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + return max(16, min(128, soft // 4)) + except Exception: + return 64 -def update_summary_with_record(taxon_id, record, out_counts, quals, lens): - """ - Updates the summary dictionaries for counts, qualities and lengths of reads from the new record Dictionaries indexed by taxon. - Args: - taxon_id (str): Key for dictionaries - record (SeqRecord): A read parsed by Bio.SeqIO - out_counts (dict): Taxon ID to read count for that taxon. - quals (dict): Taxon ID to list of quality scores (each averaged over the read) for that taxon. - lens (dict): Taxon ID to list of read lengths for that taxon. +class TaxonWriter: + """ + Buffers extracted FASTQ/FASTA records per taxon ID in memory and periodically flushes them to disk. Output file + handles are pooled with a bounded LRU cache so the number of simultaneously open files never exceeds `max_open`, + regardless of how many distinct taxon IDs are being written to - this avoids hitting the OS open-file-descriptor + limit while also avoiding buffering every read for every taxon in memory until the end of the file (as previously + happened for the low-memory=False code path). + + Records for a given taxon are always flushed in the order they were written, so read order within each output + file is preserved exactly as before. """ - out_counts[taxon_id] += 1 - name, seq, qual = record - - mean_qual = np.mean(ascii_to_phred(qual)) if qual else 0 - - quals[taxon_id].append(mean_qual) - lens[taxon_id].append(len(seq)) + def __init__( + self, + filenames, + file_index, + max_open=None, + per_taxon_bytes=1 << 20, + total_bytes=512 << 20, + ): + self.filenames = filenames + self.file_index = file_index + self.max_open = max_open if max_open else _default_max_open_files() + self.per_taxon_bytes = per_taxon_bytes + self.total_bytes = total_bytes + + self.buffers = defaultdict(list) + self.buffer_bytes = defaultdict(int) + self.total_buffered = 0 + self.handles = OrderedDict() # taxon_id -> open file handle, LRU order + + def write(self, taxon_id, record): + entry = record_to_fastq_entry(record) + self.buffers[taxon_id].append(entry) + n = len(entry) + self.buffer_bytes[taxon_id] += n + self.total_buffered += n + + if self.buffer_bytes[taxon_id] >= self.per_taxon_bytes: + self._flush(taxon_id) + elif self.total_buffered >= self.total_bytes: + self._flush_largest() + + def _flush_largest(self): + while self.total_buffered >= self.total_bytes and self.buffer_bytes: + taxon_id = max(self.buffer_bytes, key=self.buffer_bytes.get) + self._flush(taxon_id) + + def _flush(self, taxon_id): + if not self.buffers[taxon_id]: + return + handle = self._handle(taxon_id) + handle.write("".join(self.buffers[taxon_id])) + self.total_buffered -= self.buffer_bytes[taxon_id] + self.buffers[taxon_id] = [] + self.buffer_bytes[taxon_id] = 0 + + def _handle(self, taxon_id): + if taxon_id in self.handles: + self.handles.move_to_end(taxon_id) + return self.handles[taxon_id] + + if len(self.handles) >= self.max_open: + _, evicted = self.handles.popitem(last=False) + evicted.close() + + # Append mode is required, not just safe: `setup_outfiles` has already + # created every declared path empty, and a handle may be evicted from the + # LRU cache and reopened part-way through writing a taxon. + path = self.filenames[taxon_id][self.file_index] + handle = open(path, "a") + self.handles[taxon_id] = handle + return handle + + def close(self): + for taxon_id in list(self.buffers.keys()): + self._flush(taxon_id) + for handle in self.handles.values(): + handle.close() + self.handles.clear() + + +def new_taxon_stats(): + """ + A fresh, empty running-stats accumulator for a single taxon. + """ + return {"count": 0, "sum_qual": 0.0, "sum_len": 0} -def add_record( - taxon_id, - record, - out_counts, - quals, - lens, - filenames, - file_index, - out_handles={}, - out_records={}, - buffer_record_size=None, -): +def update_summary_with_record(taxon_id, record, out_counts, read_stats): """ - Add the record to the required out file and update the summary statistics with it. If out_handles has open - out_handles, reads will be written directly to the handle. If not it will be stored in memory in the out_records. If - a buffer size is set and the out_records structure has at least this number of reads, the out_records for this taxon - ID will be appended to the relevant file and cleared. + Updates the summary dictionaries for counts, qualities and lengths of reads from the new record. Rather than + keeping every individual read's quality/length, a running count/sum per taxon is kept - this bounds memory use + regardless of how many reads are seen for a given taxon. Args: taxon_id (str): Key for dictionaries record (SeqRecord): A read parsed by Bio.SeqIO out_counts (dict): Taxon ID to read count for that taxon. - quals (dict): Taxon ID to list of quality scores (each averaged over the read) for that taxon. - lens (dict): Taxon ID to list of read lengths for that taxon. - filenames (dict): Taxon ID to list of output filenames. - file_index (int): 0 or 1, depending if processing forward or reverse read, index for the filenames list. - out_handles (dict): {} if handles closed, Taxon ID to open handle if open. - out_records (dict): {} if handles open, Taxon ID to list of records if not. - buffer_record_size (int): Number of records per taxon ID before they should be output to reduce memory use. + read_stats (dict): Taxon ID to a running {"count", "sum_qual", "sum_len"} accumulator for that taxon. """ - - update_summary_with_record(taxon_id, record, out_counts, quals, lens) - + out_counts[taxon_id] += 1 name, seq, qual = record - if out_handles != {} and taxon_id in out_handles: - out_handles[taxon_id].write(record_to_fastq_entry(record)) - else: - if not buffer_record_size or len(out_records[taxon_id]) < buffer_record_size: - out_records[taxon_id].append(record) - else: - with open(filenames[taxon_id][file_index], "a") as f: - for existing_record in out_records[taxon_id]: - f.write(record_to_fastq_entry(existing_record)) - del out_records[taxon_id] - out_records[taxon_id].append(record) - - -def save_records_to_file(out_records, filenames, file_index): - """ - Save the records to file. - - Args: - out_records (dict): Taxon ID to list of records - filenames (dict): Taxon ID to list of output filenames - file_index (int): 0 or 1, depending if processing forward or reverse read, index for the filenames list. - """ - sys.stderr.write(f"Writing records for file {file_index+1}\n") - for taxon, records in out_records.items(): - with open(filenames[taxon][file_index], "a") as f: - for record in records: - f.write(record_to_fastq_entry(record)) - del out_records + taxon_stats = read_stats[taxon_id] + taxon_stats["count"] += 1 + taxon_stats["sum_qual"] += mean_phred(qual) + taxon_stats["sum_len"] += len(seq) def file_iterator( @@ -210,15 +235,12 @@ def file_iterator( subtaxa_map, inverse, out_counts, - quals, - lens, - filenames, - file_index, - out_handles, - low_memory=False, + read_stats, + writer, ): """ - Iterate through the read_file file and add the read to the appropriate file or handle. + Iterate through the read_file file and add the read to the appropriate output file, via the writer, updating + summary statistics as it goes. Args: read_file (str): Name of read file. @@ -227,71 +249,171 @@ def file_iterator( out files to be updated inverse (bool): If True, exclude all reads which match the taxon IDs in the read_map. If False, select them. out_counts (dict): Taxon ID to read count for that taxon. - quals (dict): Taxon ID to list of quality scores (each averaged over the read) for that taxon. - lens (dict): Taxon ID to list of read lengths for that taxon. - filenames (dict): Taxon ID to list of output filenames - file_index (int): 0 or 1, depending if processing forward or reverse read, index for the filenames list. - out_handles (dict): {} if handles closed, Taxon ID to open handle if open. - low_memory (bool): If True we expect to write directly to out handles, if False store records then write to file. + read_stats (dict): Taxon ID to running {"count", "sum_qual", "sum_len"} accumulator for that taxon. + writer (TaxonWriter): Buffered, bounded-filehandle writer for this file (forward or reverse read). Returns: int: Count of reads written to file. """ - reads_of_interest = set(read_map.keys()) count = 0 total_length = 0 - out_records = None - if not low_memory: - out_records = defaultdict(list) sys.stderr.write(f"Reading in {read_file}\n") for record in pyfastx.Fastq(read_file, build_index=False): name, seq, qual = record total_length += len(seq) trimmed_name = trim_read_id(name) + # A single dict.get() replaces the old `set(read_map.keys())` membership + # check followed by a second read_map[...] lookup - read_map is a + # defaultdict(str), so .get() (which never triggers the default + # factory) is used rather than [] to distinguish "not present" (None) + # from a stored value. + taxon_id = read_map.get(trimmed_name) if inverse: - if trimmed_name in reads_of_interest: - taxon_id = read_map[trimmed_name] + if taxon_id is not None: for taxon in subtaxa_map[taxon_id]: - update_summary_with_record(taxon, record, out_counts, quals, lens) + update_summary_with_record(taxon, record, out_counts, read_stats) continue else: count += 1 - add_record( - "other", - record, - out_counts, - quals, - lens, - filenames, - file_index, - out_handles=out_handles, - out_records=out_records, - ) + update_summary_with_record("other", record, out_counts, read_stats) + writer.write("other", record) elif not inverse: - if trimmed_name not in reads_of_interest: + if taxon_id is None: continue - taxon_id = read_map[trimmed_name] for taxon in subtaxa_map[taxon_id]: count += 1 - add_record( - taxon, - record, - out_counts, - quals, - lens, - filenames, - file_index, - out_handles=out_handles, - out_records=out_records, - ) + update_summary_with_record(taxon, record, out_counts, read_stats) + writer.write(taxon, record) - if not low_memory: - save_records_to_file(out_records, filenames, file_index) + writer.close() return count, total_length +def file_iterator_multi(read_file, contexts): + """ + Like file_iterator, but applies several independent fraction contexts to a single pass over + read_file, instead of requiring one pass per context. Each context is a dict with the same + per-fraction state file_iterator would otherwise take as separate arguments: 'read_map', + 'subtaxa_map', 'inverse', 'out_counts', 'read_stats', 'writer'. Every context's inclusion/ + exclusion logic is applied to every record exactly as file_iterator would for that context + alone - contexts never interact (each has its own read_map/writer/stats), only the expensive + per-record decompress+parse work is shared. + + Args: + read_file (str): Name of read file. + contexts (list): List of per-fraction context dicts, as above. + Returns: + counts (list): Count of reads written to file, one per context, in the same order. + total_length (int): Total length of all reads in read_file (the same regardless of + fraction membership, since every fraction reads the whole file). + """ + counts = [0] * len(contexts) + total_length = 0 + + sys.stderr.write(f"Reading in {read_file}\n") + for record in pyfastx.Fastq(read_file, build_index=False): + name, seq, qual = record + total_length += len(seq) + trimmed_name = trim_read_id(name) + + for i, ctx in enumerate(contexts): + taxon_id = ctx["read_map"].get(trimmed_name) + if ctx["inverse"]: + if taxon_id is not None: + for taxon in ctx["subtaxa_map"][taxon_id]: + update_summary_with_record( + taxon, record, ctx["out_counts"], ctx["read_stats"] + ) + continue + else: + counts[i] += 1 + update_summary_with_record( + "other", record, ctx["out_counts"], ctx["read_stats"] + ) + ctx["writer"].write("other", record) + else: + if taxon_id is None: + continue + for taxon in ctx["subtaxa_map"][taxon_id]: + counts[i] += 1 + update_summary_with_record( + taxon, record, ctx["out_counts"], ctx["read_stats"] + ) + ctx["writer"].write(taxon, record) + + for ctx in contexts: + ctx["writer"].close() + + return counts, total_length + + +def process_read_files_multi( + fraction_specs: list, + filetype: str, + read_file_1: Path, + read_file_2: Path = None, +) -> tuple[list, int]: + """ + Like process_read_files, but for several independent fractions sharing a single pass over + (paired/unpaired) read_files, instead of one pass per fraction. + + Args: + fraction_specs (list): A list of dicts, one per fraction, each with keys 'prefixes' + (dict), 'read_map' (dict) and 'subtaxa_map' (dict) - the same + per-fraction inputs process_read_files takes for a single fraction. + filetype (str): "fasta" or "fastq" depending on input filetype. + read_file_1 (Path): Name of read file. + read_file_2 (Path): Name of read file pair (if it exists). + + Returns: + results (list): A list of (out_counts, read_stats, filenames) tuples, one per fraction, + in the same order as fraction_specs. + total_length (int): Total length of all reads across the input file(s) - the same value + for every fraction, computed once here instead of once per fraction. + """ + contexts = [] + for spec in fraction_specs: + filenames = setup_outfiles(read_file_2, spec["prefixes"], filetype) + contexts.append( + { + "read_map": spec["read_map"], + "subtaxa_map": spec["subtaxa_map"], + "inverse": spec["inverse"], + "out_counts": defaultdict(int), + "read_stats": defaultdict(new_taxon_stats), + "filenames": filenames, + "writer": TaxonWriter(filenames, 0), + } + ) + + forward_counts, total_length = file_iterator_multi(read_file_1, contexts) + + if read_file_2: + for ctx in contexts: + ctx["writer"] = TaxonWriter(ctx["filenames"], 1) + reverse_counts, reverse_length = file_iterator_multi(read_file_2, contexts) + total_length += reverse_length + for forward_count, reverse_count in zip(forward_counts, reverse_counts): + if forward_count != reverse_count and ( + forward_count == 0 or reverse_count == 0 + ): + sys.stderr.write( + "ERROR: No reads found for one of the file pair: extracted %i an %i reads respectively" + % (forward_count, reverse_count) + ) + sys.exit(7) + + return ( + [ + (ctx["out_counts"], ctx["read_stats"], ctx["filenames"]) + for ctx in contexts + ], + total_length, + ) + + def process_read_files( prefixes: str, filetype: str, @@ -300,7 +422,6 @@ def process_read_files( read_file_1: Path, read_file_2: Path = None, inverse: bool = False, - get_handles: bool = False, ) -> tuple[dict, dict, dict, dict]: """ Iterate through (paired/unpaired) read_files and save the relevant reads, collecting summary statistics for these reads. @@ -314,50 +435,41 @@ def process_read_files( read_file_1 (Path): Name of read file. read_file_2 (Path): Name of read file pair (if it exists). inverse (bool): If True, exclude all reads which match the taxon IDs in the read_map. If False, select them. - get_handles (bool): If True we expect to write directly to out handles, if False store records then write to file. Returns: out_counts (dict): Taxon ID to read count for that taxon. - quals (dict): Taxon ID to list of quality scores (each averaged over the read) for that taxon. - lens (dict): Taxon ID to list of read lengths for that taxon. + read_stats (dict): Taxon ID to running {"count", "sum_qual", "sum_len"} accumulator for that taxon. filenames (dict): Taxon ID to list of output filenames. """ out_counts = defaultdict(int) - quals = defaultdict(list) - lens = defaultdict(list) + read_stats = defaultdict(new_taxon_stats) - filenames, out_handles_1, out_handles_2 = setup_outfiles( - read_file_2, prefixes, filetype, get_handles=get_handles - ) + filenames = setup_outfiles(read_file_2, prefixes, filetype) + # Each writer is scoped to a single read file (forward or reverse) so its + # LRU handle pool doesn't have to span both passes. + writer_1 = TaxonWriter(filenames, 0) forward_count, total_length = file_iterator( read_file_1, read_map, subtaxa_map, inverse, out_counts, - quals, - lens, - filenames, - 0, - out_handles_1, - low_memory=get_handles, + read_stats, + writer_1, ) reverse_count = 0 if read_file_2: + writer_2 = TaxonWriter(filenames, 1) reverse_count, reverse_length = file_iterator( read_file_2, read_map, subtaxa_map, inverse, out_counts, - quals, - lens, - filenames, - 1, - out_handles_2, - low_memory=get_handles, + read_stats, + writer_2, ) total_length += reverse_length if forward_count != reverse_count and ( @@ -368,8 +480,7 @@ def process_read_files( % (forward_count, reverse_count) ) sys.exit(7) - close_outfiles(out_handles_1, out_handles_2) - return (out_counts, quals, lens, filenames, total_length) + return (out_counts, read_stats, filenames, total_length) def generate_summary( @@ -377,12 +488,12 @@ def generate_summary( entries, prefix, out_counts, - quals, - lens, + read_stats, filenames, total_length, include_unclassified=False, short=False, + write_total_length=True, ): """ Generate a summary JSON file, including information about each taxon ID and corresponding read statistics. @@ -393,12 +504,16 @@ def generate_summary( inferred from the Kraken report or NCBI taxonomy. prefix (str): Prefix for the summary file. out_counts (dict): Taxon ID to read count for that taxon. - quals (dict): Taxon ID to list of quality scores (each averaged over the read) for that taxon. - lens (dict): Taxon ID to list of read lengths for that taxon. + read_stats (dict): Taxon ID to running {"count", "sum_qual", "sum_len"} accumulator for that taxon. filenames (dict): Taxon ID to list of output filenames. include_unclassified (bool): True if each output file includes unclassified reads as well as those associated with the taxon ID. short (bool): Output short form summary, excluding taxon information. + write_total_length (bool): Whether to write total_length.json. total_length is the length + of the whole input file, the same value regardless of which + fraction/report is being summarised - when generate_summary is + called once per fraction sharing the same read file(s), only + one of those calls needs to write it. filenames (dict): Taxon ID to list of output filenames. """ sys.stderr.write("Write summary\n") @@ -408,20 +523,29 @@ def generate_summary( sys.stderr.write(f"No reads extracted for taxon_id {taxon_id}\n") continue + taxon_stats = read_stats[taxon_id] + avg_qual = ( + taxon_stats["sum_qual"] / taxon_stats["count"] + if taxon_stats["count"] + else 0 + ) + mean_len = ( + taxon_stats["sum_len"] / taxon_stats["count"] + if taxon_stats["count"] + else 0 + ) + qc_metrics = { + "num_reads": out_counts[taxon_id], + "avg_qual": avg_qual, + "mean_len": mean_len, + "total_len": taxon_stats["sum_len"], + } + if short: summary.append( { "filenames": filenames[taxon_id], - "qc_metrics": { - "num_reads": out_counts[taxon_id], - "avg_qual": ( - stats.fmean(quals[taxon_id]) if quals[taxon_id] else 0 - ), - "mean_len": ( - stats.fmean(lens[taxon_id]) if lens[taxon_id] else 0 - ), - "total_len": sum(lens[taxon_id]), - }, + "qc_metrics": qc_metrics, "includes_unclassified": include_unclassified, } ) @@ -430,18 +554,9 @@ def generate_summary( { "human_readable": entries[taxon_id].name, "taxon_id": taxon_id, - "tax_level": entries[taxon_id].rank, + "tax_level": entries[taxon_id].simple_rank, "filenames": filenames[taxon_id], - "qc_metrics": { - "num_reads": out_counts[taxon_id], - "avg_qual": ( - stats.fmean(quals[taxon_id]) if quals[taxon_id] else 0 - ), - "mean_len": ( - stats.fmean(lens[taxon_id]) if lens[taxon_id] else 0 - ), - "total_len": sum(lens[taxon_id]), - }, + "qc_metrics": qc_metrics, "includes_unclassified": include_unclassified, } ) @@ -449,12 +564,16 @@ def generate_summary( with open(f"{prefix}_summary.json", "w") as f: json.dump(summary, f) - with open("total_length.json", "w") as f: - json.dump({"total_len": total_length}, f) + if write_total_length: + with open("total_length.json", "w") as f: + json.dump({"total_len": total_length}, f) - for taxon_id in quals: - if stats.fmean(quals[taxon_id]) > 60: - sys.exit( - 1, - f"Mean quality score for taxon_id {taxon_id} is {stats.fmean(quals[taxon_id])} > 60. This seems unlikely! Please report this example to the DIPI group for investigation", + for taxon_id in read_stats: + taxon_stats = read_stats[taxon_id] + if taxon_stats["count"] and taxon_stats["sum_qual"] / taxon_stats["count"] > 60: + sys.stderr.write( + f"Mean quality score for taxon_id {taxon_id} is " + f"{taxon_stats['sum_qual'] / taxon_stats['count']} > 60. This seems " + "unlikely! Please report this example to the DIPI group for investigation\n" ) + sys.exit(1) diff --git a/bin/get_total_length.py b/bin/get_total_length.py deleted file mode 100755 index 3c7f3232..00000000 --- a/bin/get_total_length.py +++ /dev/null @@ -1,46 +0,0 @@ -#!/usr/bin/env python - -import sys -import argparse -import pyfastx -import json - - -def main(): - # Parse arguments - parser = argparse.ArgumentParser() - - parser.add_argument( - "-s", - "-s1", - "-1", - dest="reads1", - required=True, - help="FASTA/FASTQ File containing the raw reads.", - ) - parser.add_argument( - "-s2", - "-2", - dest="reads2", - default="", - help="2nd FASTA/FASTQ File containing the raw reads (paired).", - ) - - parser.set_defaults(append=False) - args = parser.parse_args() - - fq = pyfastx.Fastq(args.reads1) - total_length = fq.size - - if (args.reads2): - fq = pyfastx.Fastq(args.reads2) - total_length += fq.size - - with open(f"total_length.json", "w") as f: - json.dump({"total_len": total_length}, f) - - sys.exit(0) - - -if __name__ == "__main__": - main() diff --git a/bin/make_report.py b/bin/make_report.py index 78d7b5d2..43f971bb 100755 --- a/bin/make_report.py +++ b/bin/make_report.py @@ -12,6 +12,8 @@ import json import os +from ranks import DOMAIN_RANKS, RANKS + def make_output_report( report_to_generate, template, version, sample, data_for_report={} @@ -161,6 +163,10 @@ def main(): "classifier": args.classifier, "classification_database": args.classification_database, "warnings": warnings, + # The report's JS shares this vocabulary with aggregate_lineages_bracken.py + # rather than keeping its own copy - see bin/ranks.py. + "ranks": json.dumps(RANKS), + "domain_ranks": json.dumps(DOMAIN_RANKS), } outfile = args.prefix + "_report.html" diff --git a/bin/ranks.py b/bin/ranks.py new file mode 100644 index 00000000..540c6a6e --- /dev/null +++ b/bin/ranks.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python +""" +Single source of truth for the taxonomic rank vocabulary used across the report. + +NCBI restructured its top-level ranks in 2024: `superkingdom` was retired, with +Bacteria/Archaea/Eukaryota becoming rank `domain` and Viruses becoming +`acellular root` (with `realm` clades such as Duplodnaviria beneath it). Both the +old and new names are listed here so a lineage tops out at a real domain under +either an old or a current taxonomy dump. + +This module is stdlib-free so it can be imported by scripts running in minimal +task containers. It is consumed by: + - aggregate_lineages_bracken.py, to decide which lineage nodes to keep; + - report.py, for the kraken-report domain-name heuristic; + - make_report.py, which injects RANKS and DOMAIN_RANKS into the report's + JavaScript so the two sides cannot drift apart. +""" + +# Ordered coarse -> fine. `cellular root` (taxid 131567, "cellular organisms") is +# deliberately absent: including it would collapse Bacteria, Archaea and +# Eukaryota under a single meaningless top-level node. Rank `no rank` (taxid 1, +# "root") is likewise absent. +RANKS = [ + "acellular root", + "superkingdom", + "domain", + "realm", + "clade", + "kingdom", + "phylum", + "subphylum", + "class", + "order", + "family", + "genus", + "species", + "subspecies", + "serotype", +] + +# The ranks that mark a domain boundary, i.e. the top level of a lineage. A node +# at one of these ranks names the domain its descendants belong to. +DOMAIN_RANKS = ["acellular root", "superkingdom", "domain"] + +# The rank reported for the synthetic "Unclassified" node. It is rendered as a +# top-level node alongside the real domains, so it must be one of DOMAIN_RANKS +# for the report's domain grouping to pick it up. +UNCLASSIFIED_RANK = "domain" + +# Domain/realm names used to identify the domain boundary in a *kraken report*, +# which carries single-letter rank codes rather than the rank strings above. This +# is a name-based heuristic and is intentionally separate from DOMAIN_RANKS: see +# the comment at its use site in report.py. +KNOWN_DOMAIN_NAMES = {"Bacteria", "Archaea", "Eukaryota", "Viruses"} diff --git a/bin/report.py b/bin/report.py index f6e60b85..c21f3636 100644 --- a/bin/report.py +++ b/bin/report.py @@ -4,6 +4,8 @@ import csv import sys +from ranks import KNOWN_DOMAIN_NAMES + class KrakenEntry: """ @@ -52,6 +54,20 @@ def __eq__(self, other): else: return False + @property + def simple_rank(self): + """ + The rank code collapsed to its base letter (e.g. "G1"/"G2" -> "G", "R1"/"R2" -> "R"). + + Kraken2 appends a number to a rank code for "no rank" clades inserted between two standard + ranks (this became common with the 2023+ NCBI taxonomy restructuring, which added extra + intermediate clades at many levels, not just around the domain). Code that wants to match a + taxon against one of the standard rank letters (D/K/P/C/O/F/G/S) should compare against this + rather than `rank` directly, or numbered variants like "G1" will never match "G". `rank` + itself is left untouched so it round-trips exactly for KrakenReport.save(). + """ + return self.rank[0] if self.rank else self.rank + def print(self): """ Print the attributes of KrakenEntry as a string @@ -204,7 +220,7 @@ def set_sibling_ranks(self): for entry_id, entry in self.entries.items(): if entry.sibling_rank > 0 or entry.parent is None: continue - if entry.rank in ["D", "R", "R1"]: + if entry.simple_rank in ["D", "R"]: entry.set_sibling_rank(1) elif len(self.entries[entry.parent].children) == 1: entry.set_sibling_rank(1) @@ -300,8 +316,22 @@ def load_file(self, file_name): domain = None for row in df: try: - if row["Rank"] == "D": - domain = row["Scientific Name"].strip() + name = row["Scientific Name"].strip() + # Kraken2's "D" rank code reliably marked superkingdom-level + # nodes in older NCBI taxonomy dumps, but the 2023+ taxonomy + # restructuring inserted extra no-rank clades above some + # domains (e.g. a new Kingdom rank within Bacteria), which + # pushes those domain nodes down to a "R1"/"R2"/... no-rank + # code instead of "D" - the rank code alone no longer + # reliably identifies the domain boundary (and naively taking + # "the last R-rank before a standard rank" doesn't work + # either, since some domains like Viruses have their own + # no-rank sub-clades - e.g. Duplodnaviria - below them before + # hitting a standard rank). Matching on the well-known, + # taxonomically-stable domain/realm names directly is robust + # to this rank-code churn either way. + if row["Rank"] == "D" or name in KNOWN_DOMAIN_NAMES: + domain = name self.domains[domain] = row["Taxonomy ID"] entry = KrakenEntry(row=row, domain=domain, hierarchy=hierarchy) @@ -326,14 +356,13 @@ def get_domains(self): Get a list of taxonomic domains found in the report Returns: - list: List of domains + list: List of taxon_ids of the domains found """ - domains = [] - for entry_id, entry in self.entries.items(): - if entry.rank == "D": - domains.append(entry_id) - entry.print() - return domains + # `self.domains` is populated while loading the report and already applies + # the name-based heuristic described there. Re-deriving this from + # `simple_rank == "D"` would miss any domain that modern taxonomy dumps + # push down to an "R1"/"R2" no-rank code. + return list(self.domains.values()) def get_tips(self): """ @@ -346,7 +375,6 @@ def get_tips(self): for entry_id, entry in self.entries.items(): if len(entry.children) == 0 and entry_id != "0": tips.append(entry_id) - entry.print() return tips def get_rank_entries(self, rank): @@ -360,9 +388,8 @@ def get_rank_entries(self, rank): """ subset = [] for entry_id, entry in self.entries.items(): - if entry.rank == rank: + if entry.simple_rank == rank: subset.append(entry_id) - entry.print() return subset def get_percentage(self, taxon_id, denominator="classified"): @@ -382,7 +409,12 @@ def get_percentage(self, taxon_id, denominator="classified"): elif denominator in self.domains: total = self.entries[self.domains[denominator]].count else: - print(f"Not a valid denominator {denominator}") + sys.stderr.write( + f"WARNING: '{denominator}' is not a valid denominator (not 'classified', " + "'total', or a known domain) for taxon_id " + f"{taxon_id} - treating its percentage as 0.0\n" + ) + return 0.0 if ( denominator not in ["classified", "total"] @@ -437,21 +469,7 @@ def to_source_target_df( continue # filter if an intermediate rank - if entry.rank not in [ - "K", - "D", - "D1", - "D2", - "P", - "C", - "O", - "F", - "G", - "G1", - "S", - "S1", - "S2", - ]: + if entry.simple_rank not in ["K", "D", "P", "C", "O", "F", "G", "S"]: skip.add(entry_id) continue @@ -495,7 +513,7 @@ def to_df(self, sample_id="sample_id", ranks=[]): taxon_ids = [e for e in self.entries.keys()] else: taxon_ids = [ - e for e in self.entries.keys() if self.entries[e].rank in ranks + e for e in self.entries.keys() if self.entries[e].simple_rank in ranks ] return pd.DataFrame( {sample_id: [self.entries[e].count for e in taxon_ids]}, index=taxon_ids diff --git a/bin/report_utils/sankey.js b/bin/report_utils/sankey.js index 23b01754..2fc4d22c 100644 --- a/bin/report_utils/sankey.js +++ b/bin/report_utils/sankey.js @@ -50,7 +50,10 @@ const getSampleCounts = (data) => { const _getSampleCounts = (sample_name, _domain_name, agg, fragment) => { Object.entries(fragment).forEach(([key, val]) => { - const updated_domain = val.rank == "superkingdom" ? key : _domain_name; + // `domainRanks` rather than a bare "superkingdom" test: NCBI retired + // that rank in 2024 in favour of "domain" / "acellular root", so the + // old comparison matched nothing and left every tooltip's Domain null. + const updated_domain = domainRanks.has(val.rank) ? key : _domain_name; agg[key] = { [sample_name]: val.count, taxid: val.taxid, @@ -589,7 +592,7 @@ const rankSelect = renderSelect('#rank-select', ranks); rankSelect.property('value', default_rank); // Initialise rank table -const default_rank_table = "serotype"; +const default_rank_table = "species"; // Initialise cutoff const default_cutoff = 5; diff --git a/bin/scylla.mako.html b/bin/scylla.mako.html index b74bc3c4..e4237b05 100644 --- a/bin/scylla.mako.html +++ b/bin/scylla.mako.html @@ -1071,7 +1071,12 @@