CRF Module¶

def encoder_config_from_toml(cfg: dict) -> EncoderConfig:
    """Build an `EncoderConfig` from the architecture TOML.

    One reader for the config, so the trainer, the exporter and every evaluation
    script derive geometry the same way instead of each reaching into the dict
    for the keys it happens to need. `n_base` is the alphabet minus the CTC
    blank; it is not a key of its own, and reading it as one is a `KeyError`
    that only fires on the config's first use.
    """
    enc = cfg.get("encoder", {})
    return EncoderConfig(
        n_base=len(cfg["labels"]["labels"]) - 1,
        state_len=cfg["global_norm"]["state_len"],
        features=enc.get("features", 96),
        winlen=enc.get("winlen", 31),
        stride=enc.get("stride", 10),
        scale=enc.get("scale", 5.0),
        blank_score=enc.get("blank_score", BLANK_SCORE),
    )

def load_crf_state_dict(
    model: CrfEncoder, state_dict: dict[str, Tensor], *, strict: bool = True
) -> CrfEncoder:
    """Load a checkpoint into this module, in either naming.

    Accepts both because both exist and will keep existing: every model shipped
    before the trainer moved off bonito carries bonito's positional names
    (``encoder.4.rnn.*``), and everything trained since carries this module's
    (``rnns.0.rnn.*``). The parameters are identical — same layers, same shapes —
    so the only difference is the key, and refusing one naming would strand
    either the shipped models or the new ones.

    Raises if the checkpoint contains anything unmapped, rather than silently
    leaving a layer at its random initialisation.
    """
    mapping = _legacy_key_map(model.cfg.n_layers)
    # Which naming, decided by prefix rather than by whether the keys happen to
    # map. bonito's live under its flat container (`encoder.<i>.…`); this
    # module's are `conv.*`, `rnns.*`, `linear.*`. Asking "do any keys map?"
    # instead would classify an unmapped bonito key as native naming and report
    # the mismatch as torch's generic missing-keys error rather than as the
    # unknown-layer error that says what actually went wrong.
    if not any(k.startswith("encoder.") for k in state_dict):
        model.load_state_dict(state_dict, strict=strict)
        return model

    unmapped = sorted(set(state_dict) - set(mapping))
    if unmapped and strict:
        raise KeyError(
            f"checkpoint has {len(unmapped)} key(s) this module does not know: "
            f"{unmapped[:5]}{'...' if len(unmapped) > 5 else ''}"
        )
    renamed = {mapping[k]: v for k, v in state_dict.items() if k in mapping}
    model.load_state_dict(renamed, strict=strict)
    return model

def predecessor_index(n_base: int, state_len: int, device=None) -> Tensor:
    """Incoming edges per state: `(n_states, n_base + 1)` of source-state ids.

    Column 0 is the stay edge (the state itself); columns `1..n_base` are the
    move edges, ordered by the base that falls out of the k-mer window. That
    ordering is not free choice — it has to match the score layout the encoder
    emits, `state * (n_base + 1) + label`, or the scores get paired with the
    wrong transitions and nothing complains.
    """
    n_states = n_base**state_len
    states = torch.arange(n_states, device=device)
    stay = states[:, None]
    # Predecessors share our first state_len-1 bases: shift right, then vary the
    # base that was dropped off the front.
    shifted = states // n_base
    moves = shifted[None, :] + (
        torch.arange(n_base, device=device)[:, None] * n_base ** (state_len - 1)
    )
    return torch.cat([stay, moves.T], dim=1).long()

@torch.no_grad()
def decode_batch(
    scores: Tensor,
    n_base: int = 4,
    state_len: int = 4,
    alphabet: str = ALPHABET,
) -> list[str]:
    """Scores `(T, N, n_states*(n_base+1))` -> one sequence per read.

    Drop-in for `bonito.crf.model.Model.decode_batch`, including its output type:
    a plain list of `str`, one per read, in batch order.

    The emitted sequence is `target[state_len:]` of whatever the model was
    trained on — the model cannot emit the first `state_len` bases at any window
    width (issue #36). Match decodes against `ldxlib.emitted_refs()`, not against
    the full-length target.
    """
    dest, edge = best_path(scores, n_base, state_len)
    # Emit `dest % n_base` on a move, blank (0) on a stay: see the module note on
    # why the base comes from the destination and not from the edge index.
    symbols = torch.where(edge > 0, dest % n_base + 1, torch.zeros_like(dest))
    out = symbols.t().to(torch.uint8).cpu().numpy()  # (N, T), batch-major

    # Byte lookup, then one `tobytes()` per READ. The obvious form,
    # `"".join(alphabet[s] for s in row if s)`, is one interpreter step per
    # SYMBOL — 76,800 of them for a 256-read batch at T=300, which profiled at
    # 3.79 ms, the second-largest cost in the whole decode behind the encoder.
    # This is the same computation with the per-symbol work pushed into numpy.
    lut = np.frombuffer(alphabet.encode(), dtype=np.uint8)
    chars = lut[out]
    blank = lut[0]
    return [row[row != blank].tobytes().decode() for row in chars]

