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)]