From 64f5981dec73780a05672245de76fc6710b153fc Mon Sep 17 00:00:00 2001 From: Eugenio Mattei Date: Mon, 9 Nov 2026 12:07:13 -0500 Subject: [PATCH 1/2] UPDATE the scan to drop a match before it builds one scan_fasta filtered, but scan() had already built both strings of every Hit and pushed it into a per-contig Vec, so the memory still rose with every match rather than with the ones worth keeping. The PAM is now read off the raw bytes, in guide orientation, before anything is allocated. It is not faster: 1,000 9-mers against hg38 take 33.9 s against 33.3, because the rolling lookup that finds the match runs whatever the PAM says, and with a library of every 9-mer it succeeds nearly everywhere. It halves the peak, 2.91 GB to 1.45, and that is what pays: batches of 5,000 now fit in 5.5 GB, which spreads the 19-second genome pass over 53 runs rather than 263 and takes the full library from 2.5 hours to 1.3. The per-base rule moves into code_ok, which matches_pam was the only caller of, so it goes. Output is byte-identical, counters included. Co-Authored-By: Claude Opus 5 --- rust/src/lib.rs | 23 ++++------ rust/src/scan.rs | 98 +++++++++++++++++++++++++++++++---------- tests/test_core_scan.py | 41 +++++++++++++++++ 3 files changed, 124 insertions(+), 38 deletions(-) diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 80312a5..e8f7f85 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -14,7 +14,7 @@ use std::io::{BufRead, BufReader}; use pyo3::exceptions::{PyOSError, PyValueError}; use pyo3::prelude::*; -use scan::{matches_pam, pam_length, Spacers}; +use scan::{pam_length, Spacers}; /// The crate version, so Python can check which core it loaded. #[pyfunction] @@ -56,7 +56,8 @@ fn scan_fasta( contigs: Option>, pam_only: bool, ) -> PyResult<(Vec, Vec)> { - let width = pam_length(pam).map_err(PyValueError::new_err)?; + // Validates the pattern; the scan reads it base by base from here on. + pam_length(pam).map_err(PyValueError::new_err)?; let table = Spacers::new(&spacers).map_err(PyValueError::new_err)?; let wanted: Option> = contigs.map(|names| names.into_iter().collect()); let mut seen: HashSet = HashSet::new(); @@ -84,19 +85,11 @@ fn scan_fasta( if !name.is_empty() && wanted.as_ref().is_none_or(|set| set.contains(&name)) { // Releasing the GIL here is what makes a whole genome bearable // from Python: the scan touches nothing the interpreter owns. - let found = py.detach(|| scan::scan(&sequence, &table, width, upstream)); + // The scan tallies and filters as it goes, so a match with + // no PAM costs neither a string nor a row of this list. + let found = py + .detach(|| scan::scan(&sequence, &table, pam, upstream, pam_only, &mut counts)); for hit in found { - let matched = matches_pam(&hit.pam, pam); - let tally = &mut counts[hit.guide as usize]; - tally.0 += 1; - if hit.pam.len() != width { - tally.1 += 1; - } else if !matched { - tally.2 += 1; - } - if pam_only && !matched { - continue; - } hits.push(( hit.guide, name.clone(), @@ -104,7 +97,7 @@ fn scan_fasta( if hit.reverse { "-" } else { "+" }, hit.pam, hit.upstream, - matched, + hit.matched, )); } } diff --git a/rust/src/scan.rs b/rust/src/scan.rs index 0021bd0..5fd9cc1 100644 --- a/rust/src/scan.rs +++ b/rust/src/scan.rs @@ -165,6 +165,7 @@ pub struct Hit { pub reverse: bool, pub pam: String, pub upstream: String, + pub matched: bool, } /// Every exact match of `spacers` in `sequence`, with its flanks. @@ -173,8 +174,16 @@ pub struct Hit { /// it is the left end whichever strand the guide reads on. `pam` and /// `upstream` are in guide orientation, and are empty when the flank would /// run off the end of the contig. -pub fn scan(sequence: &[u8], spacers: &Spacers, pam_length: usize, upstream: usize) -> Vec { +pub fn scan( + sequence: &[u8], + spacers: &Spacers, + pattern: &str, + upstream: usize, + pam_only: bool, + counts: &mut [(u64, u64, u64)], +) -> Vec { let length = spacers.length; + let pam_length = pattern.len(); let mut hits = Vec::new(); if sequence.len() < length { return hits; @@ -198,16 +207,37 @@ pub fn scan(sequence: &[u8], spacers: &Spacers, pam_length: usize, upstream: usi }; let start = index + 1 - length; for entry in entries { + // The guide reading right to left has its PAM 5' of the match on + // the forward strand, reverse-complemented. + let (pam_from, pam_limit) = if entry.reverse { + (start.wrapping_sub(pam_length), start) + } else { + (start + length, sequence.len()) + }; + let whole = fits(sequence, pam_from, pam_length, pam_limit); + let matched = whole && pam_ok_at(sequence, pam_from, pam_limit, pattern, entry.reverse); + + let tally = &mut counts[entry.guide as usize]; + tally.0 += 1; + if !whole { + tally.1 += 1; + } else if !matched { + tally.2 += 1; + } + // Nineteen of twenty matches stop here, before either string of + // the `Hit` they would otherwise have been built into. + if pam_only && !matched { + continue; + } + let (pam, up) = if entry.reverse { - // The guide reads right to left here: its PAM is 5' of the - // match on the forward strand, reverse-complemented. ( - revcomp_slice(sequence, start.wrapping_sub(pam_length), pam_length, start), + revcomp_slice(sequence, pam_from, pam_length, pam_limit), revcomp_slice(sequence, start + length, upstream, sequence.len()), ) } else { ( - plain_slice(sequence, start + length, pam_length, sequence.len()), + plain_slice(sequence, pam_from, pam_length, pam_limit), plain_slice(sequence, start.wrapping_sub(upstream), upstream, start), ) }; @@ -217,6 +247,7 @@ pub fn scan(sequence: &[u8], spacers: &Spacers, pam_length: usize, upstream: usi reverse: entry.reverse, pam, upstream: up, + matched, }); } } @@ -228,10 +259,7 @@ fn plain_slice(sequence: &[u8], from: usize, length: usize, limit: usize) -> Str if length == 0 { return String::new(); } - if from > limit - || from.saturating_add(length) > limit - || from.saturating_add(length) > sequence.len() - { + if !fits(sequence, from, length, limit) { return String::new(); } sequence[from..from + length] @@ -283,21 +311,45 @@ pub fn pam_length(pattern: &str) -> Result { Ok(pattern.len()) } -/// Whether a stretch of genomic bases matches an IUPAC pattern. -pub fn matches_pam(bases: &str, pattern: &str) -> bool { - if bases.len() != pattern.len() { +/// Whether one base code satisfies one IUPAC symbol. +/// +/// N stands for every base, so an unsequenced one satisfies it: NGG needs +/// nothing of its first position. Any other code names bases, and an N in +/// the reference is not known to be one of them. This is the rule +/// `base_matches` states in Python; the two have to agree. +fn code_ok(c: u8, symbol: u8) -> bool { + match iupac_mask(symbol) { + None => false, + Some(ANY) => true, + Some(mask) => c != INVALID && mask >> c & 1 == 1, + } +} + +/// Whether the PAM at `from` matches `pattern`, read in guide orientation. +/// +/// The point of it is what it does not do: a match whose PAM fails is the +/// common case by nineteen to one, and this answers without building the +/// two strings a `Hit` would carry. +fn pam_ok_at(sequence: &[u8], from: usize, limit: usize, pattern: &str, reverse: bool) -> bool { + if !fits(sequence, from, pattern.len(), limit) { return false; } - bases.bytes().zip(pattern.bytes()).all(|(base, symbol)| { - let c = code(base); - match (c, iupac_mask(symbol)) { - (_, None) => false, - // N stands for every base, so an unsequenced one satisfies it: - // NGG needs nothing of its first position. Any other code names - // bases, and an N in the reference is not known to be one of them. - (_, Some(ANY)) => true, - (INVALID, _) => false, - (c, Some(mask)) => mask >> c & 1 == 1, - } + let bases = &sequence[from..from + pattern.len()]; + pattern.bytes().enumerate().all(|(i, symbol)| { + // Reading right to left and complementing, as `revcomp_slice` does. + let raw = if reverse { + bases[pattern.len() - 1 - i] + } else { + bases[i] + }; + let c = code(raw); + code_ok(if reverse && c != INVALID { 3 - c } else { c }, symbol) }) } + +/// Whether `length` bases from `from` all exist within `limit`. +fn fits(sequence: &[u8], from: usize, length: usize, limit: usize) -> bool { + from <= limit + && from.saturating_add(length) <= limit + && from.saturating_add(length) <= sequence.len() +} diff --git a/tests/test_core_scan.py b/tests/test_core_scan.py index d682402..4cd84e8 100644 --- a/tests/test_core_scan.py +++ b/tests/test_core_scan.py @@ -241,3 +241,44 @@ def test_an_unsequenced_base_satisfies_an_n_of_the_pattern(tmp_path): assert {hit[1]: hit[6] for hit in hits} == {"wild": True, "unknown": False} assert counts == [(2, 0, 1)] + + +@pytest.mark.parametrize("seed", range(6)) +def test_pam_only_agrees_with_filtering_the_full_scan(tmp_path, seed): + """The filter moved into the scan; it must still change nothing but cost. + + The oracle checks only the unfiltered path, so this is what would catch + the tally drifting once it is kept somewhere else. + """ + rng = random.Random(seed) + spacers = [random_sequence(rng, 8) for _ in range(6)] + contigs = {f"chr{n}": random_sequence(rng, 400, "ACGTN") for n in range(3)} + path = tmp_path / "genome.fa" + write_fasta(path, contigs) + + full, counts = _core.scan_fasta(str(path), spacers, "NGG", 1) + only, same = _core.scan_fasta(str(path), spacers, "NGG", 1, None, True) + + assert counts == same == tally(oracle(contigs, spacers, "NGG", 1), spacers, "NGG") + assert sorted(only) == sorted(hit for hit in full if hit[6]) + + +def test_a_minus_strand_match_against_the_very_start_of_a_contig(tmp_path): + # Its PAM lies 5' of base 0, which does not exist: the bounds check has + # to catch `start - 3` wrapping round rather than read somewhere else. + spacer = "ACGTTGCAAGCTTCGATCGA" + contigs = {"chrEdge": reverse_complement_py(spacer) + "TTTTT"} + path = tmp_path / "genome.fa" + write_fasta(path, contigs) + + hits, counts = _core.scan_fasta(str(path), [spacer], "NGG", 1) + + assert [hit[3] for hit in hits] == ["-"] + assert hits[0][4] == "" + assert counts == [(1, 1, 0)] + only, same = _core.scan_fasta(str(path), [spacer], "NGG", 1, None, True) + assert only == [] and same == counts + + +def reverse_complement_py(sequence): + return sequence[::-1].translate(str.maketrans("ACGT", "TGCA")) From f55adf105827282c5f8181be3bf4b50e89686907 Mon Sep 17 00:00:00 2001 From: Eugenio Mattei Date: Mon, 9 Nov 2026 12:07:31 -0500 Subject: [PATCH 2/2] FIX the gotcha, which named a function the last commit removed Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 6e3d5bb..cc69e69 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -337,10 +337,12 @@ and narrowing it would buy nothing. - **`N` in a PAM pattern is satisfied by a genomic `N`.** `NGG` asks nothing of its first position, so an unsequenced base there is still a PAM; every other code names bases, and an `N` is not known to be one of - them. The rule lives in `base_matches` and in the Rust `matches_pam`, + them. The rule lives in `base_matches` and in the Rust `code_ok`, which disagreed with it until they were reconciled, and in the oracle of `tests/test_core_scan.py`, which had encoded the Rust reading. - Change one and the other two are wrong; + Change one and the other two are wrong. `code_ok` is read twice per + match, once by `pam_ok_at` on the raw bytes before anything is + allocated, so keep it branch-light; - classes count sites whose genomic PAM matches. A site whose PAM does not stays in the table with `pam_matches` `FALSE` and counts nowhere, so a spacer matching twice with one PAM is `unique`. `--pam-only`