def best_path(scores: Tensor, n_base: int = 4, state_len: int = 4) -> tuple[Tensor, Tensor]:
    """Per-timestep best edge, as `(destination_state, edge)` tensors `(T, N)`.

    `edge == 0` is the stay/blank edge and emits nothing. Exposed separately from
    `decode_batch` because the position profile analysis needs *where* bases were
    emitted, not just the string (`scripts/ldx/decode_position_profile.py`).
    """
    n_states, n_edges = n_base**state_len, n_base + 1
    t_len, batch, width = scores.shape
    if width != n_states * n_edges:
        raise ValueError(
            f"score width {width} != n_states*(n_base+1) = {n_states * n_edges}; "
            f"1024 here means the blank was never expanded per state"
        )

    ms = scores.to(torch.float32).reshape(t_len, batch, n_states, n_edges)
    _, post = _full_impl(ms, n_base, n_states)
    lp = torch.log(post + POSTERIOR_FLOOR)
    best = _best_edges_impl(lp, n_base, n_states)
    return best // n_edges, best % n_edges

def load_manifest(path: str | Path, *, require: tuple[str, ...] = ()) -> CrfManifest:
    """Read and validate a manifest from parquet, TSV or CSV.

    Args:
        path: manifest file; format is taken from the suffix.
        require: optional columns this caller additionally needs. Named up front
            so the failure is "your manifest has no `batch` column" rather than
            a `KeyError` an hour into a run.

    Raises:
        FileNotFoundError: the manifest is not there.
        ValueError: a required column is missing, the table is empty, or
            ``read_id`` is not unique.
    """
    import polars as pl  # lazy: `leech.crf` must import only torch and numpy

    path = Path(path)
    if not path.is_file():
        raise FileNotFoundError(f"manifest not found: {path}")

    suffix = path.suffix.lower()
    if suffix == ".parquet":
        frame = pl.read_parquet(path)
    elif suffix in (".tsv", ".txt"):
        frame = pl.read_csv(path, separator="\t")
    elif suffix == ".csv":
        frame = pl.read_csv(path)
    else:
        raise ValueError(
            f"unrecognised manifest format {suffix!r} ({path}); "
            f"expected .parquet, .tsv, .txt or .csv"
        )

    missing = [c for c in REQUIRED_COLUMNS if c not in frame.columns]
    if missing:
        raise ValueError(
            f"{path}: manifest is missing required column(s) {missing}. "
            f"Required: {list(REQUIRED_COLUMNS)}; found: {list(frame.columns)}. "
            f"`target` is the resolved sequence — if you are porting a table that "
            f"carries a class name instead, the class -> sequence lookup belongs "
            f"in whatever produced it, not here."
        )

    unknown_required = [c for c in require if c not in frame.columns]
    if unknown_required:
        raise ValueError(
            f"{path}: this run needs optional column(s) {unknown_required}, which "
            f"the manifest does not carry. "
            + "; ".join(
                f"without `{c}`: {OPTIONAL_COLUMNS.get(c, 'unknown column')}"
                for c in unknown_required
            )
        )

    if frame.height == 0:
        raise ValueError(f"{path}: manifest is empty")

    duplicates = frame.height - frame["read_id"].n_unique()
    if duplicates:
        raise ValueError(
            f"{path}: {duplicates} duplicate read_id(s). One row per read — a "
            f"duplicate silently double-weights those reads in training and puts "
            f"the same read on both sides of a split."
        )

    return CrfManifest(frame=frame, path=path)

