ADD pam_only to the compiled core, so a dropped match never reaches Python

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,25 +25,37 @@ 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,
     spacers: Vec<String>,
     pam: &str,
     upstream: usize,
     contigs: Option<Vec<String>>,
-) -> PyResult<Vec<PyHit>> {
+    pam_only: bool,
+) -> PyResult<(Vec<PyHit>, Vec<PyCounts>)> {
     let width = pam_length(pam).map_err(PyValueError::new_err)?;
     let table = Spacers::new(&spacers).map_err(PyValueError::new_err)?;
     let wanted: Option<HashSet<String>> = 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<PyCounts> = vec![(0, 0, 0); spacers.len()];
     let mut name = String::new();
     let mut sequence: Vec<u8> = 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]

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<u8> {
         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,
         }
     })

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)

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

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