Skip to content

refactor(rust): take the signal-level k-mer path from escapepod-signal - #222

Merged
jayhesselberth merged 2 commits into
mainfrom
feat/upstream-kmer-encoding
Aug 26, 2026
Merged

refactor(rust): take the signal-level k-mer path from escapepod-signal#222
jayhesselberth merged 2 commits into
mainfrom
feat/upstream-kmer-encoding

Conversation

@jayhesselberth

@jayhesselberth jayhesselberth commented Aug 25, 2026

Copy link
Copy Markdown
Member

Closes the loop opened in #220 and escapepod-rs#271: the signal-level k-mer
encoding now has exactly one definition, upstream.

Why it had to move

leech held the only copy, inside a crate-type = ["cdylib"] Python extension
module — so a native runtime for a leech signal_kmer model could not link it
and had to transcribe it. Two definitions of this rule diverge silently, and
this stack has the scars: extract_levels was written twice with different
centring conventions and moved 25 of 100 features; escapepod-classify
reproduced a superseded feature definition for two months, its golden missing
it because all 19 fixture reads took the other branch.

It is also the natural pair to escapepod_signal::mapping, which produces
the base-to-signal map this consumes. The producing half was already
upstream; the consuming half was not. escapepod-rs#272 fixed that — this is
leech taking the call.

What delegates

leech escapepod-signal v0.16.1
encoding.rs::encode_signal_kmer_inner seq_encoding::encode_signal_kmer
features.rs::sequence_to_int seq_encoding::sequence_to_int
signal_mapping.rs context slice seq_encoding::sequence_bases_with_context

Same A/C/G/T-U table, same -1 for anything else (upstream's UNKNOWN_BASE
is -1 — checked, not assumed).

The third one too, after an upstream change (escapepod-signal 0.16.1)

The k-mer context slice in signal_mapping.rs originally could not
delegate: leech needs the window as bases, since it serializes
sequence_with_kmer_context as a string, where upstream only offered ints.
That was filed as escapepod-rs#274 and released in 0.16.1, which now exposes
both forms over one windowing rule with sequence_to_int(bases) == ints pinned
by a test there. leech takes sequence_bases_with_context.

This is the one of the three most worth de-duplicating. It is where before
and after are not interchangeable: swap them and every k-mer is read from a
window displaced by before - after bases, silently, and encode_signal_kmer
cannot detect it because it only sees the total width.

Identical by construction and by test — the ranges agree
((seq_start - before)..(seq_end + after) is upstream's window with
n_bases = seq_end - seq_start), and both pad rather than shift, leech with
b'N' and upstream with UNKNOWN_BASE_CHAR, which is b'N'.

Evidence

A delegation's only acceptable outcome is identity, and the evidence is that
nothing had to change to accommodate it:

  • 198 parity tests pass untouched, including test_backend_parity.py, which
    compares every array in the npz between the Rust and Python backends. The
    Python reference is not touched here, so it is an independent implementation.
  • Full suite 1467 passed — run twice, once against Python escapepod 0.15.0
    and once against 0.16.0, since the Rust crate and the Python wheel are
    independent (leech_core links the crate statically; leech.io imports the
    wheel) and the pin moves to >=0.16.0.
  • cargo fmt --check, cargo clippy -D warnings, ruff, ty all clean.

Also here, because the pin moves anyway

  • uv.lock gains onnx/onnxruntime/onnxscript/protobuf. feat(export): ONNX for the classifier arms and the CRF encoder #220 added the
    onnx extra to pyproject.toml and never locked it — CI installs from
    pyproject so it worked, but uv sync --extra onnx would not have. Checked for
    the usual emscripten churn: 29 added marker lines, all on the new packages,
    plus one resolution-marker refinement. Not a rewrite.
  • The ONNX contract stops telling consumers to reimplement the encoder and names
    the crate that owns it.
  • CLAUDE.md: "four primitives now come from upstream" becomes five.

🤖 Generated with Claude Code


Update: all three overlaps now delegate

mapping produces the base-to-signal map, sequence_bases_with_context cuts
the window, encode_signal_kmer scatters it — the whole signal-level k-mer
path is upstream, none of it duplicated here.