def check_geometry(
    window: int, target_len: int, samples_per_base: float, *, state_len: int = 4
) -> None:
    """Refuse a window that cannot hold ``target_len`` bases of signal.

    Raises rather than warns. A short window does not fail loudly at training
    time — it trains, converges, and quietly discriminates on fewer bases than
    the design intended, which is how a 27-nt barcode came to be classified on
    23 of them (rnabioco/escapepod-models#36).

    ``samples_per_base`` is the read population's own translocation rate.
    Measure it (leech has dwell times; take a median) rather than carrying a
    constant — it is chemistry- and speed-dependent, and a stale constant makes
    this check pass when it should not.
    """
    if target_len <= state_len:
        raise ValueError(
            f"target of {target_len} bases emits nothing at state_len={state_len}: "
            f"the first {state_len} bases only fix the initial state, so a target "
            f"must be longer than {state_len} to decode to anything."
        )
    needed = target_len * samples_per_base
    if window < needed:
        raise ValueError(
            f"window of {window} samples cannot hold {target_len} bases at "
            f"{samples_per_base:.1f} samples/base (needs ~{needed:.0f}). "
            f"The model would train on a truncated target and report nothing: "
            f"widen the window, or shorten the target and accept that it emits "
            f"{target_len - state_len} bases."
        )

def emitted_target(target: str, state_len: int) -> str:
    """The part of ``target`` the model can actually emit: ``target[state_len:]``.

    The first ``state_len`` bases fix the initial state and are never emitted,
    at any window width — widening the signal window does not lengthen the
    decode. Match decodes against this, never against the full-length target:
    the full-length comparison still picks the right reference but inflates
    every edit distance and compresses the confidence margin that ranking
    depends on.
    """
    if state_len < 0:
        raise ValueError(f"state_len must be >= 0, got {state_len}")
    return target[state_len:]

def plan_corpus(
    manifest: CrfManifest | str | Path,
    *,
    chunk: int,
    test_frac: float = 0.1,
    per_group: int | str = "auto",
    seed: int = 0,
    groups: list[str] | None = None,
    shard_batches: list[str] | None = None,
) -> CorpusPlan:
    """Decide the corpus. Reads no signal.

    Args:
        manifest: a :class:`CrfManifest` or a path to one.
        chunk: window width in samples; also sets the ``anchor_end`` filter.
        test_frac: fraction of each class held out, before capping.
        per_group: cap per class, or ``"auto"`` for the rarest class's trainable
            depth — the only value that actually balances.
        seed: fixes the shuffle, and therefore the split.
        groups: restrict to these classes (default: every class present).
        shard_batches: keep only these batches *after* planning globally.

    Returns:
        A :class:`CorpusPlan` whose frame carries a ``split`` column.
    """
    import polars as pl

    man = manifest if isinstance(manifest, CrfManifest) else load_manifest(manifest)
    frame = _group_column(man.frame)
    if "batch" not in frame.columns:
        frame = frame.with_columns(pl.lit("all").alias("batch"))

    if groups is not None:
        frame = frame.filter(pl.col("group").is_in(groups))
        if frame.height == 0:
            raise ValueError(f"no reads for groups {groups}")

    before = frame.height
    frame = frame.filter(pl.col("anchor_end") > chunk + ANCHOR_MARGIN)
    dropped = before - frame.height
    if frame.height == 0:
        raise ValueError(
            f"no read has anchor_end > {chunk + ANCHOR_MARGIN}; every one of "
            f"{before} was dropped. The window is wider than the signal available "
            f"ahead of the anchor."
        )

    # Shuffle within batch, then interleave round-robin across batches.
    parts = []
    for _batch, part in sorted(
        ((k[0], p) for k, p in frame.partition_by("batch", as_dict=True).items()),
        key=lambda kv: str(kv[0]),
    ):
        part = part.sort("read_id").sample(fraction=1.0, shuffle=True, seed=seed)
        parts.append(part.with_columns(pl.int_range(pl.len()).alias("_pos")))
    frame = pl.concat(parts).sort("_pos", maintain_order=True)

    avail = frame.group_by("group").len().sort("group")
    counts = dict(zip(avail["group"].to_list(), avail["len"].to_list(), strict=True))
    trainable = {g: int(n * (1 - test_frac)) for g, n in counts.items()}
    rarest_group = min(trainable, key=lambda g: trainable[g])
    rarest = trainable[rarest_group]

    if per_group == "auto":
        cap = rarest
        logger.info("per_group=auto -> %d (trainable depth of %s)", cap, rarest_group)
    else:
        cap = int(per_group)
        if cap > rarest:
            short = sorted(g for g, n in trainable.items() if n < cap)
            logger.warning(
                "per_group=%d exceeds the trainable depth of %d group(s); the corpus "
                "will NOT be balanced. Rarest is %s at %d — those groups contribute "
                "everything they have while the others are capped. Use per_group='auto' "
                "(%d) for a balanced corpus.",
                cap,
                len(short),
                rarest_group,
                rarest,
                rarest,
            )

    # Rank within class over the interleaved order, so rank is a random draw.
    # Per class and global across batches: per-(batch, class) would multiply the
    # cap by the number of batches whenever classes are crossed with batch.
    ranked = frame.with_columns(
        pl.int_range(pl.len()).over("group").alias("_rank"),
        pl.len().over("group").alias("_n"),
    ).with_columns(
        # At least one held-out read per class: a class with none is invisible to
        # evaluation rather than reported as missing.
        pl.max_horizontal(pl.lit(1), (pl.col("_n") * test_frac).cast(pl.Int64)).alias("_n_test")
    )
    test = ranked.filter(pl.col("_rank") < pl.col("_n_test")).with_columns(
        pl.lit("test").alias("split")
    )
    train = (
        ranked.filter(pl.col("_rank") >= pl.col("_n_test"))
        .with_columns((pl.col("_rank") - pl.col("_n_test")).alias("_train_rank"))
        .filter(pl.col("_train_rank") < cap)
        .with_columns(pl.lit("train").alias("split"))
        .drop("_train_rank")
    )
    planned = pl.concat([test, train]).drop("_pos", "_rank", "_n", "_n_test")

    if shard_batches is not None:
        want = set(shard_batches)
        missing = want - set(planned["batch"].unique().to_list())
        if missing:
            raise ValueError(
                f"shard_batches names {len(missing)} batch(es) with no reads in the "
                f"plan: {sorted(missing)}. A shard that quietly extracted nothing "
                f"would shrink the corpus without failing."
            )
        planned = planned.filter(pl.col("batch").is_in(list(want)))

    return CorpusPlan(frame=planned, cap=cap, chunk=chunk, dropped_short_anchor=dropped)

