From f3bd2220bb82c880e363745dcb941fb817a836fc Mon Sep 17 00:00:00 2001 From: Eugenio Mattei Date: Mon, 9 Nov 2026 11:27:58 -0500 Subject: [PATCH] ADD pam_only to the core, so a dropped match never reaches Python The Rust scan pushed every hit into a Vec that crosses into Python as a list, so the peak was set at that boundary whatever Python did with it afterwards. It now skips a hit with no PAM, and returns a per-spacer tally of what it found, what ran off a contig end and what had no PAM, so the counters still say what the run saw rather than what it kept. 500 9-mers against chr21: peak 1,056 MB to 538 MB, with the tables byte-identical to the Python-side filter. matches_pam agreed with none of this: its `code` has no value for a genomic N, so it rejected one even at an N of the pattern, where base_matches accepts it by design and the pam_matches column always has. Filtering on the core's own verdict would have dropped rows the table calls TRUE. Both rules now read the same, and the oracle, which had encoded the core's, reads it too. Co-Authored-By: Claude Opus 5 --- rust/src/lib.rs | 46 +++++++++++--- rust/src/scan.rs | 33 ++++++++-- .../guide_alignment/scan.py | 48 +++++++-------- .../skills/cau-guide-alignment/SKILL.md | 5 +- tests/test_core_scan.py | 61 +++++++++++++++++-- 5 files changed, 147 insertions(+), 46 deletions(-) diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 7ff4743..80312a5 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::{Spacers, matches_pam, pam_length}; +use scan::{matches_pam, pam_length, Spacers}; /// The crate version, so Python can check which core it loaded. #[pyfunction] @@ -25,17 +25,28 @@ fn version() -> &'static str { /// One exact match, as Python receives it. type PyHit = (u32, String, usize, &'static str, String, String, bool); +/// One spacer's tally: matches found, how many ran off a contig end, and how +/// many of the rest had no PAM. Kept whether or not the matches are returned. +type PyCounts = (u64, u64, u64); + /// Every exact match of `spacers` in a reference FASTA. /// -/// Returns one tuple per match: the spacer's index in `spacers`, the contig, -/// the 0-based start of the protospacer on the forward strand, the strand the -/// guide reads on, the genomic PAM in guide orientation, the `upstream` bases -/// 5' of the spacer in guide orientation, and whether the PAM matches. +/// Returns the matches and, beside them, one `(found, off_end, no_pam)` per +/// spacer. Each match is a tuple: the spacer's index in `spacers`, the +/// contig, the 0-based start of the protospacer on the forward strand, the +/// strand the guide reads on, the genomic PAM in guide orientation, the +/// `upstream` bases 5' of the spacer in guide orientation, and whether the +/// PAM matches. /// /// A flank that would run off the end of its contig comes back empty, and a /// PAM that does is not a match. +/// +/// `pam_only` returns only the matches whose PAM matches. A short spacer +/// matches far more places than a nuclease could act at, and those are what +/// fills memory before Python ever sees them; the counts keep them, so +/// nothing the caller reports has to change. #[pyfunction] -#[pyo3(signature = (fasta, spacers, pam = "NGG", upstream = 1, contigs = None))] +#[pyo3(signature = (fasta, spacers, pam = "NGG", upstream = 1, contigs = None, pam_only = false))] fn scan_fasta( py: Python<'_>, fasta: &str, @@ -43,7 +54,8 @@ fn scan_fasta( pam: &str, upstream: usize, contigs: Option>, -) -> PyResult> { + pam_only: bool, +) -> PyResult<(Vec, Vec)> { let width = 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()); @@ -54,6 +66,7 @@ fn scan_fasta( let mut reader = BufReader::with_capacity(1 << 20, file); let mut hits = Vec::new(); + let mut counts: Vec = vec![(0, 0, 0); spacers.len()]; let mut name = String::new(); let mut sequence: Vec = Vec::new(); let mut line = Vec::new(); @@ -74,6 +87,16 @@ fn scan_fasta( let found = py.detach(|| scan::scan(&sequence, &table, width, upstream)); 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(), @@ -90,8 +113,11 @@ fn scan_fasta( // an empty result: without this the whole run would quietly // report that no guide maps anywhere. if let Some(set) = wanted.as_ref() { - let mut missing: Vec<&str> = - set.iter().filter(|n| !seen.contains(*n)).map(|n| n.as_str()).collect(); + let mut missing: Vec<&str> = set + .iter() + .filter(|n| !seen.contains(*n)) + .map(|n| n.as_str()) + .collect(); if !missing.is_empty() { missing.sort_unstable(); return Err(PyValueError::new_err(format!( @@ -116,7 +142,7 @@ fn scan_fasta( } sequence.extend(line.iter().filter(|b| !b.is_ascii_whitespace())); } - Ok(hits) + Ok((hits, counts)) } #[pymodule] diff --git a/rust/src/scan.rs b/rust/src/scan.rs index 43c223b..0021bd0 100644 --- a/rust/src/scan.rs +++ b/rust/src/scan.rs @@ -10,6 +10,9 @@ use std::collections::HashMap; /// Base codes: 0-3 for ACGT, 4 for anything else. const INVALID: u8 = 4; +/// The IUPAC mask of `N`: every base, and so every character. +const ANY: u8 = 0b1111; + fn code(base: u8) -> u8 { match base { b'A' | b'a' => 0, @@ -41,7 +44,7 @@ fn iupac_mask(symbol: u8) -> Option { b'D' => 0b1101, b'H' => 0b1011, b'V' => 0b0111, - b'N' => 0b1111, + b'N' => ANY, _ => return None, }) } @@ -105,11 +108,23 @@ impl Spacers { reverse |= ((3 - c) as u64) << (2 * offset); } let guide = index as u32; - found.insert(forward, Entry { guide, reverse: false }); + found.insert( + forward, + Entry { + guide, + reverse: false, + }, + ); // A palindromic spacer keeps both entries on purpose: it really does // read 5' to 3' on either strand, and the two readings have // different flanks, so one can have a PAM where the other does not. - found.insert(reverse, Entry { guide, reverse: true }); + found.insert( + reverse, + Entry { + guide, + reverse: true, + }, + ); } Ok(found) } @@ -213,7 +228,10 @@ 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 from > limit + || from.saturating_add(length) > limit + || from.saturating_add(length) > sequence.len() + { return String::new(); } sequence[from..from + length] @@ -273,7 +291,12 @@ pub fn matches_pam(bases: &str, pattern: &str) -> bool { bases.bytes().zip(pattern.bytes()).all(|(base, symbol)| { let c = code(base); match (c, iupac_mask(symbol)) { - (INVALID, _) | (_, None) => false, + (_, 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, } }) diff --git a/src/crispr_analysis_utils/guide_alignment/scan.py b/src/crispr_analysis_utils/guide_alignment/scan.py index b17fdb3..5159650 100644 --- a/src/crispr_analysis_utils/guide_alignment/scan.py +++ b/src/crispr_analysis_utils/guide_alignment/scan.py @@ -112,10 +112,12 @@ def scan_guides( contigs If given, read only these contigs. pam_only - Keep only the matches whose genomic PAM matches `pam`. Every match - is still counted, in `ScanResult.matches` and `n_matches`; what - changes is which of them become a `Site`. A short spacer matches - everywhere, and the sites it cannot be cut at are the bulk of them. + Keep only the matches whose genomic PAM matches `pam`. The core + drops the rest as it finds them, so they never cross into Python + and nothing holds them. Every match is still counted, in + `ScanResult.matches` and `n_matches`: a short spacer matches + everywhere, and the places it cannot be cut at are the bulk of + them, but how many there were is worth knowing. Returns ------- @@ -136,7 +138,9 @@ def scan_guides( search between them. `pam_only` shrinks what is held and what is written, not what is - searched: every match is found and counted either way. + searched: every match is found and counted either way. The counters + come from the scan rather than from the matches returned, which is what + lets them report what they never handed back. """ from .. import _core @@ -167,39 +171,33 @@ def scan_guides( queries = spacers + reversed_spacers upstream_width = max((guide.n_extra for guide in guides), default=0) listed = None if contigs is None else list(contigs) - matches = _core.scan_fasta( - os.fspath(reference), queries, primary, upstream_width, listed + matches, counts = _core.scan_fasta( + os.fspath(reference), queries, primary, upstream_width, listed, pam_only ) result = ScanResult() result.matches = dict.fromkeys(range(len(guides)), 0) + # The counts come from the scan itself rather than from the matches, so + # they say the same thing whether or not `pam_only` returned them all. + # The reversed readings are not sites and are not counted here. + for index, (found, off_end, no_pam) in enumerate(counts[: len(spacers)]): + result.n_matches += found + result.n_off_contig_end += off_end + result.n_no_pam += no_pam + result.n_sites += found - off_end + for row in rows_of[spacers[index]]: + result.matches[row] += found - off_end + reversed_has_pam: set[int] = set() for index, chrom, start, strand, genomic_pam, upstream, matched in matches: if index >= len(spacers): - # The core's verdict is enough here: this is only evidence that - # the reversed reading could cut somewhere, and no column of - # any table reports it. + # Only evidence that the reversed reading could cut somewhere. if matched: reversed_has_pam.update(flipped_rows[queries[index]]) continue - result.n_matches += 1 if len(genomic_pam) != len(primary): - result.n_off_contig_end += 1 continue - # The core's own `matched` is stricter than this over a genomic N, - # which `base_matches` reads as a PAM at an N of the pattern and - # the core reads as no PAM at all. The `pam_matches` column has - # always followed this one, so the counter and the flag do too: - # otherwise qc.tsv contradicts the table beside it. - has_pam = pattern_matches(genomic_pam, primary) - if not has_pam: - result.n_no_pam += 1 - result.n_sites += 1 - keep = has_pam or not pam_only for row in rows_of[spacers[index]]: - result.matches[row] += 1 - if not keep: - continue guide = guides[row] result.sites.setdefault(row, []).append( _site(guide, chrom, start, strand, genomic_pam, upstream, primary) diff --git a/src/crispr_analysis_utils/skills/cau-guide-alignment/SKILL.md b/src/crispr_analysis_utils/skills/cau-guide-alignment/SKILL.md index 3479f95..2dd5d45 100644 --- a/src/crispr_analysis_utils/skills/cau-guide-alignment/SKILL.md +++ b/src/crispr_analysis_utils/skills/cau-guide-alignment/SKILL.md @@ -98,7 +98,10 @@ Written out, that table is mostly rows nothing can be done with. `--pam-only` keeps the sites whose genomic PAM matches and drops the rest. -- **It changes what is written, not what is searched or counted.** +- **It changes what is written and what is held, not what is searched or + counted.** The compiled core drops a match with no PAM as it finds it, so + the peak memory of a run falls with the table rather than only after it. +- **The counts are unmoved.** `n_matches` still counts every place the protospacer occurs, which under this flag is the only record that they exist, and `qc.tsv` still reports `n_sites` and `n_no_pam`, so two runs stay comparable. The rows written diff --git a/tests/test_core_scan.py b/tests/test_core_scan.py index 3d04bda..d682402 100644 --- a/tests/test_core_scan.py +++ b/tests/test_core_scan.py @@ -59,8 +59,11 @@ def oracle(contigs, spacers, pam, upstream): up = "" if len(pam_bases) != len(pam): pam_bases = "" + # N stands for every base, so it is satisfied by an + # unsequenced one too; any other code names bases, and a + # genomic N is not known to be one of them. matched = len(pam_bases) == len(pam) and all( - base in IUPAC[code] + code == "N" or base in IUPAC[code] for base, code in zip(pam_bases, pam, strict=True) ) found.append((index, name, start, strand, pam_bases, up, matched)) @@ -78,9 +81,23 @@ def write_fasta(path, contigs, width=60): def run(tmp_path, contigs, spacers, pam="NGG", upstream=1): path = tmp_path / "genome.fa" write_fasta(path, contigs) - got = sorted(_core.scan_fasta(str(path), spacers, pam, upstream)) - assert got == oracle(contigs, spacers, pam, upstream) - return got + hits, counts = _core.scan_fasta(str(path), spacers, pam, upstream) + expected = oracle(contigs, spacers, pam, upstream) + assert sorted(hits) == expected + assert counts == tally(expected, spacers, pam) + return sorted(hits) + + +def tally(hits, spacers, pam): + """What the scan should report per spacer: found, off an end, no PAM.""" + counts = [[0, 0, 0] for _ in spacers] + for index, _, _, _, pam_bases, _, matched in hits: + counts[index][0] += 1 + if len(pam_bases) != len(pam): + counts[index][1] += 1 + elif not matched: + counts[index][2] += 1 + return [tuple(row) for row in counts] def random_sequence(rng, length, alphabet="ACGT"): @@ -188,5 +205,39 @@ def test_contig_filter(tmp_path): contigs = {"keep": "TT" + spacer + "TGG", "drop": "AA" + spacer + "AGG"} path = tmp_path / "genome.fa" write_fasta(path, contigs) - got = _core.scan_fasta(str(path), [spacer], "NGG", 1, ["keep"]) + got, _ = _core.scan_fasta(str(path), [spacer], "NGG", 1, ["keep"]) assert {hit[1] for hit in got} == {"keep"} + + +def test_pam_only_returns_the_matches_with_a_pam_and_counts_the_rest(tmp_path): + spacer = "ACGTTGCAAGCTTCGATCGA" + contigs = { + "cuts": "TT" + spacer + "TGG", + "blunt": "TT" + spacer + "TTT", + "also_blunt": "TT" + spacer + "CAT", + } + path = tmp_path / "genome.fa" + write_fasta(path, contigs) + + everything, counts = _core.scan_fasta(str(path), [spacer], "NGG", 1) + only, same = _core.scan_fasta(str(path), [spacer], "NGG", 1, None, True) + + # The counts are the scan's, not the returned matches', so they do not + # move: that is what lets the caller report what it never received. + assert counts == same == [(3, 0, 2)] + assert len(everything) == 3 + assert only == [hit for hit in everything if hit[6]] + + +def test_an_unsequenced_base_satisfies_an_n_of_the_pattern(tmp_path): + # NGG asks nothing of its first position, so an N there is still a PAM. + # Any other code names bases, and an N is not known to be one of them. + spacer = "ACGTTGCAAGCTTCGATCGA" + contigs = {"wild": "TT" + spacer + "NGG", "unknown": "TT" + spacer + "GNG"} + path = tmp_path / "genome.fa" + write_fasta(path, contigs) + + hits, counts = _core.scan_fasta(str(path), [spacer], "NGG", 1) + + assert {hit[1]: hit[6] for hit in hits} == {"wild": True, "unknown": False} + assert counts == [(2, 0, 1)]