Only rust/Cargo.toml's git tag moves to v0.16.1 (Cargo.lock pins commit
b4c9afae). The escapepod Python pin stays >=0.16.0: the new function is
in the Rust crate, not the bindings, and the two are independent — leech_core
links the crate statically while leech.io imports the wheel.

Re-verified with all three in place: 198 parity tests untouched (including
the array-by-array sequence_with_kmer_context comparison), full suite 1467
passed
, cargo fmt --check and clippy -D warnings clean.

…ignal

leech held the only copy of this rule, inside a `crate-type = ["cdylib"]`
Python extension module — so a native runtime for a leech `signal_kmer` model
could not link it and had to transcribe it, which is a second definition that
diverges silently. That is not hypothetical in this stack: `extract_levels` was
written twice with different centring conventions and moved 25 of 100 features,
and escapepod-classify reproduced a superseded feature definition for two
months, its golden missing it because all 19 fixture reads took the other
branch.

It is also the natural pair to `escapepod_signal::mapping`, which *produces*
the base-to-signal map the encoding consumes. The producing half was already
upstream and the consuming half was not. escapepod-rs#272 fixed that; this is
leech taking the call (escapepod-signal v0.16.0).

Two primitives delegate, both byte-identical to what was here:

- `encoding.rs::encode_signal_kmer_inner` -> `seq_encoding::encode_signal_kmer`
- `features.rs::sequence_to_int` -> `seq_encoding::sequence_to_int`, the same
  A/C/G/T-U table with the same `-1` for anything else (upstream's
  `UNKNOWN_BASE` is `-1`)

A third overlap is deliberately NOT delegated. The k-mer context slice in
`signal_mapping.rs` covers the same range with the same padding as upstream's
`sequence_ints_with_context`, but leech needs it as BASES — it serializes
`sequence_with_kmer_context` as a string — where upstream returns ints. Same
rule, different type; swapping it would change a serialized chunk field.

This is a delegation, so identity is the only acceptable outcome, and the
evidence is that nothing had to change to accommodate it: 198 parity tests pass
untouched, including `test_backend_parity.py`, which compares every array in the
npz between the Rust and Python backends — and the Python reference is not
touched by this commit, so it is an independent implementation.

Also here, because the pin moves anyway:

- `escapepod>=0.16.0`, and the full suite re-run against that Python package as
  well as against the Rust crate (they are independent: leech_core links the
  crate statically, `leech.io` imports the wheel). 1467 passed on both.
- `uv.lock` gains onnx/onnxruntime/onnxscript/protobuf. #220 added the `onnx`
  extra to pyproject and never locked it — CI installs from pyproject so it
  worked, but `uv sync --extra onnx` would not have.
- The ONNX contract stops telling consumers to reimplement the encoder and
  names the crate that owns it.
jayhesselberth added a commit to rnabioco/escapepod-rs that referenced this pull request Aug 25, 2026
Closes #274.