def build_corpus(
    plan: CorpusPlan,
    out: str | Path,
    *,
    state_len: int = 4,
    allow_shortfall: bool = False,
    extract_batch: int = EXTRACT_BATCH,
) -> Path:
    """Extract each planned read's window and stream it to ``<out>_X.npy``.

    Args:
        plan: from :func:`plan_corpus`.
        out: output stem; ``_X.npy`` and ``_meta.npz`` are appended.
        state_len: recorded in the metadata so consumers can derive what the
            model will emit without guessing.
        allow_shortfall: permit extracting less than half the plan.
        extract_batch: reads per POD5 call.

    Raises:
        RuntimeError: nothing was extracted, or more than half the plan is
            missing and ``allow_shortfall`` is false.
    """
    # Imported here, not at module scope: `leech.io.pod5_reader` pulls pysam and
    # escapepod, and `leech.crf` promises to cost only torch and numpy on import.
    from leech.io.pod5_reader import read_pod5_signals_batch_cached

    out = Path(out)
    out.parent.mkdir(parents=True, exist_ok=True)
    x_path = out.with_name(out.name + "_X.npy")
    frame, chunk, total = plan.frame, plan.chunk, len(plan)

    signal = np.lib.format.open_memmap(x_path, mode="w+", dtype=np.float32, shape=(total, chunk))
    kept: dict[str, list] = {k: [] for k in ("target", "group", "read_id", "split", "batch")}
    quality: dict[str, list] = {"quality_score": [], "quality_margin": []}
    has_quality = {k: k in frame.columns for k in quality}
    n = 0

    # Grouped by POD5 source: that is the unit the reader caches and the unit
    # storage order is defined within. The split above is independent of it.
    for (source,), part in sorted(
        frame.partition_by("pod5", as_dict=True).items(), key=lambda kv: str(kv[0][0])
    ):
        rows = part.to_dicts()
        for start in range(0, len(rows), extract_batch):
            block = rows[start : start + extract_batch]
            found = read_pod5_signals_batch_cached(source, [r["read_id"] for r in block])
            for row in block:
                hit = found.get(row["read_id"])
                if hit is None:
                    continue
                raw = hit[0]
                end = int(row["anchor_end"])
                if end - chunk < 0 or end > len(raw):
                    continue
                signal[n] = np.asarray(raw[end - chunk : end], dtype=np.float32)
                for key in kept:
                    kept[key].append(row[key])
                for key, present in has_quality.items():
                    quality[key].append(float(row[key]) if present else np.nan)
                n += 1
        logger.info("%s: %d extracted", source, n)
    signal.flush()

    # An empty or badly short corpus is a BROKEN build, not a small one, and
    # everything upstream can be correct and still produce it — a manifest naming
    # a parent directory rather than the one holding the .pod5 files matches zero
    # reads, writes a 0-row array, and exits cleanly. The two failures are
    # distinguished because their causes differ.
    if n == 0:
        sources = sorted({str(s) for s in frame["pod5"].unique().to_list()})[:3]
        del signal
        raise RuntimeError(
            f"extracted 0 of {total} planned reads. Nothing was read from "
            f"{', '.join(sources)}. Check that these paths hold the .pod5 files "
            f"naming these reads."
        )
    if n < 0.5 * total and not allow_shortfall:
        del signal
        raise RuntimeError(
            f"extracted {n} of {total} planned reads ({n / total:.1%}). More than "
            f"half the corpus is missing, so the manifest and the POD5s disagree "
            f"about which reads exist. Pass allow_shortfall=True if expected."
        )

    if n < total:
        _trim(x_path, n, chunk)

    targets = np.array(kept["target"], dtype=str)
    meta = {
        "y": targets,
        "code": np.array(kept["group"], dtype=str),
        "group": np.array(kept["group"], dtype=str),
        "read_id": np.array(kept["read_id"], dtype=str),
        "split": np.array(kept["split"], dtype=str),
        "batch": np.array(kept["batch"], dtype=str),
        "run": np.array(kept["batch"], dtype=str),
        "gate_score": np.array(quality["quality_score"], dtype=np.float32),
        "gate_margin": np.array(quality["quality_margin"], dtype=np.float32),
        "chunk": chunk,
        "per_code": plan.cap,
        "target_len": int(len(targets[0])) if n else 0,
        "state_len": state_len,
    }
    np.savez(out.with_name(out.name + "_meta.npz"), **meta)
    logger.info("wrote %s (%d x %d) and %s_meta.npz", x_path, n, chunk, out)
    return x_path

def load_corpus(path: str | Path, *, mmap: bool = True):
    """Read a corpus written by :func:`build_corpus`, or the legacy single ``.npz``.

    Returns ``(signal, targets, groups, read_ids, split)``; ``split`` is ``None``
    for the legacy layout, which carries none. ``mmap`` keeps the signal on disk,
    which is the only way to train on a corpus larger than RAM — batches are
    indexed out of it per step.
    """
    base = Path(str(path)[:-4] if str(path).endswith(".npz") else path)
    streamed = base.with_name(base.name + "_X.npy")
    if streamed.exists():
        signal = np.load(streamed, mmap_mode="r" if mmap else None)
        meta = np.load(base.with_name(base.name + "_meta.npz"), allow_pickle=True)
        group_key = "group" if "group" in meta else "code"
        split = meta["split"].astype(str) if "split" in meta else None
        return (
            signal,
            meta["y"].astype(str),
            meta[group_key].astype(str),
            meta["read_id"].astype(str),
            split,
        )
    legacy_path = base.with_suffix(".npz")
    if not legacy_path.exists():
        # Name BOTH layouts. The stem is not a file, so reporting the one that
        # happened to be tried last points at a path the caller never wrote.
        raise FileNotFoundError(
            f"no corpus at {base}: expected either {streamed.name} (the streamed "
            f"layout, written by build_corpus) or {legacy_path.name} (the legacy "
            f"single-file one) beside it. `--corpus` takes the STEM, without a "
            f"suffix."
        )
    legacy = np.load(legacy_path, allow_pickle=True)
    group_key = "group" if "group" in legacy else "code"
    return (
        legacy["X"],
        legacy["y"].astype(str),
        legacy[group_key].astype(str),
        legacy["read_id"].astype(str),
        None,
    )