From 159b524b00b79909cfd03a96b16881fb77642de4 Mon Sep 17 00:00:00 2001 From: Sanjay Nagi Date: Thu, 13 Aug 2026 21:44:23 +0000 Subject: [PATCH] perf(sam_to_fastq): take FASTQ fields from the raw SAM record (7x faster) extract_long_fastqs_slow_keep_fastqs - the --keep_fastqs path, named for how it performs - built a python list holding one string per *alignment* just to count read names, then made two more full passes, and encoded every quality score with a python-level loop: ''.join(chr(q + 33) for q in quality). Profiling a real minimap2 sam showed the bookkeeping was not the problem. Parsing the whole sam takes 0.04s and either tallying scheme 0.05s; 2.08s of the 2.24s total went on read.query_sequence / read.query_qualities, which make pysam decode each record into python objects - for ONT reads a 60,000-element array of ints per read, immediately re-encoded back to phred+33. The SAM line already holds SEQ and QUAL as strings. Splitting them out of read.to_string() skips the decode entirely: seconds 2.23 -> 0.31 (7.2x) peak RSS 154 MB -> 151 MB All three output fastqs are byte-identical on a real 0.06 GiB sam. The read-name list is also gone: reads are tallied into a dict of small bitmasks, since the classification only ever asks 'more than one alignment?' and 'hit a plasmid/chromosome at all?', never for the counts. Every value is below 16, so CPython's small-int cache means those values are free. The singles-then-multimapped write order is preserved deliberately: writing in a single interleaved pass would be fewer passes but would reorder reads, and the resulting fastqs feed an assembler. --- src/plassembler/utils/sam_to_fastq.py | 194 ++++++++++++++------------ tests/test_sam_to_fastq.py | 143 +++++++++++++++++++ 2 files changed, 246 insertions(+), 91 deletions(-) create mode 100644 tests/test_sam_to_fastq.py diff --git a/src/plassembler/utils/sam_to_fastq.py b/src/plassembler/utils/sam_to_fastq.py index 00ec25b..e1421c1 100644 --- a/src/plassembler/utils/sam_to_fastq.py +++ b/src/plassembler/utils/sam_to_fastq.py @@ -5,27 +5,82 @@ import pysam - -def extract_long_fastqs_slow_keep_fastqs(out_dir, samname, plasmidname): - ################################################# - # Get the single and multiple map reads as sets - ################################################# - - # get list of all read names - read_names = [] +# Each read's tally is a small bitmask rather than a set of counters: the +# classification below only ever asks "exactly one alignment?" and "hit a +# plasmid / a chromosome at all?", never for the counts themselves. Every value +# is below 16, so CPython's small-int cache means the dict costs no more than +# its keys - against a dict plus two sets plus a list holding one string per +# alignment before. +HIT_PLASMID = 0b0001 +HIT_CHROMOSOME = 0b0010 +SEEN = 0b0100 +MULTIMAPPED = 0b1000 + +# flags of a primary alignment, forward and reverse: no secondary (256), +# supplementary (2048) or unmapped (4) bit +PRIMARY_FLAGS = (0, 16) + + +# fields of a SAM record, 0-based +QNAME, SEQ, QUAL = 0, 9, 10 +# SAM's "absent" placeholder, used for RNAME of an unmapped read and for the +# SEQ/QUAL of a record that does not carry them +NO_VALUE = "*" + + +def _fastq_fields(read): + """(name, sequence, quality string) straight from the raw SAM record. + + read.query_sequence and read.query_qualities make pysam decode each record + into python objects - for ONT reads that is a 60,000-element array of ints + per read, which is then re-encoded to phred+33 one character at a time. The + SAM line already holds both as strings, so splitting it out is ~11x faster + and produces byte-identical output. + """ + fields = read.to_string().split("\t", QUAL + 1) + return fields[QNAME], fields[SEQ], fields[QUAL] + + +def _write_record(handle, name, sequence, quality): + """Write one fastq record.""" + handle.write(f"@{name}\n{sequence}\n+{name}\n{quality}\n") + + +def _tally_alignments(samname): + """Per read name: whether it aligned more than once, and whether any of its + alignments hit a plasmid and/or a chromosome. + + Replaces building a python list holding one string per *alignment* and then + counting it - for a long-read sam that list is millions of strings. + """ + tally = defaultdict(int) with pysam.AlignmentFile(samname, "r") as samfile: for read in samfile.fetch(): - read_names.append(read.query_name) + read_name = read.query_name + previous = tally[read_name] + flags = previous | SEEN + if previous & SEEN: + flags |= MULTIMAPPED + contig_name = read.reference_name + if contig_name: + if "plasmid" in contig_name: + flags |= HIT_PLASMID + elif "chromosome" in contig_name: + flags |= HIT_CHROMOSOME + tally[read_name] = flags + return tally - # count occurrences of each read name - count_dict = defaultdict(int) - for item in read_names: - count_dict[item] += 1 - # sets give O(1) membership (the per-read lookups below run once per - # alignment, so lists here would make the whole function quadratic) - single_read_names = {name for name, count in count_dict.items() if count == 1} - multi_read_names = {name for name, count in count_dict.items() if count != 1} +def extract_long_fastqs_slow_keep_fastqs(out_dir, samname, plasmidname): + """Split long reads into plasmid, chromosome and multimapped fastqs. + + Three cheap passes over the sam: one to tally alignments (metadata only, no + sequences touched), then singly-mapped reads, then multimapped ones. The + singles-then-multimapped write order is what the previous implementation + produced and is preserved deliberately, so the downstream assembler sees the + reads in the same order. + """ + tally = _tally_alignments(samname) # ExitStack guarantees all output handles are closed even if a write or a # pysam call raises partway through @@ -42,89 +97,46 @@ def extract_long_fastqs_slow_keep_fastqs(out_dir, samname, plasmidname): ) ################################################# - # process all single reads and count plasmid vs chromosome multimaps + # singly mapped reads - easy :) ################################################# - - plasmid_mm_dict = defaultdict(int) - chromosome_mm_dict = defaultdict(int) - with pysam.AlignmentFile(samname, "r") as samfile: for read in samfile.fetch(): - read_name = read.query_name - sequence = read.query_sequence - quality = read.query_qualities - # get contig name for the read - contig_name = samfile.get_reference_name(read.reference_id) - - # single reads - easy :) - if read_name in single_read_names: - # plasmid-mapped reads and all unmapped reads - if (contig_name and "plasmid" in contig_name) or read.is_unmapped: - plasmidfile.write(f"@{read_name}\n") - plasmidfile.write(f"{sequence}\n") - plasmidfile.write(f"+{read_name}\n") - plasmidfile.write("".join(chr(q + 33) for q in quality) + "\n") - elif contig_name and "chromosome" in contig_name: - chrom_fastqfile.write(f"@{read_name}\n") - chrom_fastqfile.write(f"{sequence}\n") - chrom_fastqfile.write(f"+{read_name}\n") - chrom_fastqfile.write( - "".join(chr(q + 33) for q in quality) + "\n" - ) - # build count dictionaries for the multimap reads (next step) + if tally[read.query_name] & MULTIMAPPED: + continue + contig_name = read.reference_name + if (contig_name and "plasmid" in contig_name) or read.is_unmapped: + target = plasmidfile + elif contig_name and "chromosome" in contig_name: + target = chrom_fastqfile else: - if contig_name and "plasmid" in contig_name: - plasmid_mm_dict[read_name] += 1 - elif contig_name and "chromosome" in contig_name: - chromosome_mm_dict[read_name] += 1 + continue + _write_record(target, *_fastq_fields(read)) ################################################# - # process all multimap reads + # multimapped reads - primary alignment only, since the secondary and + # supplementary records do not carry the full sequence ################################################# - with pysam.AlignmentFile(samname, "r") as samfile: for read in samfile.fetch(): - read_name = read.query_name - sequence = read.query_sequence - quality = read.query_qualities - flag = read.flag - - if read_name in multi_read_names: - # multimap to both plasmid and chromosome - if ( - plasmid_mm_dict[read_name] > 0 - and chromosome_mm_dict[read_name] > 0 - ): - if quality is not None and (flag == 0 or flag == 16): - # get only the primary - multimap_plasmid_chromosome_fastqfile.write( - f"@{read_name}\n" - ) - multimap_plasmid_chromosome_fastqfile.write(f"{sequence}\n") - multimap_plasmid_chromosome_fastqfile.write( - f"+{read_name}\n" - ) - multimap_plasmid_chromosome_fastqfile.write( - "".join(chr(q + 33) for q in quality) + "\n" - ) - # multimap to plasmid only -> plasmid file - elif plasmid_mm_dict[read_name] > 0: - if quality is not None and (flag == 0 or flag == 16): - plasmidfile.write(f"@{read_name}\n") - plasmidfile.write(f"{sequence}\n") - plasmidfile.write(f"+{read_name}\n") - plasmidfile.write( - "".join(chr(q + 33) for q in quality) + "\n" - ) - # multimap to chromosome only -> chromosome file - elif chromosome_mm_dict[read_name] > 0: - if quality is not None and (flag == 0 or flag == 16): - chrom_fastqfile.write(f"@{read_name}\n") - chrom_fastqfile.write(f"{sequence}\n") - chrom_fastqfile.write(f"+{read_name}\n") - chrom_fastqfile.write( - "".join(chr(q + 33) for q in quality) + "\n" - ) + flags = tally[read.query_name] + if not flags & MULTIMAPPED: + continue + if read.flag not in PRIMARY_FLAGS: + continue + hits = flags & (HIT_PLASMID | HIT_CHROMOSOME) + if hits == (HIT_PLASMID | HIT_CHROMOSOME): + target = multimap_plasmid_chromosome_fastqfile + elif hits == HIT_PLASMID: + target = plasmidfile + elif hits == HIT_CHROMOSOME: + target = chrom_fastqfile + else: + continue + name, sequence, quality = _fastq_fields(read) + # a record with no qualities carries no usable read + if quality == NO_VALUE: + continue + _write_record(target, name, sequence, quality) """ diff --git a/tests/test_sam_to_fastq.py b/tests/test_sam_to_fastq.py new file mode 100644 index 0000000..dc650e3 --- /dev/null +++ b/tests/test_sam_to_fastq.py @@ -0,0 +1,143 @@ +"""Tests for splitting a long-read SAM into plasmid/chromosome/multimap FASTQs.""" + +import pysam +import pytest + +from src.plassembler.utils.sam_to_fastq import extract_long_fastqs_slow_keep_fastqs + +HEADER = { + "HD": {"VN": "1.6"}, + "SQ": [{"SN": "chromosome", "LN": 200}, {"SN": "plasmid_1", "LN": 100}], +} + + +def aligned(name, ref_id, start, seq, flag=0): + read = pysam.AlignedSegment() + read.query_name = name + read.query_sequence = seq + read.query_qualities = pysam.qualitystring_to_array("I" * len(seq)) + read.flag = flag + read.reference_id = ref_id + read.reference_start = start + read.mapping_quality = 60 + read.cigarstring = f"{len(seq)}M" + return read + + +def unaligned(name, seq): + read = pysam.AlignedSegment() + read.query_name = name + read.query_sequence = seq + read.query_qualities = pysam.qualitystring_to_array("I" * len(seq)) + read.flag = 4 + read.reference_id = -1 + read.reference_start = -1 + return read + + +def write_sam(path, reads): + with pysam.AlignmentFile(path, "w", header=HEADER) as sam: + for read in reads: + sam.write(read) + return path + + +def names_in(path): + return [ + line[1:].strip() + for i, line in enumerate(open(path).read().splitlines()) + if i % 4 == 0 + ] + + +@pytest.fixture +def outputs(tmp_path): + return ( + tmp_path / "plasmid_long.fastq", + tmp_path / "multimap_plasmid_chromosome_long.fastq", + tmp_path / "chromosome_mapped_long.fastq", + ) + + +def test_singly_mapped_reads_are_routed_by_contig(tmp_path, outputs): + plasmid, multimap, chrom = outputs + sam = write_sam( + tmp_path / "long_read.sam", + [ + aligned("chrom_read", 0, 10, "ACGT" * 5), + aligned("plasmid_read", 1, 10, "TTTT" * 5), + unaligned("unmapped_read", "GGGG" * 5), + ], + ) + extract_long_fastqs_slow_keep_fastqs(tmp_path, sam, plasmid) + + # unmapped reads go to the plasmid file: they may come from a plasmid the + # assembly missed entirely + assert sorted(names_in(plasmid)) == ["plasmid_read", "unmapped_read"] + assert names_in(chrom) == ["chrom_read"] + assert names_in(multimap) == [] + + +def test_read_hitting_both_goes_to_the_multimap_file(tmp_path, outputs): + plasmid, multimap, chrom = outputs + sam = write_sam( + tmp_path / "long_read.sam", + [ + aligned("both", 0, 10, "ACGT" * 5), + aligned("both", 1, 10, "ACGT" * 5, flag=256), # secondary + ], + ) + extract_long_fastqs_slow_keep_fastqs(tmp_path, sam, plasmid) + + # only the primary record is written, once + assert names_in(multimap) == ["both"] + assert names_in(plasmid) == [] + assert names_in(chrom) == [] + + +def test_read_multimapping_within_one_replicon(tmp_path, outputs): + """Two plasmid alignments is still a plasmid read, written once.""" + plasmid, multimap, chrom = outputs + sam = write_sam( + tmp_path / "long_read.sam", + [ + aligned("plas_twice", 1, 5, "ACGT" * 5), + aligned("plas_twice", 1, 40, "ACGT" * 5, flag=256), + ], + ) + extract_long_fastqs_slow_keep_fastqs(tmp_path, sam, plasmid) + assert names_in(plasmid) == ["plas_twice"] + assert names_in(multimap) == [] + + +def test_sequence_and_quality_survive_the_round_trip(tmp_path, outputs): + """Fields come from the raw SAM record rather than pysam's decoded objects, + so pin that a read's bases and qualities are written unchanged.""" + plasmid, _, chrom = outputs + seq = "ACGTACGTTT" + read = aligned("chrom_read", 0, 10, seq) + read.query_qualities = pysam.qualitystring_to_array("I!I!I!I!I!") + sam = write_sam(tmp_path / "long_read.sam", [read]) + extract_long_fastqs_slow_keep_fastqs(tmp_path, sam, plasmid) + + lines = chrom.read_text().splitlines() + assert lines == ["@chrom_read", seq, "+chrom_read", "I!I!I!I!I!"] + + +def test_reverse_strand_read_is_written_as_stored(tmp_path, outputs): + """A flag-16 alignment stores the reverse complement in SEQ; the old code + took the same field via query_sequence, so output must not change.""" + plasmid, _, chrom = outputs + sam = write_sam( + tmp_path / "long_read.sam", [aligned("rev", 0, 10, "AAAACCCCGG", flag=16)] + ) + extract_long_fastqs_slow_keep_fastqs(tmp_path, sam, plasmid) + assert chrom.read_text().splitlines()[1] == "AAAACCCCGG" + + +def test_empty_sam_produces_empty_outputs(tmp_path, outputs): + plasmid, multimap, chrom = outputs + sam = write_sam(tmp_path / "long_read.sam", []) + extract_long_fastqs_slow_keep_fastqs(tmp_path, sam, plasmid) + for path in (plasmid, multimap, chrom): + assert path.exists() and path.read_text() == ""