#272 moved the signal-level k-mer *encoding* upstream and leech now
calls it
(rnabioco/leech#222), but one overlap survived, and it was the same
shape of
problem: leech kept its own copy of the **context windowing**, because
`sequence_ints_with_context` returns ints and a training corpus
*serialises*
the context — `sequence_with_kmer_context` is a string in the chunk
format,
which `data merge`/`load_chunks` read back as one. So the downstream
caller
needs bases, and deriving them from the ints by hand would be a third
copy of
the window rather than the end of the second.

## The equivalence this rests on

leech's window (`rust/src/inference_pipeline/signal_mapping.rs`) is
`(seq_start - before) .. (seq_end + after)`; this crate's is
`lo = core_start - before`, `hi = lo + n_bases + before + after`. With
`n_bases = seq_end - seq_start` those are the same half-open range and
the same
"pad, don't shift" rule — and the paddings are each other's, since
`base_to_int(UNKNOWN_BASE_CHAR)` is `UNKNOWN_BASE`. That is what makes
this a
refactor rather than new behaviour.

## What changed

`crates/escapepod-signal/src/seq_encoding.rs`:

- **`UNKNOWN_BASE_CHAR`** (`b'N'`) — documented as *the* reason the two
forms
  compose, not as a stray constant.
- **`sequence_bases_with_context`** — the same cut as bases, padded with
`UNKNOWN_BASE_CHAR`. Lower-case and ambiguity codes pass through
verbatim:
  this form is the sequence, not an encoding of it.
- **`context_range`** (private) — the window arithmetic, now in exactly
one
  place, called by both public forms so they cannot drift.
- `sequence_ints_with_context` — signature unchanged, body now goes
through
  `context_range`, plus a cross-reference to the bases form.

The ints form is deliberately **not**
`sequence_to_int(&sequence_bases_…)`.
That would enforce the identity structurally, at the cost of a `Vec<u8>`
per
chunk on the inference hot path; sharing `context_range` gives the same
non-drift guarantee, and the identity is pinned by a test instead.

## Tests

- `the_bases_form_is_the_same_cut_as_the_int_form` — asserts
`base_to_int(UNKNOWN_BASE_CHAR) == UNKNOWN_BASE` outright, then mirrors
the
  existing int cases (`CGTA`, `NACG`, `GTAN`, all-`N`).
- `bases_and_ints_are_one_window` — the one that matters:
`sequence_to_int(bases) == ints` over 5 contexts × every offset from 0
to
`len + 4` × 4 window widths, so both padding edges and the
wholly-off-the-end
case are covered. If it ever fails, the two windows have diverged, which
is
  exactly the failure this is meant to make impossible.
- Two doctests mirroring the int form's, so the pair reads together.

## Verification

On `rna`: `cargo nextest run --workspace` 708/708 passed, `cargo clippy
--workspace --all-targets` clean, `cargo fmt --all --check` clean,
`cargo test --doc -p escapepod-signal` 4/4.

## Downstream

leech's follow-up is a small deletion in `signal_mapping.rs` plus a
version
bump, gated on a release carrying this. Its `test_backend_parity.py`
compares
every array in the npz between the Rust and Python backends, and
`sequence_with_kmer_context` is one of the compared fields, so the swap
is
verifiable as identity on real chunks.
…l too

The third and last overlap, and the one that needed an upstream change first
(escapepod-rs#274, released in 0.16.1). It could not delegate before: leech
needs the context window as BASES, because the corpus serializes
`sequence_with_kmer_context` as a string, where upstream only offered ints.
escapepod-signal now exposes both forms over one windowing rule, with
`sequence_to_int(bases) == ints` pinned by a test there.

So all three halves of the signal-level k-mer path live upstream now:
`mapping` produces the base-to-signal map, `sequence_bases_with_context` cuts
the window it covers plus the k-mer context, and `encode_signal_kmer` scatters
the one-hot context along the signal axis. leech held the middle and the last
of those, inside a cdylib no Rust consumer could link.

This is the one of the three most worth getting out of a second copy. It is
where `before` and `after` are not interchangeable — swap them and every k-mer
is read from a window displaced by `before - after` bases, silently, and
`encode_signal_kmer` cannot detect it because it only sees the total width.

Identical by construction and by test: the ranges agree
(`(seq_start - before)..(seq_end + after)` is upstream's
`core_start - before` plus `n_bases + before + after` with
`n_bases = seq_end - seq_start`), and both pad rather than shift, leech with
`b'N'` and upstream with `UNKNOWN_BASE_CHAR`, which is `b'N'`. 198 parity tests
pass untouched, and `sequence_with_kmer_context` is one of the fields
`test_backend_parity.py` compares array-by-array between the Rust and Python
backends — the Python side is not touched here, so it is an independent check.

Only `rust/Cargo.toml`'s git tag moves to v0.16.1; `Cargo.lock` pins it at
b4c9afae. The `escapepod` Python pin stays `>=0.16.0`, because the new function
is in the Rust crate and not in the Python bindings, and the two are
independent — leech_core links the crate statically, `leech.io` imports the
wheel.

Full suite 1467 passed, 44 skipped.
@jayhesselberth jayhesselberth changed the title refactor(rust): take the signal-level k-mer encoding from escapepod-signal refactor(rust): take the signal-level k-mer path from escapepod-signal Aug 26, 2026
@jayhesselberth
jayhesselberth merged commit 07b560f into main Aug 26, 2026
3 checks passed
@jayhesselberth
jayhesselberth deleted the feat/upstream-kmer-encoding branch August 26, 2026 00:17
jayhesselberth added a commit that referenced this pull request Aug 26, 2026
`test_backend_parity.py` compares every array in the npz between the two
`data prepare` backends, which is why it has caught what narrower checks missed.
It was not comparing the padding branch of the k-mer context window at all: at
the default `signal_context=(200, 200)` no fixture read's window runs off the
end of the read, so 0 of 18 chunks contain an `N` and every case in the file
exercised only the in-range path.

That is the failure shape this stack keeps hitting. `extract_levels` was written
twice with different centring conventions; escapepod-classify reproduced a
superseded feature definition for two months and its counted golden missed the
wrong fallback because all 19 fixture reads took the other branch — found
against a real corpus, in 4 reads out of 842.

It matters more since #222 made the windowing a call into escapepod-signal. The
two sides pad against bounds that SOUND different — `len(sequence)` on the
Python side, the sequence slice on the Rust side — and nothing compared them
where the difference would show. (They are the same number: `num_bases` is
`len(self.sequence)`. It is `num_mapped_bases` that can be shorter, and it is
not the bound used here. Checked, because the docs warn the two are not
interchangeable.)

A wider signal window covers more bases and pushes the context past the read:
9 of 18 chunks padded at 2000, all 18 at 6000. The backends agree at both, so
this documents an existing property rather than fixing a bug — but it is now
demonstrated instead of resting on a reading of four lines of Rust.

The test asserts the padding actually occurred. Without that, a fixture change
that stopped producing edge chunks would leave it passing while testing
nothing — verified by running the assertion against the default config, where
it fires.

CLAUDE.md gains the general form next to the other parity rules: a parity test
only covers the branches its fixtures reach, so when you add a branch, ask which
fixture reaches it.

Full suite 1469 passed, 44 skipped.
jayhesselberth added a commit that referenced this pull request Aug 26, 2026
Minor rather than patch: new capability throughout, and two behaviour changes
-- one confined to CRF training, one to how a corpus that cannot supply
`signal_kmer` is handled.

The release is the second half of the CTC-CRF port plus ONNX export:

- `leech.crf.evaluate` (#224) -- decode a corpus, match to references by edit
  distance, report per group. The generic half of evaluation; what a panel is
  stays with whatever defines the panel.
- ONNX export for the classifier arms and the CRF encoder (#217, #222), dynamo
  exporter at opset 18, each with a contract sidecar and a round-trip check
  against torch across the serialization boundary.
- `leech model train-crf` (#219) -- the CLI for the trainer, plus the corpus
  builder (`plan_corpus`/`build_corpus`) and `CrfTrainer` itself.
- The signal-level k-mer encoding now comes from escapepod-signal (#222)
  rather than being held in a cdylib no Rust consumer could link.

Two behaviour changes, both worth reading before upgrading:

`signal_kmer` no longer degrades quietly (#230/#232). The encoding is decided
from the whole corpus rather than chunk 0, an encoding named on the command
line is no longer substituted, and the config records what the run actually
used. This one is coupled to the ONNX work above and is why the release waited
for it: the contract is derived from the config and exists so a non-Python
consumer can trust the input spec, so a config that misstates its encoding is
now refused at export rather than published.

CRF batch order (#231). `CrfTrainer.train` re-seeded `default_rng(seed)` and
replayed the permutation `resolve_split` had already drawn, so epoch 1 trained
on `pi(pi(train))`. Fixed, which means a given seed now sees different batches
-- numbers from a seed will not reproduce against 0.8.0. Batch order alone
moves a 32-epoch run's final training loss by more than 2x, so a seed is one
draw from that spread, not a fixed point.

The CRF trainer was validated against the implementation it was ported from
over six paired seeds: balanced recall differs by -0.17pp +/- 0.28pp, sign test
p = 0.688, against a within-arm seed range of 0.71pp.

Also in this commit, not from the PRs:

- README listed neither `leech model train-crf` (a shipped command missing
  from the CLI table) nor ONNX export at all, including the single-BCE-logit
  contract point that makes a misread graph silently wrong.
- CLAUDE.md said "feature-complete (v0.7.0)" while listing CRF and ONNX.
- CHANGELOG's Unreleased section had accumulated three separate `### Added`
  headings from different PRs; consolidated to one Added/Changed/Fixed set.

Full suite 1512 passed, 44 skipped. Docs build clean. Both lockfiles verified
against their manifests (`cargo metadata --locked`, `uv lock --check`).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant