From d1b7504e74a51a46d75190eb8b42f5d6ac46b166 Mon Sep 17 00:00:00 2001 From: isvock Date: Mon, 3 Aug 2026 11:56:26 -0700 Subject: [PATCH 01/12] Add window ISM analysis --- docs/api.rst | 4 + docs/usage.md | 72 +++- .../analysis/window_ism_summary.py | 244 ++++++++++++++ src/transcriptml/cli/main.py | 43 ++- src/transcriptml/interpret/__init__.py | 2 + src/transcriptml/interpret/window_ism.py | 316 ++++++++++++++++++ tests/test_cli_analysis.py | 11 +- tests/test_window_ism.py | 307 +++++++++++++++++ 8 files changed, 989 insertions(+), 10 deletions(-) create mode 100644 src/transcriptml/analysis/window_ism_summary.py create mode 100644 src/transcriptml/interpret/window_ism.py create mode 100644 tests/test_window_ism.py diff --git a/docs/api.rst b/docs/api.rst index 49a722f..ac78868 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -100,6 +100,10 @@ Interpretation :members: ISMResult, compute_ism, max_abs_effect_per_position, save_ism_result :member-order: bysource +.. automodule:: transcriptml.interpret.window_ism + :members: WindowISMResult, generate_window_starts, compute_window_ism, save_window_ism_result + :member-order: bysource + .. automodule:: transcriptml.interpret.codon_ism :members: CodonISMResult, compute_codon_ism, mutation_table_writer, save_codon_ism_result :member-order: bysource diff --git a/docs/usage.md b/docs/usage.md index 7e1d86a..22e5e79 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -363,7 +363,71 @@ windows or whole transcripts, the letters become dense and distracting; use `--no-logo` in those cases. Use `--no-isoform` if the annotation track is not needed for a particular figure. -### 4. Run Motif Analyses +### 4. Run Window ISM + +Window ISM provides a coarser complement to single-nucleotide ISM. Each window +is independently mutated several times by replacing every nucleotide with a +uniformly sampled alternative base. For replicate `r`, the effect is +`mutant_prediction - reference_prediction`; the command records its signed +mean, mean absolute value, and population standard deviation. + +Run the same scan once per fold checkpoint: + +```bash +for fold in $(seq 0 9); do + transcriptml window-ism \ + --checkpoint "runs/saluki_cv10/fold${fold}/model/best.pt" \ + --dataset data/saluki \ + --out-dir "interpret/window_ism/fold${fold}" \ + --window-size 100 \ + --stride 100 \ + --n-ablations 30 \ + --seed 123 \ + --device auto \ + --batch-size 128 \ + --mutation-batch-size 512 +done +``` + +When `--stride` is omitted it defaults to `--window-size`. If a regular tiled +scan would miss the sequence tail, an overlapping final window is shifted so it +ends exactly at the valid sequence length. Therefore every unambiguous base is +covered when the valid sequence is at least as long as the window. Use a +smaller stride for an overlapping scan; stride must not exceed the window size. + +Each fold directory contains compact `(N, Wmax)` arrays. `window_starts.npy` +stores zero-based window starts with `-1` padding, and `window_mask.npy` +distinguishes scored windows from padding or windows containing ambiguous +bases. `mean_deltas.npy` preserves effect direction, +`mean_abs_deltas.npy` is the recommended window-ranking signal, and +`std_deltas.npy` measures variability across random mutations. Effects at +masked positions are zero and must be interpreted together with the mask. + +Aggregate matching fold scans with: + +```bash +transcriptml summarize-window-ism \ + --input-dir interpret/window_ism \ + --out-dir interpret/window_ism_summary \ + --dataset data/saluki +``` + +The summary writes `average_mean_deltas.npy`, +`average_mean_abs_deltas.npy`, `within_model_std_deltas.npy`, +`fold_std_mean_deltas.npy`, and `average_reference_predictions.npy`, together +with the shared coordinates, mask, and valid lengths. The two standard +deviations separate within-model random-mutation variability from between-fold +model disagreement. Fold inputs must use identical coordinates, masks, scan +parameters, seed, and sequence order. A later hierarchical workflow can rank +windows by `average_mean_abs_deltas.npy` and apply finer window or +single-nucleotide ISM only in the most sensitive regions. + +The approximate mutant count per sequence is `n_ablations * n_windows`. With +the default tiled stride, this is approximately +`(n_ablations * valid_length) / window_size`; a stride-one scan is substantially +more expensive. + +### 5. Run Motif Analyses Motif analyses are usually much cheaper than full-transcript ISM and are designed to hone in on the context specificity and syntax of a particular motif or set of motifs. @@ -417,7 +481,7 @@ The `--region` flag can be `5utr`, `cds`, or `3utr`. Omit it to analyze motif instances across the full transcript. Region-aware analyses require Saluki-style annotation channels. -### 5. Run Codon Analyses +### 6. Run Codon Analyses Lots of work has shown that the coding sequence of an mRNA strongly influences its stability. Codon analyses are designed to dissect Saluki's understanding of this influence. @@ -646,6 +710,10 @@ transcriptml motif-ablation \ --device auto ``` +Window ISM also works unchanged with four-channel MPRA bundles. Use the same +`window-ism` and `summarize-window-ism` commands shown in the Saluki workflow, +pointing them at the MPRA checkpoint and dataset directories. + Interpret MPRA results in the assay's exact reporter context. Single-base effects may reflect cryptic splice sites, unintended promoter activity, or other construct-specific behavior in addition to the intended RNA regulatory diff --git a/src/transcriptml/analysis/window_ism_summary.py b/src/transcriptml/analysis/window_ism_summary.py new file mode 100644 index 0000000..006fd77 --- /dev/null +++ b/src/transcriptml/analysis/window_ism_summary.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Any + +import numpy as np + +from transcriptml.interpret.window_ism import sequence_ids_sha256 + + +_ARRAY_NAMES = ( + "window_starts", + "window_mask", + "mean_deltas", + "mean_abs_deltas", + "std_deltas", + "reference_predictions", + "valid_lengths", +) +_MATCHED_SETTINGS = ("input_shape", "window_size", "stride", "n_ablations", "seed", "mutation_policy") + + +def _fold_sort_key(path: Path) -> tuple[int, int | str]: + match = re.fullmatch(r"fold(\d+)", path.name) + if match: + return 0, int(match.group(1)) + return 1, path.name + + +def find_window_ism_fold_dirs(input_dir: str | Path) -> list[Path]: + """Find complete ``fold*`` window-ISM result directories.""" + + root = Path(input_dir) + if not root.exists(): + raise FileNotFoundError(f"Window-ISM input directory does not exist: {root}") + fold_dirs = [ + path + for path in root.iterdir() + if path.is_dir() + and path.name.startswith("fold") + and (path / "summary.json").exists() + and all((path / f"{name}.npy").exists() for name in _ARRAY_NAMES) + ] + return sorted(fold_dirs, key=_fold_sort_key) + + +def _load_summary(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if value.get("analysis") != "window_ism": + raise ValueError(f"Expected a window_ism summary at {path}") + return value + + +def _validate_fold_inputs(fold_dirs: list[Path]) -> tuple[dict[str, Any], tuple[int, int]]: + if not fold_dirs: + raise FileNotFoundError("No complete fold*/ window-ISM outputs found") + + reference_summary = _load_summary(fold_dirs[0] / "summary.json") + reference_arrays = { + name: np.load(fold_dirs[0] / f"{name}.npy", mmap_mode="r", allow_pickle=False) + for name in _ARRAY_NAMES + } + effect_shape = tuple(int(value) for value in reference_arrays["mean_deltas"].shape) + if len(effect_shape) != 2: + raise ValueError(f"Expected window effects with shape (N, Wmax), got {effect_shape}") + for name in ("window_starts", "window_mask", "mean_abs_deltas", "std_deltas"): + if reference_arrays[name].shape != effect_shape: + raise ValueError(f"{fold_dirs[0] / f'{name}.npy'} shape does not match mean_deltas.npy") + if reference_arrays["valid_lengths"].shape != (effect_shape[0],): + raise ValueError("valid_lengths.npy must have shape (N,)") + if reference_arrays["reference_predictions"].shape != (effect_shape[0],): + raise ValueError("reference_predictions.npy must have shape (N,)") + + for fold_dir in fold_dirs[1:]: + summary = _load_summary(fold_dir / "summary.json") + for setting in _MATCHED_SETTINGS: + if summary.get(setting) != reference_summary.get(setting): + raise ValueError( + f"Window-ISM setting mismatch for {setting}: " + f"{fold_dir} has {summary.get(setting)!r}, expected {reference_summary.get(setting)!r}" + ) + ref_ids = reference_summary.get("sequence_ids_sha256") + fold_ids = summary.get("sequence_ids_sha256") + if fold_ids != ref_ids: + raise ValueError(f"Sequence ID ordering mismatch for {fold_dir}") + + arrays = { + name: np.load(fold_dir / f"{name}.npy", mmap_mode="r", allow_pickle=False) + for name in _ARRAY_NAMES + } + for name in ("mean_deltas", "mean_abs_deltas", "std_deltas", "reference_predictions"): + if arrays[name].shape != reference_arrays[name].shape: + raise ValueError( + f"Shape mismatch for {fold_dir / f'{name}.npy'}: " + f"{arrays[name].shape} != {reference_arrays[name].shape}" + ) + for name in ("window_starts", "window_mask", "valid_lengths"): + if not np.array_equal(arrays[name], reference_arrays[name]): + raise ValueError(f"Coordinate or mask mismatch for {fold_dir / f'{name}.npy'}") + return reference_summary, effect_shape + + +def summarize_window_ism_folds( + *, + input_dir: str | Path, + out_dir: str | Path, + dataset: str | Path | None = None, + batch_size: int = 256, + dtype: str | np.dtype = "float32", +) -> dict[str, Any]: + """Aggregate matching fold-level window-ISM tracks. + + Signed means and mean absolute effects are averaged across models. Random- + mutation variability is summarized as the root mean within-model variance, + while model disagreement is the population standard deviation of fold-level + signed means. + """ + + if int(batch_size) <= 0: + raise ValueError("batch_size must be positive") + output_dtype = np.dtype(dtype) + if not np.issubdtype(output_dtype, np.floating): + raise ValueError(f"dtype must be floating-point, got {output_dtype}") + + fold_dirs = find_window_ism_fold_dirs(input_dir) + reference_summary, effect_shape = _validate_fold_inputs(fold_dirs) + n_sequences, max_windows = effect_shape + n_folds = len(fold_dirs) + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + + starts = np.load(fold_dirs[0] / "window_starts.npy", mmap_mode="r", allow_pickle=False) + mask = np.load(fold_dirs[0] / "window_mask.npy", mmap_mode="r", allow_pickle=False) + lengths = np.load(fold_dirs[0] / "valid_lengths.npy", mmap_mode="r", allow_pickle=False) + np.save(out / "window_starts.npy", starts) + np.save(out / "window_mask.npy", mask) + np.save(out / "valid_lengths.npy", lengths) + + if dataset is not None: + from transcriptml.data.bundle import load_bundle + + bundle = load_bundle(dataset, mmap_mode="r") + if int(bundle.X.shape[0]) != n_sequences: + raise ValueError(f"Dataset has N={bundle.X.shape[0]}, expected N={n_sequences}") + expected_input_shape = tuple(int(value) for value in reference_summary["input_shape"]) + if tuple(int(value) for value in bundle.X.shape) != expected_input_shape: + raise ValueError( + f"Dataset X shape {bundle.X.shape} does not match saved input shape {expected_input_shape}" + ) + ids_digest = sequence_ids_sha256(bundle.ids) + saved_digest = reference_summary.get("sequence_ids_sha256") + if saved_digest is not None and ids_digest != saved_digest: + raise ValueError("Dataset sequence ID ordering does not match the window-ISM fold outputs") + (out / "ids.txt").write_text("\n".join(str(value) for value in bundle.ids) + "\n", encoding="utf-8") + + output_specs = { + "average_mean_deltas": effect_shape, + "average_mean_abs_deltas": effect_shape, + "within_model_std_deltas": effect_shape, + "fold_std_mean_deltas": effect_shape, + "average_reference_predictions": (n_sequences,), + } + outputs = { + name: np.lib.format.open_memmap(out / f"{name}.npy", mode="w+", dtype=output_dtype, shape=shape) + for name, shape in output_specs.items() + } + fold_means = [np.load(path / "mean_deltas.npy", mmap_mode="r", allow_pickle=False) for path in fold_dirs] + fold_abs = [np.load(path / "mean_abs_deltas.npy", mmap_mode="r", allow_pickle=False) for path in fold_dirs] + fold_std = [np.load(path / "std_deltas.npy", mmap_mode="r", allow_pickle=False) for path in fold_dirs] + + for start in range(0, n_sequences, int(batch_size)): + end = min(start + int(batch_size), n_sequences) + chunk_shape = (end - start, max_windows) + mean_sum = np.zeros(chunk_shape, dtype=np.float64) + mean_sumsq = np.zeros(chunk_shape, dtype=np.float64) + abs_sum = np.zeros(chunk_shape, dtype=np.float64) + within_variance_sum = np.zeros(chunk_shape, dtype=np.float64) + for means, magnitudes, deviations in zip(fold_means, fold_abs, fold_std): + mean_values = np.asarray(means[start:end], dtype=np.float64) + std_values = np.asarray(deviations[start:end], dtype=np.float64) + mean_sum += mean_values + mean_sumsq += mean_values * mean_values + abs_sum += np.asarray(magnitudes[start:end], dtype=np.float64) + within_variance_sum += std_values * std_values + + average_mean = mean_sum / n_folds + between_variance = (mean_sumsq / n_folds) - (average_mean * average_mean) + np.maximum(between_variance, 0.0, out=between_variance) + outputs["average_mean_deltas"][start:end] = average_mean + outputs["average_mean_abs_deltas"][start:end] = abs_sum / n_folds + outputs["within_model_std_deltas"][start:end] = np.sqrt(within_variance_sum / n_folds) + outputs["fold_std_mean_deltas"][start:end] = np.sqrt(between_variance) + + ref_sum = np.zeros(n_sequences, dtype=np.float64) + for fold_dir in fold_dirs: + ref_sum += np.load(fold_dir / "reference_predictions.npy", mmap_mode="r", allow_pickle=False) + outputs["average_reference_predictions"][:] = ref_sum / n_folds + for values in outputs.values(): + values.flush() + + summary = { + "analysis": "window_ism_summary", + "input_dir": str(input_dir), + "out_dir": str(out), + "fold_count": n_folds, + "fold_dirs": [str(path) for path in fold_dirs], + "effect_definition": "mutant_prediction - reference_prediction", + "window_size": reference_summary["window_size"], + "stride": reference_summary["stride"], + "n_ablations": reference_summary["n_ablations"], + "seed": reference_summary["seed"], + "mutation_policy": reference_summary["mutation_policy"], + "shape": [n_sequences, max_windows], + "dtype": str(output_dtype), + "batch_size": int(batch_size), + "same_sequence_order_required": True, + "dataset": str(dataset) if dataset is not None else None, + "ids_written": dataset is not None, + "within_model_std_definition": "sqrt(mean_folds(std_deltas ** 2))", + "fold_std_definition": "population std across fold mean_deltas", + "outputs": { + **{name: str(out / f"{name}.npy") for name in output_specs}, + "window_starts": str(out / "window_starts.npy"), + "window_mask": str(out / "window_mask.npy"), + "valid_lengths": str(out / "valid_lengths.npy"), + "ids": str(out / "ids.txt") if dataset is not None else None, + }, + } + (out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + return summary + + +def run_window_ism_summary_from_args(args: Any) -> dict[str, Any]: + """Run fold-level window-ISM aggregation from parsed CLI arguments.""" + + return summarize_window_ism_folds( + input_dir=args.input_dir, + out_dir=args.out_dir, + dataset=args.dataset, + batch_size=args.batch_size, + dtype=args.dtype, + ) diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index fcf277a..1c27596 100644 --- a/src/transcriptml/cli/main.py +++ b/src/transcriptml/cli/main.py @@ -186,6 +186,7 @@ def build_parser() -> argparse.ArgumentParser: for name, help_text in [ ("ism", "Run single-nucleotide ISM"), + ("window-ism", "Run window-level random-mutagenesis ISM"), ("codon-ism", "Run CDS codon-level ISM"), ("motif-ablation", "Run motif ablation"), ("motif-context", "Run motif context scan"), @@ -200,7 +201,7 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument("--out-dir", dest="out_dir_flag", help="Output directory") p.add_argument("--device", default="cpu") p.add_argument("--batch-size", type=int, default=128) - if name not in {"ism", "codon-ism"}: + if name in {"motif-ablation", "motif-context", "epistasis"}: p.add_argument("--motif", required=True) p.add_argument( "--region", @@ -214,8 +215,13 @@ def build_parser() -> argparse.ArgumentParser: choices=["random_different", "shuffle", "dinuc_shuffle"], ) p.add_argument("--seed", type=int, default=123) - if name in {"ism", "codon-ism"}: + if name in {"ism", "window-ism", "codon-ism"}: p.add_argument("--mutation-batch-size", type=int, default=512) + if name == "window-ism": + p.add_argument("--window-size", type=int, required=True) + p.add_argument("--stride", type=int, help="Window stride; defaults to --window-size") + p.add_argument("--n-ablations", type=int, default=30) + p.add_argument("--seed", type=int, default=123) if name == "codon-ism": p.add_argument( "--mutation-policy", @@ -319,6 +325,13 @@ def build_parser() -> argparse.ArgumentParser: help="Also write centered scores projected onto reference bases; requires --dataset", ) + p = sub.add_parser("summarize-window-ism", help="Aggregate matching fold-level window-ISM tracks") + p.add_argument("--input-dir", type=Path, required=True, help="Directory containing fold*/ window-ISM outputs") + p.add_argument("--out-dir", type=Path, required=True, help="Directory for aggregated window-ISM arrays") + p.add_argument("--dataset", type=Path, help="Optional dataset bundle used to validate and write sequence IDs") + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--dtype", default="float32") + p = sub.add_parser("summarize-codon-ism", help="Summarize codon-ISM mutation tables") p.add_argument("--mode", required=True, choices=["synonymous", "all-codons"]) p.add_argument("--input-dir", type=Path, required=True) @@ -428,6 +441,11 @@ def main(argv: list[str] | None = None) -> None: run_ism_summary_from_args(args) return + if args.command == "summarize-window-ism": + from transcriptml.analysis.window_ism_summary import run_window_ism_summary_from_args + + run_window_ism_summary_from_args(args) + return if args.command == "summarize-codon-ism": common = [ "--out-dir", @@ -541,7 +559,7 @@ def main(argv: list[str] | None = None) -> None: out_dir = interpret_paths["out_dir"] log_progress(f"{args.command}: loading dataset {dataset}") - bundle = load_bundle(dataset, mmap_mode="r" if args.command == "codon-ism" else None) + bundle = load_bundle(dataset, mmap_mode="r" if args.command in {"codon-ism", "window-ism"} else None) log_progress(f"{args.command}: loading checkpoint {checkpoint}") predictor = Predictor.from_checkpoint(checkpoint, device=args.device, batch_size=args.batch_size) cds_channel = _maybe_int(getattr(args, "cds_channel", None)) @@ -550,6 +568,25 @@ def main(argv: list[str] | None = None) -> None: result = compute_ism(bundle.X, predictor, mutation_batch_size=args.mutation_batch_size) save_ism_result(result, out_dir) + elif args.command == "window-ism": + from transcriptml.interpret.window_ism import compute_window_ism, save_window_ism_result + + result = compute_window_ism( + bundle.X, + predictor, + window_size=args.window_size, + stride=args.stride, + n_ablations=args.n_ablations, + seed=args.seed, + mutation_batch_size=args.mutation_batch_size, + ) + save_window_ism_result( + result, + out_dir, + checkpoint=checkpoint, + dataset=dataset, + sequence_ids=bundle.ids, + ) elif args.command == "codon-ism": from transcriptml.interpret.codon_ism import compute_codon_ism, mutation_table_writer, save_codon_ism_result diff --git a/src/transcriptml/interpret/__init__.py b/src/transcriptml/interpret/__init__.py index fcdbcf3..8e57d9d 100644 --- a/src/transcriptml/interpret/__init__.py +++ b/src/transcriptml/interpret/__init__.py @@ -6,12 +6,14 @@ from transcriptml.interpret.epistasis import motif_epistasis from transcriptml.interpret.ism import compute_ism from transcriptml.interpret.predictor import EnsemblePredictor, Predictor +from transcriptml.interpret.window_ism import compute_window_ism __all__ = [ "EnsemblePredictor", "Predictor", "compute_codon_ism", "compute_ism", + "compute_window_ism", "motif_ablation", "motif_context_scan", "motif_epistasis", diff --git a/src/transcriptml/interpret/window_ism.py b/src/transcriptml/interpret/window_ism.py new file mode 100644 index 0000000..5852e1c --- /dev/null +++ b/src/transcriptml/interpret/window_ism.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np + +from transcriptml.data.encoding import infer_valid_lengths +from transcriptml.interpret.edits import scramble_window_inplace, valid_base_window +from transcriptml.interpret.predictor import Predictor +from transcriptml.progress import ProgressReporter, log_progress + + +@dataclass +class WindowISMResult: + """Window-level random-mutagenesis effects for a sequence batch.""" + + window_starts: np.ndarray + window_mask: np.ndarray + mean_deltas: np.ndarray + mean_abs_deltas: np.ndarray + std_deltas: np.ndarray + reference_predictions: np.ndarray + valid_lengths: np.ndarray + input_shape: tuple[int, int, int] + window_size: int + stride: int + n_ablations: int + seed: int + + +def generate_window_starts(valid_length: int, window_size: int, stride: int) -> np.ndarray: + """Generate fixed-width window starts, including a terminally anchored window. + + The final window ends exactly at ``valid_length``. Together with the + requirement ``stride <= window_size``, this guarantees full base coverage + for any sequence at least as long as the requested window. + + Args: + valid_length: Number of represented sequence positions. + window_size: Width of every window. + stride: Distance between regular window starts. + """ + + length = int(valid_length) + width = int(window_size) + step = int(stride) + if width <= 0: + raise ValueError("window_size must be positive") + if step <= 0 or step > width: + raise ValueError("stride must satisfy 1 <= stride <= window_size") + if length < width: + return np.empty((0,), dtype=np.int64) + + terminal_start = length - width + starts = list(range(0, terminal_start + 1, step)) + if starts[-1] != terminal_start: + starts.append(terminal_start) + return np.asarray(starts, dtype=np.int64) + + +def _normalize_valid_lengths(X: np.ndarray, valid_lengths: Sequence[int] | None) -> np.ndarray: + """Infer or validate one valid length per encoded sequence.""" + + lengths = infer_valid_lengths(X) if valid_lengths is None else np.asarray(valid_lengths, dtype=np.int64) + if lengths.shape != (X.shape[0],): + raise ValueError(f"valid_lengths must have shape ({X.shape[0]},), got {lengths.shape}") + if np.any(lengths < 0) or np.any(lengths > X.shape[-1]): + raise ValueError(f"valid_lengths entries must be between 0 and encoded length {X.shape[-1]}") + return lengths.astype(np.int64, copy=False) + + +def compute_window_ism( + X: np.ndarray, + predictor: Predictor, + *, + window_size: int, + stride: int | None = None, + n_ablations: int = 30, + seed: int = 123, + valid_lengths: Sequence[int] | None = None, + mutation_batch_size: int = 512, + progress: bool = True, +) -> WindowISMResult: + """Compute repeated random-mutagenesis effects for fixed-width windows. + + Every nucleotide in a scored window is independently replaced by a + uniformly sampled alternative base. Effects are signed mutant-minus- + reference prediction differences. Replicate-level effects are summarized + online as their mean, mean absolute value, and population standard + deviation. + + Args: + X: Encoded ``(N, C, L)`` sequence batch with at least four base + channels. + predictor: Predictor used to score reference and mutant sequences. + window_size: Number of bases mutated in each window. + stride: Distance between regular starts. Defaults to ``window_size``. + n_ablations: Number of independently mutated sequences per window. + seed: Non-negative base seed for deterministic per-window generators. + valid_lengths: Optional represented length for each sequence. + mutation_batch_size: Maximum number of mutants queued per prediction + call. + progress: Whether to emit progress messages. + """ + + arr = np.asarray(X) + if arr.ndim != 3 or arr.shape[1] < 4: + raise ValueError(f"Expected X with shape (N, C>=4, L), got {arr.shape}") + width = int(window_size) + step = width if stride is None else int(stride) + if width <= 0: + raise ValueError("window_size must be positive") + if step <= 0 or step > width: + raise ValueError("stride must satisfy 1 <= stride <= window_size") + if int(n_ablations) <= 0: + raise ValueError("n_ablations must be positive") + if int(mutation_batch_size) <= 0: + raise ValueError("mutation_batch_size must be positive") + if int(seed) < 0: + raise ValueError("seed must be non-negative") + + n_sequences = int(arr.shape[0]) + lengths = _normalize_valid_lengths(arr, valid_lengths) + starts_by_sequence = [generate_window_starts(int(length), width, step) for length in lengths] + max_windows = max((len(starts) for starts in starts_by_sequence), default=0) + starts_out = np.full((n_sequences, max_windows), -1, dtype=np.int64) + mask = np.zeros((n_sequences, max_windows), dtype=bool) + for seq_i, starts in enumerate(starts_by_sequence): + starts_out[seq_i, : len(starts)] = starts + for window_i, start in enumerate(starts.tolist()): + mask[seq_i, window_i] = valid_base_window(arr[seq_i], int(start), int(start) + width) + + log_progress(f"window-ism: predicting {n_sequences} reference sequences", enabled=progress) + reference = predictor.predict(arr).astype(np.float32, copy=False) + if reference.shape != (n_sequences,): + raise ValueError(f"predictor must return one scalar per sequence; got shape {reference.shape}") + + sums = np.zeros(mask.shape, dtype=np.float64) + abs_sums = np.zeros(mask.shape, dtype=np.float64) + square_sums = np.zeros(mask.shape, dtype=np.float64) + counts = np.zeros(mask.shape, dtype=np.int32) + mutant_batch: list[np.ndarray] = [] + mutant_meta: list[tuple[int, int]] = [] + + def flush() -> None: + """Predict queued mutants and update per-window moments.""" + + if not mutant_batch: + return + predictions = predictor.predict(np.stack(mutant_batch, axis=0)) + if predictions.shape[0] != len(mutant_meta): + raise ValueError("predictor returned an unexpected number of mutant predictions") + for prediction, (seq_i, window_i) in zip(predictions, mutant_meta): + delta = float(prediction - reference[seq_i]) + sums[seq_i, window_i] += delta + abs_sums[seq_i, window_i] += abs(delta) + square_sums[seq_i, window_i] += delta * delta + counts[seq_i, window_i] += 1 + mutant_batch.clear() + mutant_meta.clear() + + n_scored_windows = int(mask.sum()) + reporter = ProgressReporter( + "window-ism: scan windows", + total=n_scored_windows, + unit="windows", + enabled=progress, + ) + for seq_i in range(n_sequences): + for window_i in np.flatnonzero(mask[seq_i]).tolist(): + start = int(starts_out[seq_i, window_i]) + rng = np.random.default_rng(np.random.SeedSequence([int(seed), seq_i, start])) + for _ in range(int(n_ablations)): + mutant = arr[seq_i].copy() + scramble_window_inplace( + mutant, + start=start, + window_size=width, + strategy="random_different", + rng=rng, + ) + mutant_batch.append(mutant) + mutant_meta.append((seq_i, window_i)) + if len(mutant_batch) >= int(mutation_batch_size): + flush() + reporter.update() + flush() + reporter.close(extra=f"{int(counts.sum())} mutants predicted") + + if n_scored_windows and not np.all(counts[mask] == int(n_ablations)): + raise RuntimeError("Not all valid windows received the requested number of ablations") + mean = np.zeros(mask.shape, dtype=np.float32) + mean_abs = np.zeros(mask.shape, dtype=np.float32) + std = np.zeros(mask.shape, dtype=np.float32) + if n_scored_windows: + mean_values = sums[mask] / counts[mask] + mean_abs_values = abs_sums[mask] / counts[mask] + variance = (square_sums[mask] / counts[mask]) - (mean_values * mean_values) + variance = np.maximum(variance, 0.0) + mean[mask] = mean_values.astype(np.float32) + mean_abs[mask] = mean_abs_values.astype(np.float32) + std[mask] = np.sqrt(variance).astype(np.float32) + + return WindowISMResult( + window_starts=starts_out, + window_mask=mask, + mean_deltas=mean, + mean_abs_deltas=mean_abs, + std_deltas=std, + reference_predictions=reference, + valid_lengths=lengths, + input_shape=tuple(int(value) for value in arr.shape), + window_size=width, + stride=step, + n_ablations=int(n_ablations), + seed=int(seed), + ) + + +def sequence_ids_sha256(sequence_ids: Sequence[str]) -> str: + """Return a stable digest for an ordered collection of sequence IDs.""" + + digest = hashlib.sha256() + for sequence_id in sequence_ids: + digest.update(str(sequence_id).encode("utf-8")) + digest.update(b"\0") + return digest.hexdigest() + + +def _coverage_counts(result: WindowISMResult) -> tuple[int, int]: + """Return covered and uncovered valid-base counts for one result.""" + + covered_total = 0 + valid_total = int(np.asarray(result.valid_lengths, dtype=np.int64).sum()) + for seq_i, valid_length in enumerate(result.valid_lengths.tolist()): + covered = np.zeros(int(valid_length), dtype=bool) + for window_i in np.flatnonzero(result.window_mask[seq_i]).tolist(): + start = int(result.window_starts[seq_i, window_i]) + covered[start : start + int(result.window_size)] = True + covered_total += int(covered.sum()) + return covered_total, valid_total - covered_total + + +def save_window_ism_result( + result: WindowISMResult, + out_dir: str | Path, + *, + checkpoint: str | Path | None = None, + dataset: str | Path | None = None, + sequence_ids: Sequence[str] | None = None, + progress: bool = True, +) -> None: + """Save window-ISM arrays and reproducibility metadata.""" + + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + log_progress(f"window-ism: saving results to {out}", enabled=progress) + arrays = { + "window_starts": result.window_starts, + "window_mask": result.window_mask, + "mean_deltas": result.mean_deltas, + "mean_abs_deltas": result.mean_abs_deltas, + "std_deltas": result.std_deltas, + "reference_predictions": result.reference_predictions, + "valid_lengths": result.valid_lengths, + } + for name, values in arrays.items(): + np.save(out / f"{name}.npy", np.asarray(values)) + + candidate_mask = result.window_starts >= 0 + covered_bases, uncovered_bases = _coverage_counts(result) + n_short = int(np.count_nonzero(result.valid_lengths < result.window_size)) + ids_digest = None + if sequence_ids is not None: + if len(sequence_ids) != int(result.valid_lengths.shape[0]): + raise ValueError("sequence_ids length must match the number of result sequences") + ids_digest = sequence_ids_sha256(sequence_ids) + summary = { + "analysis": "window_ism", + "effect_definition": "mutant_prediction - reference_prediction", + "mutation_policy": "random_different_every_base", + "alternative_base_sampling": "uniform_over_other_three_bases", + "window_size": int(result.window_size), + "stride": int(result.stride), + "n_ablations": int(result.n_ablations), + "seed": int(result.seed), + "coordinate_convention": "zero_based_half_open", + "window_interval": "[start, start + window_size)", + "terminal_window_anchored": True, + "padding_start_value": -1, + "masked_effect_fill_value": 0.0, + "raw_replicates_saved": False, + "checkpoint": str(checkpoint) if checkpoint is not None else None, + "dataset": str(dataset) if dataset is not None else None, + "sequence_ids_sha256": ids_digest, + "n_sequences": int(result.valid_lengths.shape[0]), + "input_shape": list(result.input_shape), + "window_effect_shape": list(result.mean_deltas.shape), + "n_candidate_windows": int(candidate_mask.sum()), + "n_scored_windows": int(result.window_mask.sum()), + "n_ambiguous_windows": int(np.count_nonzero(candidate_mask & ~result.window_mask)), + "n_sequences_shorter_than_window": n_short, + "n_valid_bases": int(result.valid_lengths.sum()), + "n_covered_valid_bases": covered_bases, + "n_uncovered_valid_bases": uncovered_bases, + "arrays": { + name: {"shape": list(np.asarray(values).shape), "dtype": str(np.asarray(values).dtype)} + for name, values in arrays.items() + }, + } + (out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + log_progress("window-ism: done", enabled=progress) diff --git a/tests/test_cli_analysis.py b/tests/test_cli_analysis.py index a3fbd44..fdcac0b 100644 --- a/tests/test_cli_analysis.py +++ b/tests/test_cli_analysis.py @@ -103,13 +103,14 @@ def test_evaluate_cli_rejects_conflicting_named_and_positional_args(): def test_interpret_cli_resolves_named_and_legacy_positional_args(): parser = build_parser() - motif_extra = { + command_extra = { + "window-ism": ["--window-size", "100"], "motif-ablation": ["--motif", "AUG"], "motif-context": ["--motif", "AUG"], "epistasis": ["--motif", "AUG"], } - for command in ["ism", "codon-ism", "motif-ablation", "motif-context", "epistasis"]: + for command in ["ism", "window-ism", "codon-ism", "motif-ablation", "motif-context", "epistasis"]: named = parser.parse_args( [ command, @@ -119,7 +120,7 @@ def test_interpret_cli_resolves_named_and_legacy_positional_args(): "data/saluki", "--out-dir", f"interpret/{command}", - *motif_extra.get(command, []), + *command_extra.get(command, []), ] ) assert _resolve_interpret_args(named, parser) == { @@ -134,7 +135,7 @@ def test_interpret_cli_resolves_named_and_legacy_positional_args(): "model/best.pt", "data/saluki", f"interpret/{command}", - *motif_extra.get(command, []), + *command_extra.get(command, []), ] ) assert _resolve_interpret_args(positional, parser) == { @@ -150,7 +151,7 @@ def test_interpret_cli_resolves_named_and_legacy_positional_args(): "data/saluki", "--out-dir", f"interpret/{command}", - *motif_extra.get(command, []), + *command_extra.get(command, []), ] ) assert _resolve_interpret_args(mixed, parser)["out_dir"] == f"interpret/{command}" diff --git a/tests/test_window_ism.py b/tests/test_window_ism.py new file mode 100644 index 0000000..18dd394 --- /dev/null +++ b/tests/test_window_ism.py @@ -0,0 +1,307 @@ +import json + +import numpy as np +import pytest +import torch + +from transcriptml.analysis.window_ism_summary import summarize_window_ism_folds +from transcriptml.cli.main import main +from transcriptml.data.encoding import encode_rna_sequence, encode_saluki_transcript, encode_sequences +from transcriptml.interpret.predictor import Predictor +from transcriptml.interpret.window_ism import ( + WindowISMResult, + compute_window_ism, + generate_window_starts, + save_window_ism_result, +) + + +class BaseWeightModel(torch.nn.Module): + def __init__(self, weights): + super().__init__() + self.register_buffer("weights", torch.tensor(weights, dtype=torch.float32).view(1, 4, 1)) + + def forward(self, x): + return (x[:, :4, :] * self.weights).sum(dim=(1, 2)) + + +class PositionAModel(torch.nn.Module): + def __init__(self, weights): + super().__init__() + self.register_buffer("weights", torch.tensor(weights, dtype=torch.float32).view(1, 1, -1)) + + def forward(self, x): + return (x[:, 0:1, :] * self.weights).sum(dim=(1, 2)) + + +class RecordingCallable: + def __init__(self): + self.calls = [] + + def __call__(self, X): + values = np.asarray(X) + self.calls.append(values.copy()) + return values[:, :4, :].sum(axis=(1, 2), dtype=np.float32) + + +def test_generate_window_starts_anchors_terminal_window_and_covers_sequence(): + starts = generate_window_starts(valid_length=10, window_size=4, stride=4) + assert starts.tolist() == [0, 4, 6] + covered = np.zeros(10, dtype=bool) + for start in starts: + covered[start : start + 4] = True + assert covered.all() + + assert generate_window_starts(valid_length=3, window_size=4, stride=4).size == 0 + with pytest.raises(ValueError, match="stride"): + generate_window_starts(valid_length=10, window_size=4, stride=5) + + +@pytest.mark.parametrize(("length", "stride"), [(4, 4), (5, 4), (10, 4), (10, 3), (11, 1)]) +def test_scored_windows_cover_every_unambiguous_base(length, stride): + X = encode_rna_sequence("A" * length)[None].astype(np.float32) + result = compute_window_ism( + X, + Predictor(BaseWeightModel([1, 0, 0, 0])), + window_size=4, + stride=stride, + n_ablations=1, + progress=False, + ) + covered = np.zeros(length, dtype=bool) + for window_i in np.flatnonzero(result.window_mask[0]): + start = int(result.window_starts[0, window_i]) + covered[start : start + result.window_size] = True + assert covered.all() + + +def test_window_ism_exact_additive_effects_and_terminal_coordinates(): + X = encode_rna_sequence("AAAAA")[None].astype(np.float32) + result = compute_window_ism( + X, + Predictor(BaseWeightModel([1, 0, 0, 0])), + window_size=2, + stride=2, + n_ablations=3, + mutation_batch_size=2, + progress=False, + ) + + assert result.window_starts.tolist() == [[0, 2, 3]] + assert result.window_mask.tolist() == [[True, True, True]] + np.testing.assert_allclose(result.mean_deltas, -2.0) + np.testing.assert_allclose(result.mean_abs_deltas, 2.0) + np.testing.assert_allclose(result.std_deltas, 0.0) + assert result.reference_predictions.tolist() == [5.0] + + +def test_window_ism_compact_padding_ambiguous_mask_and_short_sequence(): + X = encode_sequences(["ACGUACGUAA", "ACGUACGU", "ACN"], length=10).astype(np.float32) + result = compute_window_ism( + X, + Predictor(BaseWeightModel([1, 2, 4, 8])), + window_size=4, + stride=4, + n_ablations=2, + progress=False, + ) + + assert result.window_starts.tolist() == [[0, 4, 6], [0, 4, -1], [-1, -1, -1]] + assert result.window_mask.tolist() == [[True, True, True], [True, True, False], [False, False, False]] + assert np.all(result.mean_deltas[~result.window_mask] == 0) + + ambiguous = encode_rna_sequence("ACNU")[None].astype(np.float32) + ambiguous_result = compute_window_ism( + ambiguous, + Predictor(BaseWeightModel([1, 2, 4, 8])), + window_size=2, + stride=2, + n_ablations=2, + progress=False, + ) + assert ambiguous_result.window_starts.tolist() == [[0, 2]] + assert ambiguous_result.window_mask.tolist() == [[True, False]] + + +def test_window_ism_preserves_annotation_channels(): + X = encode_saluki_transcript( + "ACGUAC", + length=6, + cds_positions=[0, 3], + splice_positions=[2, 5], + )[None].astype(np.float32) + recorder = RecordingCallable() + compute_window_ism( + X, + Predictor(recorder), + window_size=3, + stride=3, + n_ablations=2, + mutation_batch_size=3, + progress=False, + ) + + mutant_calls = recorder.calls[1:] + assert mutant_calls + for call in mutant_calls: + expected_annotations = np.repeat(X[:, 4:, :], call.shape[0], axis=0) + np.testing.assert_array_equal(call[:, 4:, :], expected_annotations) + reference_bases = np.argmax(X[0, :4, :], axis=0) + for mutant in call: + mutant_bases = np.argmax(mutant[:4, :], axis=0) + assert np.count_nonzero(mutant_bases != reference_bases) == 3 + + +def test_window_ism_is_reproducible_across_mutation_batch_sizes_and_seeded(): + X = encode_rna_sequence("AAAAAA")[None].astype(np.float32) + predictor = Predictor(BaseWeightModel([0, 1, 4, 9])) + kwargs = dict(window_size=3, stride=2, n_ablations=12, seed=17, progress=False) + first = compute_window_ism(X, predictor, mutation_batch_size=1, **kwargs) + second = compute_window_ism(X, predictor, mutation_batch_size=7, **kwargs) + np.testing.assert_array_equal(first.mean_deltas, second.mean_deltas) + np.testing.assert_array_equal(first.mean_abs_deltas, second.mean_abs_deltas) + np.testing.assert_array_equal(first.std_deltas, second.std_deltas) + + different = compute_window_ism(X, predictor, mutation_batch_size=7, **{**kwargs, "seed": 18}) + assert not np.array_equal(first.mean_deltas, different.mean_deltas) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"window_size": 0}, "window_size"), + ({"window_size": 3, "stride": 4}, "stride"), + ({"window_size": 3, "n_ablations": 0}, "n_ablations"), + ({"window_size": 3, "mutation_batch_size": 0}, "mutation_batch_size"), + ({"window_size": 3, "seed": -1}, "seed"), + ], +) +def test_window_ism_rejects_invalid_parameters(kwargs, message): + X = encode_rna_sequence("AAAA")[None].astype(np.float32) + with pytest.raises(ValueError, match=message): + compute_window_ism(X, Predictor(BaseWeightModel([1, 0, 0, 0])), progress=False, **kwargs) + + +def test_window_ism_mean_absolute_effect_ranks_high_weight_window(): + X = encode_rna_sequence("AAAAAA")[None].astype(np.float32) + predictor = Predictor(PositionAModel([1, 1, 10, 10, 1, 1])) + result = compute_window_ism( + X, + predictor, + window_size=2, + stride=2, + n_ablations=3, + progress=False, + ) + assert int(np.argmax(result.mean_abs_deltas[0])) == 1 + np.testing.assert_allclose(result.mean_abs_deltas[0], [2, 20, 2]) + + +def _synthetic_result(mean, magnitude, deviation, reference, *, starts=None): + mean = np.asarray(mean, dtype=np.float32) + starts_array = np.asarray(starts if starts is not None else [[0, 2]], dtype=np.int64) + return WindowISMResult( + window_starts=starts_array, + window_mask=starts_array >= 0, + mean_deltas=mean, + mean_abs_deltas=np.asarray(magnitude, dtype=np.float32), + std_deltas=np.asarray(deviation, dtype=np.float32), + reference_predictions=np.asarray(reference, dtype=np.float32), + valid_lengths=np.asarray([4], dtype=np.int64), + input_shape=(1, 4, 4), + window_size=2, + stride=2, + n_ablations=30, + seed=123, + ) + + +def test_save_window_ism_result_writes_arrays_metadata_and_coverage(tmp_path): + X = encode_rna_sequence("AAAAA")[None].astype(np.float32) + result = compute_window_ism( + X, + Predictor(BaseWeightModel([1, 0, 0, 0])), + window_size=2, + n_ablations=2, + progress=False, + ) + save_window_ism_result( + result, + tmp_path, + checkpoint="model.pt", + dataset="data/bundle", + sequence_ids=["seq0"], + progress=False, + ) + + for name in ( + "window_starts", + "window_mask", + "mean_deltas", + "mean_abs_deltas", + "std_deltas", + "reference_predictions", + "valid_lengths", + ): + assert (tmp_path / f"{name}.npy").exists() + summary = json.loads((tmp_path / "summary.json").read_text(encoding="utf-8")) + assert summary["analysis"] == "window_ism" + assert summary["n_valid_bases"] == 5 + assert summary["n_covered_valid_bases"] == 5 + assert summary["n_uncovered_valid_bases"] == 0 + assert summary["sequence_ids_sha256"] is not None + + +def test_summarize_window_ism_separates_within_and_between_model_variation(tmp_path): + input_dir = tmp_path / "window_ism" + fold0 = input_dir / "fold0" + fold1 = input_dir / "fold1" + save_window_ism_result( + _synthetic_result([[1, 3]], [[1, 3]], [[2, 4]], [10]), + fold0, + sequence_ids=["seq0"], + progress=False, + ) + save_window_ism_result( + _synthetic_result([[3, 7]], [[3, 7]], [[4, 8]], [14]), + fold1, + sequence_ids=["seq0"], + progress=False, + ) + + out_dir = tmp_path / "summary" + summary = summarize_window_ism_folds(input_dir=input_dir, out_dir=out_dir, batch_size=1) + np.testing.assert_allclose(np.load(out_dir / "average_mean_deltas.npy"), [[2, 5]]) + np.testing.assert_allclose(np.load(out_dir / "average_mean_abs_deltas.npy"), [[2, 5]]) + np.testing.assert_allclose(np.load(out_dir / "within_model_std_deltas.npy"), [[np.sqrt(10), np.sqrt(40)]]) + np.testing.assert_allclose(np.load(out_dir / "fold_std_mean_deltas.npy"), [[1, 2]]) + np.testing.assert_allclose(np.load(out_dir / "average_reference_predictions.npy"), [12]) + assert summary["fold_count"] == 2 + + +def test_summarize_window_ism_rejects_coordinate_mismatch(tmp_path): + input_dir = tmp_path / "window_ism" + save_window_ism_result( + _synthetic_result([[1, 2]], [[1, 2]], [[0, 0]], [1]), + input_dir / "fold0", + progress=False, + ) + save_window_ism_result( + _synthetic_result([[1, 2]], [[1, 2]], [[0, 0]], [1], starts=[[0, 1]]), + input_dir / "fold1", + progress=False, + ) + with pytest.raises(ValueError, match="Coordinate or mask mismatch"): + summarize_window_ism_folds(input_dir=input_dir, out_dir=tmp_path / "summary") + + +def test_summarize_window_ism_cli_smoke(tmp_path): + input_dir = tmp_path / "window_ism" + save_window_ism_result( + _synthetic_result([[1, 2]], [[1, 2]], [[0, 0]], [1]), + input_dir / "fold0", + progress=False, + ) + out_dir = tmp_path / "summary" + main(["summarize-window-ism", "--input-dir", str(input_dir), "--out-dir", str(out_dir)]) + assert (out_dir / "average_mean_deltas.npy").exists() From 989296d4d38de8221fcbd8d37c761e462a54c83e Mon Sep 17 00:00:00 2001 From: isvock Date: Tue, 11 Aug 2026 14:49:15 -0700 Subject: [PATCH 02/12] Add RBPNet preprocessing --- README.md | 8 +- docs/api.rst | 23 + docs/conf.py | 2 +- docs/index.rst | 6 +- docs/installation.md | 3 +- docs/rbpnet.md | 303 +++++++++ docs/usage.md | 4 + pyproject.toml | 12 +- src/transcriptml/cli/main.py | 9 + src/transcriptml/data/__init__.py | 96 ++- src/transcriptml/data/bundle.py | 53 +- src/transcriptml/rbpnet/__init__.py | 44 ++ src/transcriptml/rbpnet/_progress.py | 32 + src/transcriptml/rbpnet/annotation.py | 101 +++ src/transcriptml/rbpnet/bundle.py | 373 ++++++++++++ src/transcriptml/rbpnet/cli.py | 206 +++++++ src/transcriptml/rbpnet/coordinates.py | 148 +++++ src/transcriptml/rbpnet/experiment.py | 317 ++++++++++ src/transcriptml/rbpnet/fasta.py | 111 ++++ src/transcriptml/rbpnet/preprocessing.py | 263 ++++++++ src/transcriptml/rbpnet/selection.py | 741 +++++++++++++++++++++++ src/transcriptml/rbpnet/serialization.py | 127 ++++ src/transcriptml/rbpnet/signals.py | 360 +++++++++++ src/transcriptml/rbpnet/windows.py | 307 ++++++++++ tests/test_data.py | 15 + tests/test_rbpnet.py | 415 +++++++++++++ 26 files changed, 4018 insertions(+), 61 deletions(-) create mode 100644 docs/rbpnet.md create mode 100644 src/transcriptml/rbpnet/__init__.py create mode 100644 src/transcriptml/rbpnet/_progress.py create mode 100644 src/transcriptml/rbpnet/annotation.py create mode 100644 src/transcriptml/rbpnet/bundle.py create mode 100644 src/transcriptml/rbpnet/cli.py create mode 100644 src/transcriptml/rbpnet/coordinates.py create mode 100644 src/transcriptml/rbpnet/experiment.py create mode 100644 src/transcriptml/rbpnet/fasta.py create mode 100644 src/transcriptml/rbpnet/preprocessing.py create mode 100644 src/transcriptml/rbpnet/selection.py create mode 100644 src/transcriptml/rbpnet/serialization.py create mode 100644 src/transcriptml/rbpnet/signals.py create mode 100644 src/transcriptml/rbpnet/windows.py create mode 100644 tests/test_rbpnet.py diff --git a/README.md b/README.md index edcb2d0..d5cb359 100644 --- a/README.md +++ b/README.md @@ -6,16 +6,20 @@ APIs for preparing sequence datasets, training models, evaluating held-out predictions, and investigating learned sequence features with analyses such as in silico mutagenesis, motif ablation, context scans, etc. -TranscriptML currently supports two main workflows: +TranscriptML currently supports three main workflows: - **Saluki** predicts transcriptome-wide RNA stability from transcript sequence, coding-frame annotations, and splice sites. - **MPRA-LegNet** models MPRA measurements from variable sequence inserts and supports targets such as RNA stability, translation, protein output, etc. +- **RBPNet/eCLIP data** converts FASTA/GTF/BAM inputs into a canonical + transcript-space experiment, descriptive windows, explicit selection + manifests, and memory-mappable model-ready NumPy bundles. The RBPNet model + and trainer are not implemented yet. In the future, I plan to also support [RiboNN](https://www.nature.com/articles/s41587-025-02712-x) modeling of translation efficiency measurements -and [RBPNet](https://link.springer.com/article/10.1186/s13059-023-03015-7) modeling of RBP binding assays like eCLIP. +and complete [RBPNet](https://link.springer.com/article/10.1186/s13059-023-03015-7) model training and interpretation. ## Installation diff --git a/docs/api.rst b/docs/api.rst index ac78868..b83f4be 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -43,6 +43,29 @@ Sequence controls :members: SequenceControlOperation, SequenceControlConfig, normalize_sequence_control_config, apply_sequence_controls_array, apply_sequence_controls_to_bundle :member-order: bysource +RBPNet/eCLIP data +----------------- + +.. automodule:: transcriptml.rbpnet.preprocessing + :members: Sample, PipelineConfig, preprocess_eclip + :member-order: bysource + +.. automodule:: transcriptml.rbpnet.experiment + :members: ProcessedECLIPDataset, TranscriptRecord, SampleRecord, RegionRecord, GenomicBlock + :member-order: bysource + +.. automodule:: transcriptml.rbpnet.windows + :members: WindowScanConfig, generate_window_bounds, calculate_gc_fraction, summarize_regions, scan_windows + :member-order: bysource + +.. automodule:: transcriptml.rbpnet.selection + :members: SelectionConfig, SelectionManifest, select_regions, load_selection_manifest + :member-order: bysource + +.. automodule:: transcriptml.rbpnet.bundle + :members: RBPNetBundleConfig, make_rbpnet_bundle, load_rbpnet_bundle + :member-order: bysource + Models ------ diff --git a/docs/conf.py b/docs/conf.py index 72ffa46..3a493dd 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -47,7 +47,7 @@ # Useful autodoc behavior autosummary_generate = True -autodoc_mock_imports = ["torch", "typing_extensions"] +autodoc_mock_imports = ["torch", "typing_extensions", "h5py", "pyarrow", "pysam", "scipy"] autodoc_typehints = "description" napoleon_google_docstring = True napoleon_numpy_docstring = True diff --git a/docs/index.rst b/docs/index.rst index ece506b..c5342bb 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -15,8 +15,9 @@ motif ablations, motif context scans, motif epistasis analyses, and Saluki-specific codon ISM. These analyses can expose learned regulatory sequence features as well as technical artifacts in the model or assay. -RiboNN support for translation measurements and RBPNet support for RBP binding -measurements are planned but not yet implemented. +RBPNet/eCLIP data preprocessing, descriptive scanning, region selection, and +materialized dataset construction are supported. RBPNet model training remains +planned. Start here ---------- @@ -33,5 +34,6 @@ guide describes every Saluki, LegNet, and shared training option. Use the installation usage + rbpnet training_configuration api diff --git a/docs/installation.md b/docs/installation.md index efb1533..0e28cf8 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -44,6 +44,7 @@ Optional extras are available for a few heavier workflows: | --- | --- | | Write codon-ISM tables as Parquet or Arrow | `python -m pip install -e ".[arrow]"` | | Summarize and plot codon-ISM tables | `python -m pip install -e ".[analysis]"` | +| Preprocess eCLIP and build RBPNet datasets | `python -m pip install -e ".[rbpnet]"` | | Run the test suite | `python -m pip install -e ".[dev]"` | Extras can be combined. The `analysis` extra already includes `pyarrow`, so you @@ -51,5 +52,5 @@ do not need to install both `analysis` and `arrow`. A full analysis and development install is: ```bash -python -m pip install -e ".[analysis,dev]" +python -m pip install -e ".[analysis,rbpnet,dev]" ``` diff --git a/docs/rbpnet.md b/docs/rbpnet.md new file mode 100644 index 0000000..eb01ef6 --- /dev/null +++ b/docs/rbpnet.md @@ -0,0 +1,303 @@ +# RBPNet/eCLIP data workflow + +TranscriptML includes the complete data path needed before implementing an +RBPNet model. It starts from ordinary eCLIP alignments and ends with fixed-shape, +memory-mappable NumPy arrays. It does **not** implement an RBPNet architecture, +loss, trainer, peak caller intended for general use, GC matching, or training +example sampling. + +```text +FASTA + one-transcript-per-gene GTF + IP BAM(s) + SMInput BAM + | + v + preprocess + | + v + canonical transcript-space HDF5 + FASTA + metadata + | + v + descriptive window scan + | + v + region selection + | + v + versioned selection manifest + | + v + RBPNet DatasetBundle + | + v + fixed-shape, memory-mappable .npy arrays +``` + +Install the optional assay dependencies with: + +```bash +python -m pip install -e ".[rbpnet]" +``` + +## 1. Canonical preprocessing + +The GTF must select one transcript per gene. GTF coordinates are converted to +zero-based, half-open intervals. Exons are spliced in transcript 5′→3′ order; +minus-strand exon sequence is reverse-complemented and minus-strand transcript +position zero is therefore the highest-genomic-coordinate mature RNA base. +Protein-coding transcripts are partitioned into `5putr`, `cds`, and `3putr`. +Transcripts without a CDS are `noncoding_exon` throughout. + +```bash +transcriptml rbpnet preprocess \ + --genome-fasta ../RBPNet2/Data/chr21_test/chr21.fa \ + --gtf ../RBPNet2/Data/chr21_test/gencode_v50_MANE_select_chr21.gtf \ + --ip-bam ip1=../RBPNet2/Data/chr21_test/ip1_chr21.bam \ + --ip-bam ip2=../RBPNet2/Data/chr21_test/ip2_chr21.bam \ + --sminput-bam sminput=../RBPNet2/Data/chr21_test/sminput_chr21.bam \ + --output-dir processed/chr21 +``` + +Only read1 alignments are considered. The crosslink-position signal is the +aligned 5′ reference base: `reference_start` for forward alignments and +`reference_end - 1` for reverse alignments. In the eCLIP libraries used during +development, read1 aligns opposite the RNA strand, so `--read1-rna-strand +opposite` is the default. Use `same` or `unstranded` for libraries with another +convention. This read1 convention describes the supplied BAMs and should not be +confused with papers that name the crosslink-bearing FASTQ mate R2 before BAM +construction. + +Unmapped, secondary, supplementary, QC-failed, duplicate (by default), and +low-MAPQ alignments are filtered. An event is assigned only when its 5′ base +maps to one strand-compatible selected transcript and the complete alignment is +compatible with that mature transcript. Aligned/deleted reference segments +must remain in exons and each CIGAR `N` must exactly match a selected adjacent +exon junction. Intronic/pre-mRNA alignments are reported separately as +`transcript_incompatible` in `qc.json`. + +Missing `.fai`/`.bai` indexes are created when the source files and their +directories are writable. GTF transcripts on contigs absent from the analysis +FASTA are reported and skipped; retained FASTA contigs must be represented in +every BAM. + +The canonical directory contains: + +| File | Contract | +| --- | --- | +| `manifest.json` | Format version, exact sample order/roles, input file stat records, configuration, and effective library sizes | +| `qc.json` | Annotation, sequence, read-filter, assignment, and per-sample totals | +| `transcripts.tsv` | Gene/transcript metadata, length, strand, raw counts, SMInput TPM, and compact region annotations | +| `exons.tsv.gz` | Transcript interval ↔ genomic exon block mappings | +| `regions.tsv.gz` | Long-form transcript region intervals | +| `transcripts.fa` + `.fai` | Indexed mature-transcript sequences | +| `signals.h5` | Canonical base-resolution retained read1 5′ counts | + +SMInput TPM is calculated from retained transcript counts divided by mature +transcript length, followed by normalization of those rates to one million. +Each sample's `effective_library_size` is exactly the number of retained read1 +5′ events used to construct its HDF5 track. That field is the CPM denominator. +Pooled-IP CPM uses the sum of IP counts divided by the sum of IP effective +library sizes. + +### HDF5 signal layout + +HDF5 is a hierarchical binary container: datasets behave like typed, +multidimensional arrays stored inside a file, can be compressed and chunked, +and can be sliced without loading the entire array. `signals.h5` uses a compact +concatenated-transcript layout: + +| Dataset | Shape/dtype | Meaning | +| --- | --- | --- | +| `counts` | `(S, total_transcript_bases)`, `uint32` | One row per manifest sample | +| `ip_pooled` | `(total_transcript_bases,)`, `uint32` | Chunkwise sum of all IP rows | +| `sample_names` / `sample_roles` | `(S,)`, UTF-8 | HDF5 row identity | +| `transcript_ids` | `(T,)`, UTF-8 | Transcript order | +| `transcript_offsets` / `transcript_lengths` | `(T,)`, `int64` | Slice boundaries in concatenated space | + +Users normally do not need to calculate flat offsets. The lazy reader handles +FASTA, HDF5, sample order, and exon blocks: + +```python +from transcriptml.rbpnet import ProcessedECLIPDataset + +with ProcessedECLIPDataset("processed/chr21") as ds: + print(ds.transcripts) + print(ds.samples) + seq = ds.get_sequence("ENST...", 100, 400) + input_profile = ds.get_profile("ENST...", 100, 400, sample="sminput") + ip1_profile = ds.get_profile("ENST...", 100, 400, sample="ip1") + pooled = ds.get_pooled_ip_profile("ENST...", 100, 400) + blocks = ds.get_genomic_blocks("ENST...", 100, 400) +``` + +## 2. Descriptive window scanning + +The scanner reads the canonical experiment, not BAMs. A new window size or +stride therefore requires only a new scan: + +```bash +transcriptml rbpnet scan-windows \ + --processed-dir processed/chr21 \ + --window-size 100 \ + --stride 50 \ + --min-sminput-tpm 0 \ + --pseudocount 1 \ + --output-prefix processed/chr21_windows_100nt +``` + +It writes equivalent `*.tsv.gz` and `*.parquet` tables plus `*.scan.json`. +Incomplete terminal windows are omitted unless +`--include-incomplete-terminal-windows` is given. Per-transcript cumulative +sums make count aggregation efficient even for stride-1 scans. + +Rows include transcript/genomic coordinates and exon blocks; exact overlap +counts/fractions for every region class; GC; SMInput TPM; dynamic per-sample +counts, CPM, and maximum positional counts; pooled-IP values; combined +coverage; and a CPM-scale pooled-IP/SMInput log ratio. Boundary-crossing +windows are `mixed`, never silently assigned to one region. + +The scanner is deliberately descriptive. It does not label peaks, negatives, +or training examples. + +## 3. Region selection + +Selection asks which experimental loci are eligible and why. It writes a +versioned `*.parquet` manifest, equivalent `*.tsv.gz`, and a +`*.selection.json` provenance sidecar. + +### Published RBPNet v1 + +First make the published 100-nt, stride-1 descriptive scan, then select: + +```bash +transcriptml rbpnet scan-windows \ + --processed-dir processed/chr21 --window-size 100 --stride 1 \ + --output-prefix processed/chr21_windows_v1 + +transcriptml rbpnet select-regions \ + --processed-dir processed/chr21 \ + --windows processed/chr21_windows_v1.parquet \ + --strategy original_rbpnet \ + --output-prefix processed/chr21_original +``` + +The default preset follows [Horlacher et al. 2023](https://doi.org/10.1186/s13059-023-03015-7) +and its [reference implementation](https://github.com/mhorlacher/rbpnet): a one-sided Poisson test +against the transcript-level pooled-IP rate (`p < 0.01`), at least 8 pooled +counts, a maximum positional count of at least 2, and a 50-nt advance after an +accepted candidate. The selected interval remains 100 nt; a later 300-nt +bundle context is independent. All thresholds are CLI-configurable. + +### Broad measured windows (`yeo_2026`) + +This strategy applies coverage thresholds without a peak test: + +```bash +transcriptml rbpnet select-regions \ + --processed-dir processed/chr21 \ + --windows processed/chr21_windows_100nt.parquet \ + --strategy yeo_2026 \ + --min-total-count 8 --min-sminput-count 1 --min-ip-count 1 \ + --output-prefix processed/chr21_measured +``` + +`--replicate-mode combined` retains one row while preserving every replicate +column. `per_ip` emits a replicate-identified row for each IP satisfying the +thresholds. This is a configurable implementation of the broad inclusion +philosophy; its defaults are not an assertion that one coverage cutoff is +universally optimal. + +### Peak / gray / confident negative + +```bash +transcriptml rbpnet select-regions \ + --processed-dir processed/chr21 \ + --windows processed/chr21_windows_100nt.parquet \ + --strategy peak_gray_negative \ + --min-total-count 8 --min-sminput-count 1 --min-ip-count 1 \ + --peak-fdr 0.05 --peak-min-log2-ratio 1 \ + --negative-fdr 0.05 --negative-max-log2-ratio -0.5 \ + --stitch-gap 0 \ + --output-prefix processed/chr21_peak_gray_negative +``` + +Adequately measured windows are tested by conditioning on pooled-IP + SMInput +counts. The null IP probability is determined by their effective library +sizes. One-sided exact binomial enrichment/depletion p-values are +Benjamini–Hochberg corrected. Peaks require enrichment plus a minimum log2 +effect; confident negatives require depletion plus a maximum log2 effect; the +remaining adequate windows are gray. Low-information windows are omitted. +Overlapping/nearby windows are stitched only when transcript, state, and region +type agree. Peak anchors are pooled-signal maxima; negative and gray anchors +are interval midpoints. These defaults are transparent starting choices, not a +definitive CLIP peak caller. + +### Selection manifest v1 + +Every row has a deterministic content-derived `example_id`; gene, transcript, +chromosome, strand, anchor, and half-open selection interval; region overlap; +strategy/state and optional replicate identity; sample/pooled signal summaries; +statistical fields; and explicit gene/transcript/chromosome grouping columns. +Parquet metadata and the sidecar preserve the complete selection and scan +configuration. Consumers must key by `example_id`, not row order. + +## 4. Materialized RBPNet bundle + +Bundle construction separates the selected biological interval from future +model context: + +```bash +transcriptml rbpnet make-bundle \ + --processed-dir processed/chr21 \ + --selection-manifest processed/chr21_measured.parquet \ + --output-dir data/rbpnet_chr21 \ + --input-length 300 \ + --profile-length 300 \ + --max-jitter 32 \ + --transcript-end-policy pad +``` + +For requested input length `L`, profile length `P`, and maximum future jitter +`J`, the bundle writes: + +| File | Shape/dtype | +| --- | --- | +| `X.npy` | `(N, 4, L + 2J)`, `uint8`, TranscriptML RNA4 A/C/G/U encoding | +| `sminput_profiles.npy` | `(N, P + 2J)`, `uint32` | +| `ip_profiles.npy` | `(N, R, P + 2J)`, `uint32`; `R` follows `ip_axis_order` | +| `sequence_valid_mask.npy` | `(N, L + 2J)`, `uint8` | +| `profile_valid_mask.npy` | `(N, P + 2J)`, `uint8` | +| `profile_sminput_totals.npy` | `(N,)`, `uint64` | +| `profile_ip_totals.npy` | `(N, R)`, `uint64` | +| `selection_sminput_counts.npy` | `(N,)`, `uint64` | +| `selection_ip_counts.npy` | `(N, R)`, `uint64` | + +Pooled IP is `ip_profiles.sum(axis=1)` and is not the only stored +representation. The ordinary TranscriptML sidecars (`ids.txt`, +`metadata.json`, `schema.json`, `config.json`) accompany the arrays, and the +full selected-example metadata is copied to `examples.parquet`. All arrays can +be loaded with NumPy `mmap_mode="r"`. + +With `--transcript-end-policy drop`, examples whose full materialized sequence +or profile context crosses a transcript end are removed and counted. With +`pad`, fixed widths are preserved using all-zero sequence/profile padding and +the validity masks distinguish padding from valid ambiguous sequence or true +zero signal. + +A future jitter shift `s` in `[-J, +J]` takes each requested crop starting at +`J + s` in its materialized array. Bundle construction does not perform random +augmentation. Original v1 can use `J=0`; a jitter-ready workflow can use +`J=32`. Candidate-scan advance and training-time jitter are unrelated. + +```python +from transcriptml.rbpnet.bundle import load_rbpnet_bundle + +bundle = load_rbpnet_bundle("data/rbpnet_chr21", mmap_mode="r") +X = bundle.X +ip = bundle.arrays["ip_profiles"] +input_profile = bundle.arrays["sminput_profiles"] +print(bundle.config["sample_metadata"]["ip_axis_order"]) +``` + +The canonical experiment uses HDF5 because it provides compressed lazy slicing +over an entire transcriptome. The model bundle uses separate `.npy` files +because its selected fixed-shape arrays are simple to inspect and memory-map. +Changing selection, context, or jitter does not require reprocessing BAMs. diff --git a/docs/usage.md b/docs/usage.md index 22e5e79..ca1ce4f 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -5,6 +5,10 @@ This page walks through the two main TranscriptML workflows: - **Saluki**, for transcriptome-derived RNA stability measurements. - **MPRA-LegNet**, for single-insert MPRA-style measurements. +The assay-aware RBPNet/eCLIP preprocessing, window scanning, selection, and +bundle workflow has its own [RBPNet data guide](rbpnet.md). RBPNet model +training is intentionally not implemented yet. + For each workflow, the basic pattern is the same: 1. Build a TranscriptML dataset bundle. diff --git a/pyproject.toml b/pyproject.toml index 2e185a0..69e99f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,11 @@ dependencies = [ [project.optional-dependencies] dev = [ - "pytest>=7" + "pytest>=7", + "h5py>=3.10", + "pyarrow>=12", + "pysam>=0.22", + "scipy>=1.10" ] arrow = [ "pyarrow>=12" @@ -32,6 +36,12 @@ analysis = [ "polars>=0.20", "seaborn>=0.13" ] +rbpnet = [ + "h5py>=3.10", + "pyarrow>=12", + "pysam>=0.22", + "scipy>=1.10" +] [project.scripts] transcriptml = "transcriptml.cli.main:main" diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index 1c27596..5264e20 100644 --- a/src/transcriptml/cli/main.py +++ b/src/transcriptml/cli/main.py @@ -90,6 +90,10 @@ def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="transcriptml") sub = parser.add_subparsers(dest="command", required=True) + from transcriptml.rbpnet.cli import add_rbpnet_parser + + add_rbpnet_parser(sub) + p = sub.add_parser("init-run", help="Write starter run configuration files") p.add_argument("--workflow", required=True, choices=["saluki", "legnet"]) p.add_argument("--out-dir", required=True) @@ -354,6 +358,11 @@ def main(argv: list[str] | None = None) -> None: parser = build_parser() args = parser.parse_args(argv) + if args.command == "rbpnet": + from transcriptml.rbpnet.cli import run_rbpnet_command + + run_rbpnet_command(args, parser) + return if args.command == "init-run": from transcriptml.workflows import init_run diff --git a/src/transcriptml/data/__init__.py b/src/transcriptml/data/__init__.py index a4d74c5..7476189 100644 --- a/src/transcriptml/data/__init__.py +++ b/src/transcriptml/data/__init__.py @@ -1,54 +1,46 @@ -"""Data processing utilities.""" +"""Data processing utilities, exposed lazily to keep assay CLIs lightweight.""" -from transcriptml.data.bundle import DatasetBundle, load_bundle, save_bundle -from transcriptml.data.controls import ( - SequenceControlConfig, - SequenceControlOperation, - apply_sequence_controls_array, - apply_sequence_controls_to_bundle, - normalize_sequence_control_config, -) -from transcriptml.data.encoding import ( - DEFAULT_SALUKI_LENGTH, - encode_rna_sequence, - encode_saluki_transcript, - encode_sequences, - infer_valid_length, - infer_valid_lengths, -) -from transcriptml.data.genomics import ( - TranscriptFeature, - TranscriptRecord, - extract_transcript_records, - iter_gtf_records, - load_transcript_features, - parse_gtf_attributes, -) -from transcriptml.data.schemas import RNA4, SALUKI6, SequenceSchema, get_schema +from __future__ import annotations -__all__ = [ - "DatasetBundle", - "DEFAULT_SALUKI_LENGTH", - "RNA4", - "SALUKI6", - "SequenceSchema", - "SequenceControlConfig", - "SequenceControlOperation", - "TranscriptFeature", - "TranscriptRecord", - "apply_sequence_controls_array", - "apply_sequence_controls_to_bundle", - "encode_rna_sequence", - "encode_saluki_transcript", - "encode_sequences", - "extract_transcript_records", - "get_schema", - "infer_valid_length", - "infer_valid_lengths", - "iter_gtf_records", - "load_bundle", - "load_transcript_features", - "normalize_sequence_control_config", - "parse_gtf_attributes", - "save_bundle", -] +from importlib import import_module + +_EXPORTS = { + "DatasetBundle": ("transcriptml.data.bundle", "DatasetBundle"), + "load_bundle": ("transcriptml.data.bundle", "load_bundle"), + "save_bundle": ("transcriptml.data.bundle", "save_bundle"), + "SequenceControlConfig": ("transcriptml.data.controls", "SequenceControlConfig"), + "SequenceControlOperation": ("transcriptml.data.controls", "SequenceControlOperation"), + "apply_sequence_controls_array": ("transcriptml.data.controls", "apply_sequence_controls_array"), + "apply_sequence_controls_to_bundle": ("transcriptml.data.controls", "apply_sequence_controls_to_bundle"), + "normalize_sequence_control_config": ("transcriptml.data.controls", "normalize_sequence_control_config"), + "DEFAULT_SALUKI_LENGTH": ("transcriptml.data.encoding", "DEFAULT_SALUKI_LENGTH"), + "encode_rna_sequence": ("transcriptml.data.encoding", "encode_rna_sequence"), + "encode_saluki_transcript": ("transcriptml.data.encoding", "encode_saluki_transcript"), + "encode_sequences": ("transcriptml.data.encoding", "encode_sequences"), + "infer_valid_length": ("transcriptml.data.encoding", "infer_valid_length"), + "infer_valid_lengths": ("transcriptml.data.encoding", "infer_valid_lengths"), + "TranscriptFeature": ("transcriptml.data.genomics", "TranscriptFeature"), + "TranscriptRecord": ("transcriptml.data.genomics", "TranscriptRecord"), + "extract_transcript_records": ("transcriptml.data.genomics", "extract_transcript_records"), + "iter_gtf_records": ("transcriptml.data.genomics", "iter_gtf_records"), + "load_transcript_features": ("transcriptml.data.genomics", "load_transcript_features"), + "parse_gtf_attributes": ("transcriptml.data.genomics", "parse_gtf_attributes"), + "RNA4": ("transcriptml.data.schemas", "RNA4"), + "SALUKI6": ("transcriptml.data.schemas", "SALUKI6"), + "SequenceSchema": ("transcriptml.data.schemas", "SequenceSchema"), + "get_schema": ("transcriptml.data.schemas", "get_schema"), +} + +__all__ = sorted(_EXPORTS) + + +def __getattr__(name: str): + """Import a requested data symbol without importing unrelated Torch code.""" + + try: + module_name, attribute = _EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module 'transcriptml.data' has no attribute {name!r}") from exc + value = getattr(import_module(module_name), attribute) + globals()[name] = value + return value diff --git a/src/transcriptml/data/bundle.py b/src/transcriptml/data/bundle.py index cfd2f11..aaaec9e 100644 --- a/src/transcriptml/data/bundle.py +++ b/src/transcriptml/data/bundle.py @@ -12,7 +12,13 @@ @dataclass class DatasetBundle: - """Self-describing processed dataset.""" + """Self-describing processed dataset. + + ``X`` and optional ``y`` retain TranscriptML's original compact contract. + Workflows with additional aligned targets (for example, RBPNet count + profiles) may use ``arrays``. Every named array must share ``X``'s first + dimension and is serialized as an ordinary ``.npy`` file. + """ X: np.ndarray y: np.ndarray | None = None @@ -21,6 +27,7 @@ class DatasetBundle: metadata: Sequence[Mapping[str, Any]] | None = None splits: Mapping[str, Sequence[int]] | None = None config: Mapping[str, Any] = field(default_factory=dict) + arrays: Mapping[str, np.ndarray] = field(default_factory=dict) def __post_init__(self) -> None: """Normalize schema and validate array-aligned fields.""" @@ -32,6 +39,14 @@ def __post_init__(self) -> None: raise ValueError("ids length must match X.shape[0]") if self.y is not None and int(self.y.shape[0]) != int(self.X.shape[0]): raise ValueError("y length must match X.shape[0]") + if self.metadata is not None and len(self.metadata) != int(self.X.shape[0]): + raise ValueError("metadata length must match X.shape[0]") + reserved = {"X", "y"} + for name, array in self.arrays.items(): + if not name or name in reserved or not name.replace("_", "").isalnum(): + raise ValueError(f"invalid named array key: {name!r}") + if int(array.shape[0]) != int(self.X.shape[0]): + raise ValueError(f"named array {name!r} length must match X.shape[0]") def _json_default(obj: Any) -> Any: @@ -78,6 +93,15 @@ def save_bundle_metadata(bundle: DatasetBundle, out_dir: str | Path) -> None: config = dict(bundle.config) config.setdefault("n_examples", int(bundle.X.shape[0])) config.setdefault("shape", [int(x) for x in bundle.X.shape]) + if bundle.arrays: + config["named_arrays"] = { + name: { + "file": f"{name}.npy", + "shape": [int(x) for x in array.shape], + "dtype": str(array.dtype), + } + for name, array in bundle.arrays.items() + } (out / "config.json").write_text(json.dumps(config, indent=2, default=_json_default), encoding="utf-8") @@ -95,6 +119,8 @@ def save_bundle(bundle: DatasetBundle, out_dir: str | Path) -> None: np.save(out / "X.npy", bundle.X) if bundle.y is not None: np.save(out / "y.npy", bundle.y) + for name, array in bundle.arrays.items(): + np.save(out / f"{name}.npy", array) save_bundle_metadata(bundle, out) @@ -118,4 +144,27 @@ def load_bundle(path: str | Path, *, mmap_mode: str | None = None) -> DatasetBun splits = json.loads(splits_path.read_text(encoding="utf-8")) if splits_path.exists() else None config_path = root / "config.json" config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.exists() else {} - return DatasetBundle(X=X, y=y, ids=ids, schema=schema, metadata=metadata, splits=splits, config=config) + arrays = {} + for name, spec in config.get("named_arrays", {}).items(): + filename = Path(spec["file"]) + if filename.is_absolute() or len(filename.parts) != 1 or filename.suffix != ".npy": + raise ValueError(f"invalid named array file for {name!r}: {filename}") + array = np.load(root / filename, mmap_mode=mmap_mode) + expected_shape = tuple(int(value) for value in spec.get("shape", array.shape)) + expected_dtype = np.dtype(spec.get("dtype", array.dtype)) + if array.shape != expected_shape or array.dtype != expected_dtype: + raise ValueError( + f"named array {name!r} does not match config metadata: " + f"found {array.shape}/{array.dtype}, expected {expected_shape}/{expected_dtype}" + ) + arrays[name] = array + return DatasetBundle( + X=X, + y=y, + ids=ids, + schema=schema, + metadata=metadata, + splits=splits, + config=config, + arrays=arrays, + ) diff --git a/src/transcriptml/rbpnet/__init__.py b/src/transcriptml/rbpnet/__init__.py new file mode 100644 index 0000000..390c299 --- /dev/null +++ b/src/transcriptml/rbpnet/__init__.py @@ -0,0 +1,44 @@ +"""RBPNet/eCLIP transcript-space data preparation. + +The public API deliberately stops at model-ready data. Neural-network +architectures, losses, and training are not part of this module. +""" + +__all__ = [ + "PipelineConfig", + "ProcessedECLIPDataset", + "RBPNetBundleConfig", + "Sample", + "SelectionConfig", + "WindowScanConfig", + "make_rbpnet_bundle", + "preprocess_eclip", + "scan_windows", + "select_regions", +] + + +def __getattr__(name: str): + """Lazily expose APIs so base TranscriptML installs can still show CLI help.""" + + if name in {"PipelineConfig", "Sample", "preprocess_eclip"}: + from transcriptml.rbpnet.preprocessing import PipelineConfig, Sample, preprocess_eclip + + return {"PipelineConfig": PipelineConfig, "Sample": Sample, "preprocess_eclip": preprocess_eclip}[name] + if name == "ProcessedECLIPDataset": + from transcriptml.rbpnet.experiment import ProcessedECLIPDataset + + return ProcessedECLIPDataset + if name in {"WindowScanConfig", "scan_windows"}: + from transcriptml.rbpnet.windows import WindowScanConfig, scan_windows + + return {"WindowScanConfig": WindowScanConfig, "scan_windows": scan_windows}[name] + if name in {"SelectionConfig", "select_regions"}: + from transcriptml.rbpnet.selection import SelectionConfig, select_regions + + return {"SelectionConfig": SelectionConfig, "select_regions": select_regions}[name] + if name in {"RBPNetBundleConfig", "make_rbpnet_bundle"}: + from transcriptml.rbpnet.bundle import RBPNetBundleConfig, make_rbpnet_bundle + + return {"RBPNetBundleConfig": RBPNetBundleConfig, "make_rbpnet_bundle": make_rbpnet_bundle}[name] + raise AttributeError(f"module 'transcriptml.rbpnet' has no attribute {name!r}") diff --git a/src/transcriptml/rbpnet/_progress.py b/src/transcriptml/rbpnet/_progress.py new file mode 100644 index 0000000..17496dd --- /dev/null +++ b/src/transcriptml/rbpnet/_progress.py @@ -0,0 +1,32 @@ +"""Progress helpers built on TranscriptML's native reporter.""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from typing import TypeVar + +from transcriptml.progress import ProgressReporter, log_progress + +T = TypeVar("T") + + +def track( + iterable: Iterable[T], + label: str, + *, + total: int | None = None, + unit: str = "items", + enabled: bool = True, +) -> Iterator[T]: + """Yield an iterable while emitting throttled TranscriptML progress.""" + + reporter = ProgressReporter(label, total=total, unit=unit, enabled=enabled) + try: + for value in iterable: + yield value + reporter.update() + finally: + reporter.close() + + +__all__ = ["ProgressReporter", "log_progress", "track"] diff --git a/src/transcriptml/rbpnet/annotation.py b/src/transcriptml/rbpnet/annotation.py new file mode 100644 index 0000000..37ad1ee --- /dev/null +++ b/src/transcriptml/rbpnet/annotation.py @@ -0,0 +1,101 @@ +"""Strict parsing of one-transcript-per-gene GENCODE-style GTF files.""" + +from __future__ import annotations + +import gzip +import re +from collections import Counter +from pathlib import Path + +from transcriptml.rbpnet._progress import track +from transcriptml.rbpnet.coordinates import Exon, Transcript, annotate_regions + +_ATTR_RE = re.compile(r'([^\s;]+)\s+(?:"([^"]*)"|([^;\s]+))') + + +def parse_attributes(text: str) -> dict[str, str]: + """Parse GTF attributes into a string dictionary.""" + + return {key: quoted or bare for key, quoted, bare in _ATTR_RE.findall(text)} + + +def _open_text(path: Path): + return gzip.open(path, "rt") if path.suffix == ".gz" else path.open(encoding="utf-8") + + +def parse_gtf(path: str | Path, *, progress: bool = True) -> list[Transcript]: + """Parse and validate a one-transcript-per-gene annotation.""" + + path = Path(path) + if not path.is_file(): + raise FileNotFoundError(f"GTF not found: {path}") + transcripts: dict[str, Transcript] = {} + exon_rows: dict[str, list[Exon]] = {} + feature_rows: dict[str, dict[str, list[tuple[int, int]]]] = {} + with _open_text(path) as handle: + lines = track(handle, "rbpnet preprocess: parse GTF", unit="lines", enabled=progress) + for line_no, line in enumerate(lines, 1): + if not line.strip() or line.startswith("#"): + continue + fields = line.rstrip("\n").split("\t") + if len(fields) != 9: + raise ValueError(f"{path}:{line_no}: expected 9 tab-separated GTF fields") + chrom, _, feature, start_s, end_s, _, strand, _, attrs_s = fields + attrs = parse_attributes(attrs_s) + tx_id = attrs.get("transcript_id") + if not tx_id: + continue + try: + start, end = int(start_s) - 1, int(end_s) + except ValueError as exc: + raise ValueError(f"{path}:{line_no}: invalid coordinates") from exc + if feature == "transcript": + if tx_id in transcripts: + raise ValueError(f"duplicate transcript row for {tx_id}") + transcripts[tx_id] = Transcript( + transcript_id=tx_id, + gene_id=attrs.get("gene_id", ""), + gene_name=attrs.get("gene_name", ""), + transcript_name=attrs.get("transcript_name", ""), + transcript_type=attrs.get("transcript_type", attrs.get("gene_type", "")), + chrom=chrom, + strand=strand, + ) + elif feature == "exon": + exon_rows.setdefault(tx_id, []).append( + Exon( + chrom, + start, + end, + strand, + attrs.get("exon_number", ""), + attrs.get("exon_id", ""), + ) + ) + elif feature in {"CDS", "UTR", "start_codon", "stop_codon"}: + feature_rows.setdefault(tx_id, {}).setdefault(feature, []).append((start, end)) + if not transcripts: + raise ValueError(f"no transcript records found in {path}") + orphan_exons = set(exon_rows) - set(transcripts) + if orphan_exons: + raise ValueError(f"exons found without transcript row (example: {sorted(orphan_exons)[0]})") + gene_counts = Counter(tx.gene_id for tx in transcripts.values() if tx.gene_id) + duplicates = {gene_id for gene_id, count in gene_counts.items() if count > 1} + if duplicates: + raise ValueError( + "annotation is not one-transcript-per-gene; multiple transcript rows found for " + + sorted(duplicates)[0] + ) + result: list[Transcript] = [] + for tx in transcripts.values(): + tx.exons = exon_rows.get(tx.transcript_id, []) + tx.feature_intervals = feature_rows.get(tx.transcript_id, {}) + tx.finalize() + tx.regions = annotate_regions(tx) + result.append(tx) + result.sort(key=lambda tx: (tx.chrom, min(e.start for e in tx.exons), tx.transcript_id)) + offset = 0 + for tx in result: + tx.offset = offset + offset += tx.length + return result diff --git a/src/transcriptml/rbpnet/bundle.py b/src/transcriptml/rbpnet/bundle.py new file mode 100644 index 0000000..7034fa7 --- /dev/null +++ b/src/transcriptml/rbpnet/bundle.py @@ -0,0 +1,373 @@ +"""Materialize selected eCLIP loci as a TranscriptML RBPNet array bundle.""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + +from transcriptml.data.bundle import DatasetBundle, load_bundle, save_bundle_metadata +from transcriptml.data.encoding import encode_rna_sequence +from transcriptml.data.schemas import RNA4 +from transcriptml.progress import ProgressReporter, log_progress +from transcriptml.rbpnet.experiment import ProcessedECLIPDataset +from transcriptml.rbpnet.selection import SelectionManifest, load_selection_manifest + + +@dataclass(frozen=True) +class RBPNetBundleConfig: + """Configuration for fixed-shape RBPNet bundle materialization.""" + + processed_dir: Path + selection_manifest: Path + output_dir: Path + input_length: int = 300 + profile_length: int = 300 + max_jitter: int = 0 + transcript_end_policy: str = "drop" + overwrite: bool = False + progress: bool = True + + +def _materialized_interval(anchor: int, length: int, jitter: int) -> tuple[int, int]: + width = length + 2 * jitter + start = anchor - width // 2 + return start, start + width + + +def _source_and_destination(start: int, end: int, transcript_length: int) -> tuple[int, int, int, int]: + source_start = max(0, start) + source_end = min(transcript_length, end) + destination_start = source_start - start + destination_end = destination_start + max(0, source_end - source_start) + return source_start, source_end, destination_start, destination_end + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _prepare_output(path: Path, overwrite: bool) -> None: + path.mkdir(parents=True, exist_ok=True) + known = { + "X.npy", "sminput_profiles.npy", "ip_profiles.npy", + "sequence_valid_mask.npy", "profile_valid_mask.npy", + "profile_sminput_totals.npy", "profile_ip_totals.npy", + "selection_sminput_counts.npy", "selection_ip_counts.npy", + "ids.txt", "metadata.json", "schema.json", "config.json", "examples.parquet", + } + existing = [path / name for name in known if (path / name).exists()] + if existing and not overwrite: + raise FileExistsError(f"bundle output already exists ({existing[0]}); pass --overwrite") + for item in existing: + item.unlink() + + +def _validate_config(config: RBPNetBundleConfig) -> None: + if config.input_length <= 0 or config.profile_length <= 0: + raise ValueError("input_length and profile_length must be positive") + if config.max_jitter < 0: + raise ValueError("max_jitter must be non-negative") + if config.transcript_end_policy not in {"drop", "pad"}: + raise ValueError("transcript_end_policy must be drop or pad") + + +def _sorted_rows(manifest: SelectionManifest) -> list[dict]: + order = pc.sort_indices(manifest.table, sort_keys=[("example_id", "ascending")]) + return pc.take(manifest.table, order).to_pylist() + + +def _validate_manifest_dataset( + manifest: SelectionManifest, + ds: ProcessedECLIPDataset, +) -> None: + scan = manifest.metadata.get("window_scan", {}) + observed = { + str(name): int(value) + for name, value in scan.get("effective_library_sizes", {}).items() + } + expected = { + sample.name: int(sample.effective_library_size) + for sample in ds.samples + } + if observed != expected: + raise ValueError( + "selection manifest effective library sizes do not match the processed experiment" + ) + required = { + f"{sample.name}_{suffix}" + for sample in ds.samples + for suffix in ("count", "cpm") + } + missing = sorted(required - set(manifest.table.column_names)) + if missing: + raise ValueError( + f"selection manifest lacks sample columns required by this experiment: {', '.join(missing)}" + ) + + +def make_rbpnet_bundle(config: RBPNetBundleConfig) -> DatasetBundle: + """Materialize selected loci into memory-mappable NumPy arrays. + + ``input_length`` and ``profile_length`` describe future training crops. + The stored widths add ``2 * max_jitter`` so a future loader can choose a + shared positional shift without reopening FASTA or HDF5 files. + """ + + _validate_config(config) + manifest = load_selection_manifest(config.selection_manifest) + rows = _sorted_rows(manifest) + sequence_width = config.input_length + 2 * config.max_jitter + profile_width = config.profile_length + 2 * config.max_jitter + _prepare_output(config.output_dir, config.overwrite) + + log_progress( + f"rbpnet make-bundle: validate {len(rows):,} selected examples", + enabled=config.progress, + ) + with ProcessedECLIPDataset(config.processed_dir) as ds: + _validate_manifest_dataset(manifest, ds) + ip_samples = ds.ip_samples + if not ip_samples: + raise ValueError("processed dataset contains no IP samples") + kept: list[dict] = [] + dropped = 0 + for row in rows: + tx = ds.get_transcript(row["transcript_id"]) + anchor = int(row["transcript_anchor"]) + selection_start = int(row["selection_start"]) + selection_end = int(row["selection_end"]) + if selection_start < 0 or selection_end <= selection_start or selection_end > tx.length: + raise ValueError( + f"invalid selection interval {tx.transcript_id}:{selection_start}-{selection_end}; " + f"transcript length is {tx.length}" + ) + if anchor < 0 or anchor >= tx.length: + raise ValueError( + f"selection anchor {anchor} is outside transcript {tx.transcript_id} length {tx.length}" + ) + replicate_id = str(row["replicate_id"]) + if replicate_id and replicate_id not in {sample.name for sample in ip_samples}: + raise ValueError(f"selection manifest has unknown IP replicate_id {replicate_id!r}") + seq_start, seq_end = _materialized_interval(anchor, config.input_length, config.max_jitter) + profile_start, profile_end = _materialized_interval(anchor, config.profile_length, config.max_jitter) + in_bounds = ( + seq_start >= 0 and seq_end <= tx.length + and profile_start >= 0 and profile_end <= tx.length + ) + if config.transcript_end_policy == "drop" and not in_bounds: + dropped += 1 + continue + row = dict(row) + row.update({ + "sequence_context_start": seq_start, + "sequence_context_end": seq_end, + "profile_context_start": profile_start, + "profile_context_end": profile_end, + }) + kept.append(row) + if not kept: + raise ValueError( + "no examples remain after transcript-end handling; use --transcript-end-policy pad " + "or reduce context/jitter lengths" + ) + + n_examples = len(kept) + n_ip = len(ip_samples) + X = np.lib.format.open_memmap( + config.output_dir / "X.npy", mode="w+", dtype=np.uint8, + shape=(n_examples, 4, sequence_width), + ) + sminput_profiles = np.lib.format.open_memmap( + config.output_dir / "sminput_profiles.npy", mode="w+", dtype=np.uint32, + shape=(n_examples, profile_width), + ) + ip_profiles = np.lib.format.open_memmap( + config.output_dir / "ip_profiles.npy", mode="w+", dtype=np.uint32, + shape=(n_examples, n_ip, profile_width), + ) + sequence_valid_mask = np.lib.format.open_memmap( + config.output_dir / "sequence_valid_mask.npy", mode="w+", dtype=np.uint8, + shape=(n_examples, sequence_width), + ) + profile_valid_mask = np.lib.format.open_memmap( + config.output_dir / "profile_valid_mask.npy", mode="w+", dtype=np.uint8, + shape=(n_examples, profile_width), + ) + profile_sminput_totals = np.lib.format.open_memmap( + config.output_dir / "profile_sminput_totals.npy", mode="w+", dtype=np.uint64, + shape=(n_examples,), + ) + profile_ip_totals = np.lib.format.open_memmap( + config.output_dir / "profile_ip_totals.npy", mode="w+", dtype=np.uint64, + shape=(n_examples, n_ip), + ) + selection_sminput_counts = np.lib.format.open_memmap( + config.output_dir / "selection_sminput_counts.npy", mode="w+", dtype=np.uint64, + shape=(n_examples,), + ) + selection_ip_counts = np.lib.format.open_memmap( + config.output_dir / "selection_ip_counts.npy", mode="w+", dtype=np.uint64, + shape=(n_examples, n_ip), + ) + arrays = { + "sminput_profiles": sminput_profiles, + "ip_profiles": ip_profiles, + "sequence_valid_mask": sequence_valid_mask, + "profile_valid_mask": profile_valid_mask, + "profile_sminput_totals": profile_sminput_totals, + "profile_ip_totals": profile_ip_totals, + "selection_sminput_counts": selection_sminput_counts, + "selection_ip_counts": selection_ip_counts, + } + metadata: list[dict] = [] + reporter = ProgressReporter( + "rbpnet make-bundle: materialize examples", + total=n_examples, + unit="examples", + enabled=config.progress, + ) + input_name = ds.sminput_sample.name + for index, row in enumerate(kept): + tx = ds.get_transcript(row["transcript_id"]) + seq_start = int(row["sequence_context_start"]) + seq_end = int(row["sequence_context_end"]) + src_start, src_end, dst_start, dst_end = _source_and_destination( + seq_start, seq_end, tx.length + ) + if src_end > src_start: + encoded = encode_rna_sequence(ds.get_sequence(tx.transcript_id, src_start, src_end)) + X[index, :, dst_start:dst_end] = encoded + sequence_valid_mask[index, dst_start:dst_end] = 1 + + profile_start = int(row["profile_context_start"]) + profile_end = int(row["profile_context_end"]) + psrc_start, psrc_end, pdst_start, pdst_end = _source_and_destination( + profile_start, profile_end, tx.length + ) + if psrc_end > psrc_start: + sminput = ds.get_profile(tx.transcript_id, psrc_start, psrc_end, input_name) + sminput_profiles[index, pdst_start:pdst_end] = sminput + for ip_index, sample in enumerate(ip_samples): + ip_profiles[index, ip_index, pdst_start:pdst_end] = ds.get_profile( + tx.transcript_id, psrc_start, psrc_end, sample.name + ) + profile_valid_mask[index, pdst_start:pdst_end] = 1 + profile_sminput_totals[index] = sminput_profiles[index].sum(dtype=np.uint64) + profile_ip_totals[index] = ip_profiles[index].sum(axis=1, dtype=np.uint64) + selection_sminput_counts[index] = int(row[f"{input_name}_count"]) + selection_ip_counts[index] = np.asarray( + [int(row[f"{sample.name}_count"]) for sample in ip_samples], dtype=np.uint64 + ) + metadata.append({ + "example_id": row["example_id"], + "gene_id": row["gene_id"], + "transcript_id": row["transcript_id"], + "chromosome": row["chromosome"], + "strand": row["strand"], + "transcript_anchor": int(row["transcript_anchor"]), + "selection_start": int(row["selection_start"]), + "selection_end": int(row["selection_end"]), + "region_type": row["region_type"], + "selection_strategy": row["selection_strategy"], + "selection_state": row["selection_state"], + "replicate_id": row["replicate_id"], + "group_gene_id": row["group_gene_id"], + "group_transcript_id": row["group_transcript_id"], + "group_chromosome": row["group_chromosome"], + "sequence_context_start": seq_start, + "sequence_context_end": seq_end, + "sequence_left_pad": max(0, -seq_start), + "sequence_right_pad": max(0, seq_end - tx.length), + "profile_context_start": profile_start, + "profile_context_end": profile_end, + "profile_left_pad": max(0, -profile_start), + "profile_right_pad": max(0, profile_end - tx.length), + }) + reporter.update() + reporter.close() + X.flush() + for array in arrays.values(): + array.flush() + + # Keep a self-contained scalable copy of the selected-example contract, + # sorted in the exact same stable-ID order as the arrays. + example_table = pa.Table.from_pylist(kept, schema=manifest.table.schema) + pq.write_table(example_table, config.output_dir / "examples.parquet", compression="zstd") + config_payload = { + "builder": "rbpnet", + "bundle_format": "transcriptml-rbpnet-bundle", + "bundle_format_version": "1", + "source_processed_dir": str(config.processed_dir.resolve()), + "source_selection_manifest": str(manifest.path.resolve()), + "source_selection_sha256": _sha256(manifest.path), + "selection": manifest.metadata, + "input_length": config.input_length, + "profile_length": config.profile_length, + "max_jitter": config.max_jitter, + "materialized_sequence_length": sequence_width, + "materialized_profile_length": profile_width, + "jitter_contract": ( + "future shift s in [-max_jitter,+max_jitter] takes sequence/profile crop " + "starting at max_jitter+s from their respective materialized arrays" + ), + "transcript_end_policy": config.transcript_end_policy, + "n_selected_manifest_rows": len(rows), + "n_dropped_at_transcript_ends": dropped, + "sample_metadata": { + "sminput": { + "name": ds.sminput_sample.name, + "effective_library_size": ds.sminput_sample.effective_library_size, + }, + "ip": [ + {"name": sample.name, "effective_library_size": sample.effective_library_size} + for sample in ip_samples + ], + "ip_axis_order": [sample.name for sample in ip_samples], + "pooled_ip_definition": "sum ip_profiles across axis 1", + }, + "example_metadata_file": "examples.parquet", + } + bundle = DatasetBundle( + X=X, + y=None, + ids=[row["example_id"] for row in kept], + schema=RNA4, + metadata=metadata, + config=config_payload, + arrays=arrays, + ) + save_bundle_metadata(bundle, config.output_dir) + log_progress( + f"rbpnet make-bundle: wrote {n_examples:,} examples to {config.output_dir}", + enabled=config.progress, + ) + return bundle + + +def load_rbpnet_bundle(path: str | Path, *, mmap_mode: str | None = "r") -> DatasetBundle: + """Load and validate a materialized RBPNet bundle.""" + + bundle = load_bundle(path, mmap_mode=mmap_mode) + if bundle.config.get("bundle_format") != "transcriptml-rbpnet-bundle": + raise ValueError(f"not a TranscriptML RBPNet bundle: {path}") + if str(bundle.config.get("bundle_format_version")) != "1": + raise ValueError("unsupported RBPNet bundle format version") + required = { + "sminput_profiles", "ip_profiles", "sequence_valid_mask", "profile_valid_mask", + "profile_sminput_totals", "profile_ip_totals", + "selection_sminput_counts", "selection_ip_counts", + } + missing = sorted(required - set(bundle.arrays)) + if missing: + raise ValueError(f"RBPNet bundle lacks named arrays: {', '.join(missing)}") + return bundle diff --git a/src/transcriptml/rbpnet/cli.py b/src/transcriptml/rbpnet/cli.py new file mode 100644 index 0000000..7dc16b6 --- /dev/null +++ b/src/transcriptml/rbpnet/cli.py @@ -0,0 +1,206 @@ +"""TranscriptML-native CLI wiring for the RBPNet data workflow.""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +_VALID_SAMPLE = re.compile(r"^[A-Za-z][A-Za-z0-9_.-]*$") + + +def add_rbpnet_parser(subparsers) -> None: + """Add the nested ``transcriptml rbpnet`` command family.""" + + root = subparsers.add_parser("rbpnet", help="Prepare eCLIP data for future RBPNet models") + commands = root.add_subparsers(dest="rbpnet_command", required=True) + + preprocess = commands.add_parser( + "preprocess", help="Build a canonical mature-transcript eCLIP experiment" + ) + preprocess.add_argument("--genome-fasta", required=True, type=Path) + preprocess.add_argument("--gtf", required=True, type=Path, help="one-transcript-per-gene GTF") + preprocess.add_argument( + "--ip-bam", action="append", required=True, metavar="[LABEL=]PATH", + help="IP BAM; repeat once per replicate", + ) + preprocess.add_argument("--sminput-bam", required=True, metavar="[LABEL=]PATH") + preprocess.add_argument("--output-dir", required=True, type=Path) + preprocess.add_argument( + "--read1-rna-strand", choices=("opposite", "same", "unstranded"), default="opposite", + help="RNA strand relative to read1 alignment (default: opposite for eCLIP)", + ) + preprocess.add_argument("--min-mapq", type=int, default=1) + preprocess.add_argument("--include-duplicates", action="store_true") + preprocess.add_argument("--overwrite", action="store_true") + preprocess.add_argument("--no-progress", action="store_true") + + scan = commands.add_parser( + "scan-windows", help="Create descriptive transcript-window TSV and Parquet tables" + ) + scan.add_argument("--processed-dir", required=True, type=Path) + scan.add_argument("--window-size", type=int, default=100) + scan.add_argument("--stride", type=int, default=50) + scan.add_argument("--min-sminput-tpm", type=float, default=0.0) + scan.add_argument("--pseudocount", type=float, default=1.0) + terminal = scan.add_mutually_exclusive_group() + terminal.add_argument( + "--omit-incomplete-terminal-windows", dest="omit_incomplete", action="store_true", + default=True, + ) + terminal.add_argument( + "--include-incomplete-terminal-windows", dest="omit_incomplete", action="store_false", + ) + scan.add_argument("--output-prefix", required=True, type=Path) + scan.add_argument("--overwrite", action="store_true") + scan.add_argument("--no-progress", action="store_true") + + select = commands.add_parser( + "select-regions", help="Select eligible loci and write a lightweight manifest" + ) + select.add_argument("--processed-dir", required=True, type=Path) + select.add_argument("--windows", required=True, type=Path, help="window Parquet path or prefix") + select.add_argument("--output-prefix", required=True, type=Path) + select.add_argument( + "--strategy", required=True, + choices=("original_rbpnet", "yeo_2026", "peak_gray_negative"), + ) + select.add_argument("--original-min-pvalue", type=float, default=0.01) + select.add_argument("--original-min-count", type=int, default=8) + select.add_argument("--original-min-height", type=int, default=2) + select.add_argument("--original-advance", type=int, default=50) + select.add_argument("--min-total-count", type=int, default=8) + select.add_argument("--min-sminput-count", type=int, default=1) + select.add_argument("--min-ip-count", type=int, default=1) + select.add_argument("--min-sminput-tpm", type=float, default=0.0) + select.add_argument("--replicate-mode", choices=("combined", "per_ip"), default="combined") + select.add_argument("--peak-fdr", type=float, default=0.05) + select.add_argument("--peak-min-log2-ratio", type=float, default=1.0) + select.add_argument("--negative-fdr", type=float, default=0.05) + select.add_argument("--negative-max-log2-ratio", type=float, default=-0.5) + select.add_argument("--stitch-gap", type=int, default=0) + select.add_argument("--overwrite", action="store_true") + select.add_argument("--no-progress", action="store_true") + + bundle = commands.add_parser( + "make-bundle", help="Materialize selected loci as NumPy RBPNet arrays" + ) + bundle.add_argument("--processed-dir", required=True, type=Path) + bundle.add_argument("--selection-manifest", required=True, type=Path) + bundle.add_argument("--output-dir", required=True, type=Path) + bundle.add_argument("--input-length", type=int, default=300) + bundle.add_argument("--profile-length", type=int, default=300) + bundle.add_argument("--max-jitter", type=int, default=0) + bundle.add_argument("--transcript-end-policy", choices=("drop", "pad"), default="drop") + bundle.add_argument("--overwrite", action="store_true") + bundle.add_argument("--no-progress", action="store_true") + + +def _sample(value: str, role: str): + from transcriptml.rbpnet.preprocessing import Sample + + if "=" in value: + name, path_text = value.split("=", 1) + else: + path_text = value + name = Path(value).name.removesuffix(".bam") + if not _VALID_SAMPLE.fullmatch(name): + raise ValueError( + f"invalid sample name {name!r}; use LABEL=/path/file.bam with letters/numbers/._-" + ) + return Sample(name=name, path=Path(path_text), role=role) + + +def _install_message() -> str: + return "This command requires the RBPNet extra: pip install 'TranscriptML[rbpnet]'" + + +def run_rbpnet_command(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None: + """Dispatch one parsed RBPNet subcommand with consistent errors.""" + + try: + if args.rbpnet_command == "preprocess": + from transcriptml.rbpnet.preprocessing import PipelineConfig, preprocess_eclip + + qc = preprocess_eclip(PipelineConfig( + genome_fasta=args.genome_fasta, + gtf=args.gtf, + sminput=_sample(args.sminput_bam, "sminput"), + ips=tuple(_sample(value, "ip") for value in args.ip_bam), + output_dir=args.output_dir, + read1_rna_strand=args.read1_rna_strand, + min_mapq=args.min_mapq, + exclude_duplicates=not args.include_duplicates, + overwrite=args.overwrite, + progress=not args.no_progress, + )) + result = { + "output_dir": str(args.output_dir), + "transcripts": qc["annotation"]["transcripts"], + "samples": {name: values["retained"] for name, values in qc["samples"].items()}, + } + elif args.rbpnet_command == "scan-windows": + from transcriptml.rbpnet.windows import WindowScanConfig, scan_windows + + result = scan_windows(WindowScanConfig( + processed_dir=args.processed_dir, + output_prefix=args.output_prefix, + window_size=args.window_size, + stride=args.stride, + min_sminput_tpm=args.min_sminput_tpm, + pseudocount=args.pseudocount, + omit_incomplete_terminal_windows=args.omit_incomplete, + overwrite=args.overwrite, + progress=not args.no_progress, + )) + elif args.rbpnet_command == "select-regions": + from transcriptml.rbpnet.selection import SelectionConfig, select_regions + + result = select_regions(SelectionConfig( + processed_dir=args.processed_dir, + windows=args.windows, + output_prefix=args.output_prefix, + strategy=args.strategy, + original_min_pvalue=args.original_min_pvalue, + original_min_count=args.original_min_count, + original_min_height=args.original_min_height, + original_advance=args.original_advance, + min_total_count=args.min_total_count, + min_sminput_count=args.min_sminput_count, + min_ip_count=args.min_ip_count, + min_sminput_tpm=args.min_sminput_tpm, + replicate_mode=args.replicate_mode, + peak_fdr=args.peak_fdr, + peak_min_log2_ratio=args.peak_min_log2_ratio, + negative_fdr=args.negative_fdr, + negative_max_log2_ratio=args.negative_max_log2_ratio, + stitch_gap=args.stitch_gap, + overwrite=args.overwrite, + progress=not args.no_progress, + )) + else: + from transcriptml.rbpnet.bundle import RBPNetBundleConfig, make_rbpnet_bundle + + built = make_rbpnet_bundle(RBPNetBundleConfig( + processed_dir=args.processed_dir, + selection_manifest=args.selection_manifest, + output_dir=args.output_dir, + input_length=args.input_length, + profile_length=args.profile_length, + max_jitter=args.max_jitter, + transcript_end_policy=args.transcript_end_policy, + overwrite=args.overwrite, + progress=not args.no_progress, + )) + result = { + "output_dir": str(args.output_dir), + "n_examples": len(built.ids), + "X_shape": list(built.X.shape), + "ip_profiles_shape": list(built.arrays["ip_profiles"].shape), + } + except ModuleNotFoundError as exc: + raise SystemExit(f"Missing optional dependency {exc.name!r}. {_install_message()}") from exc + except (FileNotFoundError, FileExistsError, KeyError, IndexError, ValueError, RuntimeError, OSError) as exc: + parser.exit(2, f"error: {exc}\n") + print(json.dumps(result, indent=2, sort_keys=True)) diff --git a/src/transcriptml/rbpnet/coordinates.py b/src/transcriptml/rbpnet/coordinates.py new file mode 100644 index 0000000..39faf06 --- /dev/null +++ b/src/transcriptml/rbpnet/coordinates.py @@ -0,0 +1,148 @@ +"""Spliced transcript coordinate models and interval conversion.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Iterable + + +@dataclass +class Exon: + """Half-open genomic exon with an assigned transcript interval.""" + + chrom: str + start: int + end: int + strand: str + exon_number: str = "" + exon_id: str = "" + tx_start: int = 0 + tx_end: int = 0 + + def __post_init__(self) -> None: + if self.start < 0 or self.end <= self.start: + raise ValueError(f"invalid exon interval {self.chrom}:{self.start}-{self.end}") + if self.strand not in {"+", "-"}: + raise ValueError(f"invalid exon strand: {self.strand!r}") + + +@dataclass(frozen=True) +class Region: + """Half-open transcript interval carrying a biological region label.""" + + start: int + end: int + label: str + + +@dataclass +class Transcript: + """One mature transcript in transcript 5-prime to 3-prime orientation.""" + + transcript_id: str + gene_id: str + gene_name: str + transcript_name: str + transcript_type: str + chrom: str + strand: str + exons: list[Exon] = field(default_factory=list) + feature_intervals: dict[str, list[tuple[int, int]]] = field(default_factory=dict) + regions: list[Region] = field(default_factory=list) + offset: int = 0 + + def finalize(self) -> None: + """Validate exons, order them 5-prime to 3-prime, and assign coordinates.""" + + if not self.exons: + raise ValueError(f"transcript {self.transcript_id} has no exons") + if any(e.chrom != self.chrom or e.strand != self.strand for e in self.exons): + raise ValueError(f"inconsistent chromosome/strand in {self.transcript_id}") + genomic = sorted(self.exons, key=lambda e: (e.start, e.end)) + for left, right in zip(genomic, genomic[1:]): + if left.end > right.start: + raise ValueError(f"overlapping exons in {self.transcript_id}") + ordered = genomic if self.strand == "+" else list(reversed(genomic)) + cursor = 0 + for exon in ordered: + exon.tx_start = cursor + cursor += exon.end - exon.start + exon.tx_end = cursor + self.exons = ordered + + @property + def length(self) -> int: + return sum(e.end - e.start for e in self.exons) + + def genome_to_transcript(self, chrom: str, pos: int) -> int | None: + """Map one zero-based genomic base to transcript space.""" + + if chrom != self.chrom: + return None + for exon in self.exons: + if exon.start <= pos < exon.end: + delta = pos - exon.start if self.strand == "+" else exon.end - 1 - pos + return exon.tx_start + delta + return None + + def transcript_to_genome(self, pos: int) -> tuple[str, int, str]: + """Map one zero-based transcript base to genomic coordinates.""" + + if pos < 0 or pos >= self.length: + raise IndexError(f"transcript position {pos} outside [0,{self.length})") + for exon in self.exons: + if exon.tx_start <= pos < exon.tx_end: + delta = pos - exon.tx_start + genomic = exon.start + delta if self.strand == "+" else exon.end - 1 - delta + return self.chrom, genomic, self.strand + raise AssertionError("finalized transcript has a coordinate gap") + + def genomic_interval_to_transcript(self, start: int, end: int) -> list[tuple[int, int]]: + """Map a half-open genomic interval to covered transcript pieces.""" + + pieces: list[tuple[int, int]] = [] + for exon in self.exons: + lo, hi = max(start, exon.start), min(end, exon.end) + if lo >= hi: + continue + if self.strand == "+": + pieces.append((exon.tx_start + lo - exon.start, exon.tx_start + hi - exon.start)) + else: + pieces.append((exon.tx_start + exon.end - hi, exon.tx_start + exon.end - lo)) + return sorted(pieces) + + +def merge_intervals(intervals: Iterable[tuple[int, int]]) -> list[tuple[int, int]]: + """Merge overlapping or directly adjacent half-open intervals.""" + + merged: list[list[int]] = [] + for start, end in sorted(intervals): + if not merged or start > merged[-1][1]: + merged.append([start, end]) + else: + merged[-1][1] = max(merged[-1][1], end) + return [(start, end) for start, end in merged] + + +def annotate_regions(tx: Transcript) -> list[Region]: + """Partition a mature transcript into UTR/CDS or noncoding-exon sequence.""" + + cds_genomic = tx.feature_intervals.get("CDS", []) + tx.feature_intervals.get("stop_codon", []) + cds = merge_intervals( + piece + for interval in cds_genomic + for piece in tx.genomic_interval_to_transcript(*interval) + ) + if not cds: + if tx.transcript_type == "protein_coding": + raise ValueError(f"protein-coding transcript {tx.transcript_id} has no CDS annotation") + return [Region(0, tx.length, "noncoding_exon")] + cds_start = min(start for start, _ in cds) + cds_end = max(end for _, end in cds) + result: list[Region] = [] + if cds_start: + result.append(Region(0, cds_start, "5putr")) + result.append(Region(cds_start, cds_end, "cds")) + if cds_end < tx.length: + result.append(Region(cds_end, tx.length, "3putr")) + return result diff --git a/src/transcriptml/rbpnet/experiment.py b/src/transcriptml/rbpnet/experiment.py new file mode 100644 index 0000000..85412ff --- /dev/null +++ b/src/transcriptml/rbpnet/experiment.py @@ -0,0 +1,317 @@ +"""Lazy, ergonomic access to a processed transcript-space eCLIP experiment.""" + +from __future__ import annotations + +import csv +import gzip +import json +from dataclasses import dataclass +from pathlib import Path + +import h5py +import numpy as np +import pysam + + +@dataclass(frozen=True) +class RegionRecord: + start: int + end: int + region_type: str + + +@dataclass(frozen=True) +class TranscriptRecord: + gene_id: str + transcript_id: str + chromosome: str + strand: str + length: int + signal_offset: int + sminput_tpm: float + regions: tuple[RegionRecord, ...] + + +@dataclass(frozen=True) +class SampleRecord: + name: str + role: str + effective_library_size: int | None + + +@dataclass(frozen=True) +class GenomicBlock: + chromosome: str + start: int + end: int + + +class ProcessedECLIPDataset: + """Thin lazy reader for one canonical processed eCLIP directory. + + Small metadata tables are loaded at construction. FASTA and HDF5 handles + are opened only on first access and are never inherited when the reader is + pickled, making the object safe to construct before worker processes. + """ + + SUPPORTED_FORMATS = {"transcriptml-rbpnet-experiment", "rbpnet-preprocess-dataset"} + + def __init__(self, processed_dir: str | Path): + self.processed_dir = Path(processed_dir) + manifest_path = self.processed_dir / "manifest.json" + if not manifest_path.is_file(): + raise FileNotFoundError(f"processed manifest not found: {manifest_path}") + self.manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + data_format = self.manifest.get("format") + if data_format not in self.SUPPORTED_FORMATS: + raise ValueError( + f"unsupported processed experiment format {data_format!r}; " + f"supported: {', '.join(sorted(self.SUPPORTED_FORMATS))}" + ) + if str(self.manifest.get("format_version")) != "1": + raise ValueError( + f"unsupported processed experiment format_version {self.manifest.get('format_version')!r}" + ) + files = self.manifest.get("files", {}) + required = {"metadata", "exon_mapping", "sequences", "signals"} + missing_keys = sorted(required - set(files)) + if missing_keys: + raise ValueError(f"manifest is missing file entries: {', '.join(missing_keys)}") + self._metadata_path = self.processed_dir / files["metadata"] + self._exon_path = self.processed_dir / files["exon_mapping"] + self._fasta_path = self.processed_dir / files["sequences"] + self._signal_path = self.processed_dir / files["signals"] + for path in (self._metadata_path, self._exon_path, self._fasta_path, self._signal_path): + if not path.is_file(): + raise FileNotFoundError(f"processed dataset file not found: {path}") + + self.transcripts = self._read_transcripts() + self._transcript_by_id = {tx.transcript_id: tx for tx in self.transcripts} + if len(self._transcript_by_id) != len(self.transcripts): + raise ValueError("duplicate transcript IDs in metadata") + self.samples = tuple( + SampleRecord( + name=sample["name"], + role=sample["role"], + effective_library_size=( + int(sample["effective_library_size"]) + if "effective_library_size" in sample else None + ), + ) + for sample in self.manifest.get("samples", []) + ) + self._sample_index = {sample.name: index for index, sample in enumerate(self.samples)} + if len(self._sample_index) != len(self.samples): + raise ValueError("duplicate sample names in manifest") + self._fasta: pysam.FastaFile | None = None + self._h5: h5py.File | None = None + self._exons_by_transcript: dict[str, list[dict]] | None = None + self._validate_store() + + def _read_transcripts(self) -> tuple[TranscriptRecord, ...]: + records = [] + with self._metadata_path.open(encoding="utf-8") as handle: + for row in csv.DictReader(handle, delimiter="\t"): + regions = tuple( + RegionRecord(int(region["start"]), int(region["end"]), region["type"]) + for region in json.loads(row["region_annotations"]) + ) + record = TranscriptRecord( + gene_id=row["gene_id"], + transcript_id=row["transcript_id"], + chromosome=row["chrom"], + strand=row["strand"], + length=int(row["transcript_length"]), + signal_offset=int(row["signal_offset"]), + sminput_tpm=float(row["sm_input_tpm"]), + regions=regions, + ) + if record.length <= 0: + raise ValueError(f"transcript {record.transcript_id} has non-positive length") + if ( + not regions + or regions[0].start != 0 + or regions[-1].end != record.length + or any( + region.start < 0 + or region.end <= region.start + or (index and regions[index - 1].end != region.start) + for index, region in enumerate(regions) + ) + ): + raise ValueError( + f"region annotations do not partition transcript {record.transcript_id}" + ) + records.append(record) + return tuple(records) + + @staticmethod + def _decode(values) -> list[str]: + return [value.decode() if isinstance(value, bytes) else str(value) for value in values] + + def _open_handles(self) -> None: + if self._fasta is None: + self._fasta = pysam.FastaFile(str(self._fasta_path)) + if self._h5 is None: + self._h5 = h5py.File(self._signal_path, "r") + + def _validate_store(self) -> None: + with h5py.File(self._signal_path, "r") as store: + required = { + "counts", "ip_pooled", "sample_names", "sample_roles", "transcript_ids", + "transcript_offsets", "transcript_lengths", + } + missing = sorted(required - set(store)) + if missing: + raise ValueError(f"signals.h5 is missing datasets: {', '.join(missing)}") + h5_samples = self._decode(store["sample_names"][:]) + if h5_samples != [sample.name for sample in self.samples]: + raise ValueError("sample order differs between manifest and signals.h5") + h5_roles = self._decode(store["sample_roles"][:]) + if h5_roles != [sample.role for sample in self.samples]: + raise ValueError("sample roles differ between manifest and signals.h5") + h5_transcripts = self._decode(store["transcript_ids"][:]) + if h5_transcripts != [tx.transcript_id for tx in self.transcripts]: + raise ValueError("transcript order differs between metadata and signals.h5") + expected_offsets = np.asarray([tx.signal_offset for tx in self.transcripts], dtype=np.int64) + expected_lengths = np.asarray([tx.length for tx in self.transcripts], dtype=np.int64) + if not np.array_equal(store["transcript_offsets"][:], expected_offsets): + raise ValueError("transcript offsets differ between metadata and signals.h5") + if not np.array_equal(store["transcript_lengths"][:], expected_lengths): + raise ValueError("transcript lengths differ between metadata and signals.h5") + total_length = int(expected_lengths.sum()) + if store["counts"].shape != (len(self.samples), total_length): + raise ValueError("counts shape disagrees with sample and transcript metadata") + if store["ip_pooled"].shape != (total_length,): + raise ValueError("ip_pooled shape disagrees with transcript metadata") + with pysam.FastaFile(str(self._fasta_path)) as fasta: + if tuple(fasta.references) != tuple(tx.transcript_id for tx in self.transcripts): + raise ValueError("transcript order differs between metadata and transcript FASTA") + if tuple(fasta.lengths) != tuple(tx.length for tx in self.transcripts): + raise ValueError("transcript lengths differ between metadata and transcript FASTA") + + @property + def sample_names(self) -> tuple[str, ...]: + return tuple(sample.name for sample in self.samples) + + @property + def ip_samples(self) -> tuple[SampleRecord, ...]: + return tuple(sample for sample in self.samples if sample.role == "ip") + + @property + def sminput_sample(self) -> SampleRecord: + inputs = [sample for sample in self.samples if sample.role == "sminput"] + if len(inputs) != 1: + raise ValueError(f"expected exactly one sminput sample, found {len(inputs)}") + return inputs[0] + + @property + def pooled_ip_effective_library_size(self) -> int: + sizes = [sample.effective_library_size for sample in self.ip_samples] + if not sizes or any(size is None for size in sizes): + raise ValueError("manifest lacks effective library size for one or more IP samples") + return sum(int(size) for size in sizes) + + def get_transcript(self, transcript_id: str) -> TranscriptRecord: + try: + return self._transcript_by_id[transcript_id] + except KeyError as exc: + raise KeyError(f"unknown transcript: {transcript_id}") from exc + + def _slice(self, transcript_id: str, start: int, end: int) -> tuple[TranscriptRecord, slice]: + tx = self.get_transcript(transcript_id) + if start < 0 or end < start or end > tx.length: + raise IndexError( + f"invalid interval {transcript_id}:{start}-{end}; transcript length is {tx.length}" + ) + return tx, slice(tx.signal_offset + start, tx.signal_offset + end) + + def get_sequence(self, transcript_id: str, start: int, end: int) -> str: + self._slice(transcript_id, start, end) + self._open_handles() + assert self._fasta is not None + return self._fasta.fetch(transcript_id, start, end) + + def get_profile(self, transcript_id: str, start: int, end: int, sample: str) -> np.ndarray: + _, flat_slice = self._slice(transcript_id, start, end) + try: + sample_index = self._sample_index[sample] + except KeyError as exc: + raise KeyError(f"unknown sample {sample!r}; available: {', '.join(self.sample_names)}") from exc + self._open_handles() + assert self._h5 is not None + return self._h5["counts"][sample_index, flat_slice] + + def get_profiles(self, transcript_id: str, start: int, end: int) -> np.ndarray: + """Return all sample profiles in manifest order for one interval.""" + + _, flat_slice = self._slice(transcript_id, start, end) + self._open_handles() + assert self._h5 is not None + return self._h5["counts"][:, flat_slice] + + def get_pooled_ip_profile(self, transcript_id: str, start: int, end: int) -> np.ndarray: + _, flat_slice = self._slice(transcript_id, start, end) + self._open_handles() + assert self._h5 is not None + return self._h5["ip_pooled"][flat_slice] + + def _load_exons(self) -> None: + exons: dict[str, list[dict]] = {} + with gzip.open(self._exon_path, "rt") as handle: + for row in csv.DictReader(handle, delimiter="\t"): + exons.setdefault(row["transcript_id"], []).append({ + "tx_start": int(row["tx_start"]), + "tx_end": int(row["tx_end"]), + "chromosome": row["chrom"], + "genomic_start": int(row["genomic_start"]), + "genomic_end": int(row["genomic_end"]), + "strand": row["strand"], + }) + self._exons_by_transcript = exons + + def get_genomic_blocks( + self, transcript_id: str, start: int, end: int + ) -> tuple[GenomicBlock, ...]: + """Map one transcript interval to compact genomic exon blocks.""" + + tx, _ = self._slice(transcript_id, start, end) + if self._exons_by_transcript is None: + self._load_exons() + assert self._exons_by_transcript is not None + blocks = [] + for exon in self._exons_by_transcript.get(transcript_id, []): + lo = max(start, exon["tx_start"]) + hi = min(end, exon["tx_end"]) + if lo >= hi: + continue + if tx.strand == "+": + genomic_start = exon["genomic_start"] + lo - exon["tx_start"] + genomic_end = exon["genomic_start"] + hi - exon["tx_start"] + else: + genomic_start = exon["genomic_end"] - (hi - exon["tx_start"]) + genomic_end = exon["genomic_end"] - (lo - exon["tx_start"]) + blocks.append(GenomicBlock(exon["chromosome"], genomic_start, genomic_end)) + if sum(block.end - block.start for block in blocks) != end - start: + raise ValueError(f"exon mapping does not cover {transcript_id}:{start}-{end}") + return tuple(blocks) + + def close(self) -> None: + if self._h5 is not None: + self._h5.close() + self._h5 = None + if self._fasta is not None: + self._fasta.close() + self._fasta = None + + def __getstate__(self): + state = self.__dict__.copy() + state["_h5"] = None + state["_fasta"] = None + return state + + def __enter__(self) -> "ProcessedECLIPDataset": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self.close() diff --git a/src/transcriptml/rbpnet/fasta.py b/src/transcriptml/rbpnet/fasta.py new file mode 100644 index 0000000..06f3554 --- /dev/null +++ b/src/transcriptml/rbpnet/fasta.py @@ -0,0 +1,111 @@ +"""FASTA validation and mature-transcript sequence extraction.""" + +from __future__ import annotations + +from collections import Counter +from pathlib import Path + +import pysam + +from transcriptml.rbpnet._progress import track +from transcriptml.rbpnet.coordinates import Transcript + +_COMPLEMENT = str.maketrans("ACGTRYMKBDHVacgtrymkbdhv", "TGCAYRKMVHDBtgcayrkmvhdb") + + +def ensure_fasta_index(path: str | Path) -> Path: + """Return a FASTA index path, creating the index when reasonable.""" + + path = Path(path) + if not path.is_file(): + raise FileNotFoundError(f"FASTA not found: {path}") + index = Path(str(path) + ".fai") + if not index.exists(): + try: + pysam.faidx(str(path)) + except Exception as exc: + raise RuntimeError(f"could not create FASTA index {index}: {exc}") from exc + return index + + +def reverse_complement(sequence: str) -> str: + """Return the IUPAC-aware DNA reverse complement.""" + + return sequence.translate(_COMPLEMENT)[::-1] + + +def retain_fasta_transcripts( + genome_fasta: str | Path, transcripts: list[Transcript] +) -> tuple[list[Transcript], dict]: + """Skip GTF transcripts on contigs intentionally absent from the FASTA.""" + + ensure_fasta_index(genome_fasta) + with pysam.FastaFile(str(genome_fasta)) as fasta: + fasta_contigs = set(fasta.references) + retained = [tx for tx in transcripts if tx.chrom in fasta_contigs] + missing_counts = Counter(tx.chrom for tx in transcripts if tx.chrom not in fasta_contigs) + if not retained: + missing = ", ".join(sorted(missing_counts)[:5]) or "none" + raise ValueError( + "no annotated transcripts remain after intersecting GTF and FASTA contigs; " + f"GTF-only contig examples: {missing}" + ) + offset = 0 + for tx in retained: + tx.offset = offset + offset += tx.length + return retained, { + "gtf_transcripts_total": len(transcripts), + "transcripts_retained": len(retained), + "transcripts_skipped_missing_fasta_contig": len(transcripts) - len(retained), + "gtf_contigs_missing_from_fasta": dict(sorted(missing_counts.items())), + } + + +def transcript_sequence(fasta: pysam.FastaFile, tx: Transcript) -> str: + """Assemble one mature transcript in transcript 5-prime to 3-prime order.""" + + chunks = [] + for exon in tx.exons: + chunk = fasta.fetch(exon.chrom, exon.start, exon.end) + chunks.append(chunk if tx.strand == "+" else reverse_complement(chunk)) + sequence = "".join(chunks).upper() + if len(sequence) != tx.length: + raise RuntimeError(f"sequence length mismatch for {tx.transcript_id}") + return sequence + + +def write_transcript_fasta( + path: str | Path, + genome_fasta: str | Path, + transcripts: list[Transcript], + *, + progress: bool = True, +) -> dict: + """Write indexed mature-transcript FASTA and return sequence QC.""" + + path = Path(path) + ensure_fasta_index(genome_fasta) + stats = {"transcripts": len(transcripts), "bases": 0, "non_acgtn_bases": 0} + with pysam.FastaFile(str(genome_fasta)) as fasta, path.open("w", encoding="utf-8") as out: + missing = sorted({tx.chrom for tx in transcripts} - set(fasta.references)) + if missing: + raise ValueError( + f"GTF chromosome(s) absent from FASTA: {', '.join(missing[:5])}; " + f"FASTA examples: {', '.join(fasta.references[:5])}" + ) + for tx in track( + transcripts, + "rbpnet preprocess: extract transcript sequences", + total=len(transcripts), + unit="transcripts", + enabled=progress, + ): + sequence = transcript_sequence(fasta, tx) + stats["bases"] += len(sequence) + stats["non_acgtn_bases"] += sum(base not in "ACGTN" for base in sequence) + out.write(f">{tx.transcript_id}\n") + for start in range(0, len(sequence), 80): + out.write(sequence[start : start + 80] + "\n") + pysam.faidx(str(path)) + return stats diff --git a/src/transcriptml/rbpnet/preprocessing.py b/src/transcriptml/rbpnet/preprocessing.py new file mode 100644 index 0000000..0a0c127 --- /dev/null +++ b/src/transcriptml/rbpnet/preprocessing.py @@ -0,0 +1,263 @@ +"""Universal eCLIP preprocessing orchestration.""" + +from __future__ import annotations + +import logging +import platform +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +import h5py +import numpy as np +import pysam + +from transcriptml import __version__ +from transcriptml.progress import log_progress +from transcriptml.rbpnet.annotation import parse_gtf +from transcriptml.rbpnet.fasta import retain_fasta_transcripts, write_transcript_fasta +from transcriptml.rbpnet.serialization import write_exons, write_json, write_metadata, write_regions +from transcriptml.rbpnet.signals import ExonBinIndex, create_signal_store, extract_bam_to_store, write_ip_pooled + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class Sample: + """One named BAM sample and its experimental role.""" + + name: str + path: Path + role: str + + +@dataclass(frozen=True) +class PipelineConfig: + """Configuration for canonical transcript-space eCLIP preprocessing.""" + + genome_fasta: Path + gtf: Path + sminput: Sample + ips: tuple[Sample, ...] + output_dir: Path + read1_rna_strand: str = "opposite" + min_mapq: int = 1 + exclude_duplicates: bool = True + overwrite: bool = False + progress: bool = True + + +_OUTPUT_FILES = ( + "transcripts.tsv", "exons.tsv.gz", "regions.tsv.gz", "transcripts.fa", + "transcripts.fa.fai", "signals.h5", "qc.json", "manifest.json", +) + + +def _prepare_output(path: Path, overwrite: bool) -> None: + path.mkdir(parents=True, exist_ok=True) + conflicts = [path / name for name in _OUTPUT_FILES if (path / name).exists()] + if conflicts and not overwrite: + raise FileExistsError( + f"output files already exist in {path}; choose a new directory or pass --overwrite" + ) + for conflict in conflicts: + conflict.unlink() + + +def _input_record(path: Path) -> dict: + resolved = path.resolve() + stat = resolved.stat() + return {"path": str(resolved), "size_bytes": stat.st_size, "mtime_ns": stat.st_mtime_ns} + + +def preprocess_eclip(config: PipelineConfig) -> dict: + """Create a reusable canonical transcript-space eCLIP experiment. + + The HDF5 track is concatenated transcript space and stays lazy on read. + This stage deliberately performs no peak calling or region selection. + """ + + started = time.time() + if not config.ips: + raise ValueError("at least one IP BAM is required") + if config.sminput.role != "sminput" or any(sample.role != "ip" for sample in config.ips): + raise ValueError("sample roles must be one 'sminput' followed by one or more 'ip' samples") + if config.read1_rna_strand not in {"opposite", "same", "unstranded"}: + raise ValueError("read1_rna_strand must be opposite, same, or unstranded") + if config.min_mapq < 0: + raise ValueError("min_mapq must be non-negative") + sample_names = [config.sminput.name] + [sample.name for sample in config.ips] + if len(sample_names) != len(set(sample_names)): + raise ValueError("sample names must be unique") + + log_progress(f"rbpnet preprocess: prepare {config.output_dir}", enabled=config.progress) + _prepare_output(config.output_dir, config.overwrite) + all_transcripts = parse_gtf(config.gtf, progress=config.progress) + transcripts, contig_filter_qc = retain_fasta_transcripts(config.genome_fasta, all_transcripts) + skipped = contig_filter_qc["transcripts_skipped_missing_fasta_contig"] + if skipped: + missing_contigs = list(contig_filter_qc["gtf_contigs_missing_from_fasta"]) + missing = ", ".join(missing_contigs[:10]) + if len(missing_contigs) > 10: + missing += f", ... ({len(missing_contigs)} contigs total)" + logger.warning( + "Skipping %d transcript(s) on GTF contigs absent from the FASTA: %s", + skipped, + missing, + ) + + sequence_qc = write_transcript_fasta( + config.output_dir / "transcripts.fa", + config.genome_fasta, + transcripts, + progress=config.progress, + ) + write_exons(config.output_dir / "exons.tsv.gz", transcripts, progress=config.progress) + write_regions(config.output_dir / "regions.tsv.gz", transcripts, progress=config.progress) + + samples = [config.sminput, *config.ips] + index = ExonBinIndex(transcripts) + sample_counts: list[np.ndarray] = [] + bam_qc: dict[str, dict] = {} + log_progress("rbpnet preprocess: create base-resolution signal store", enabled=config.progress) + with create_signal_store( + config.output_dir / "signals.h5", + transcripts, + sample_names, + [sample.role for sample in samples], + ) as store: + for row, sample in enumerate(samples): + log_progress( + f"rbpnet preprocess: process {sample.name} ({sample.role})", + enabled=config.progress, + ) + sample_qc, counts = extract_bam_to_store( + sample.path, + row, + store["counts"], + transcripts, + index, + config.read1_rna_strand, + config.min_mapq, + config.exclude_duplicates, + temp_dir=config.output_dir, + progress=config.progress, + ) + sample_qc["bam"] = str(sample.path.resolve()) + sample_qc["role"] = sample.role + sample_qc["retained_fraction_of_read1"] = ( + sample_qc["retained"] / sample_qc["read1_seen"] if sample_qc["read1_seen"] else 0.0 + ) + sample_qc["retained_fraction_of_passing_filters"] = ( + sample_qc["retained"] / sample_qc["passing_filters"] + if sample_qc["passing_filters"] else 0.0 + ) + bam_qc[sample.name] = sample_qc + sample_counts.append(counts) + write_ip_pooled(store, list(range(1, len(samples))), progress=config.progress) + + sm_tpm = write_metadata( + config.output_dir / "transcripts.tsv", + transcripts, + sample_names, + sample_counts, + 0, + progress=config.progress, + ) + region_counts: dict[str, int] = {} + for tx in transcripts: + for region in tx.regions: + region_counts[region.label] = region_counts.get(region.label, 0) + region.end - region.start + qc = { + "format_version": "1", + "pipeline_version": __version__, + "configuration": { + "read1_rna_strand": config.read1_rna_strand, + "min_mapq": config.min_mapq, + "exclude_duplicates": config.exclude_duplicates, + "assignment": ( + "unique strand-compatible mature transcript at read1 5-prime aligned base, " + "with all aligned CIGAR segments compatible with selected exons and junctions" + ), + }, + "annotation": { + **contig_filter_qc, + "transcripts": len(transcripts), + "genes": len({tx.gene_id for tx in transcripts}), + "exons": sum(len(tx.exons) for tx in transcripts), + "transcriptome_bases": sum(tx.length for tx in transcripts), + "chromosomes": sorted({tx.chrom for tx in transcripts}), + "region_bases": dict(sorted(region_counts.items())), + "sminput_transcripts_nonzero": int(np.count_nonzero(sample_counts[0])), + "sminput_transcripts_tpm_ge_1": int(np.count_nonzero(sm_tpm >= 1.0)), + }, + "sequence": sequence_qc, + "samples": bam_qc, + "elapsed_seconds": round(time.time() - started, 3), + } + write_json(config.output_dir / "qc.json", qc) + manifest = { + "format": "transcriptml-rbpnet-experiment", + "format_version": "1", + "created_at": datetime.now(timezone.utc).isoformat(), + "coordinate_system": ( + "all intervals are 0-based, half-open; sequences/tracks are transcript 5-prime to 3-prime" + ), + "inputs": { + "genome_fasta": _input_record(config.genome_fasta), + "gtf": _input_record(config.gtf), + "sminput_bam": _input_record(config.sminput.path), + "ip_bams": [_input_record(sample.path) for sample in config.ips], + }, + "configuration": qc["configuration"], + "samples": [ + { + "name": sample.name, + "role": sample.role, + "effective_library_size": bam_qc[sample.name]["retained"], + "effective_library_size_count_type": "retained_read1_5prime_events", + } + for sample in samples + ], + "derived_signals": { + "ip_pooled": { + "source_samples": [sample.name for sample in config.ips], + "effective_library_size": sum(bam_qc[sample.name]["retained"] for sample in config.ips), + "effective_library_size_count_type": "sum_of_source_effective_library_sizes", + } + }, + "normalization": { + "cpm_formula": "window_count / effective_library_size * 1e6", + "sample_denominator_field": "samples[].effective_library_size", + "effective_library_size_definition": ( + "retained read1 5-prime events used to construct the transcript-space signal track" + ), + "pooled_ip_denominator": "sum of effective_library_size over IP source samples", + }, + "files": { + "metadata": "transcripts.tsv", + "exon_mapping": "exons.tsv.gz", + "regions": "regions.tsv.gz", + "sequences": "transcripts.fa", + "signals": "signals.h5", + "qc": "qc.json", + }, + "software": { + "transcriptml": __version__, + "python": platform.python_version(), + "pysam": pysam.__version__, + "h5py": h5py.__version__, + "numpy": np.__version__, + }, + } + write_json(config.output_dir / "manifest.json", manifest) + log_progress( + f"rbpnet preprocess: complete in {time.time() - started:.1f}s", + enabled=config.progress, + ) + return qc + + +# A familiar alias for callers migrating from the standalone implementation. +run_pipeline = preprocess_eclip diff --git a/src/transcriptml/rbpnet/selection.py b/src/transcriptml/rbpnet/selection.py new file mode 100644 index 0000000..7009eb7 --- /dev/null +++ b/src/transcriptml/rbpnet/selection.py @@ -0,0 +1,741 @@ +"""Explicit, versioned region selection over descriptive eCLIP windows.""" + +from __future__ import annotations + +import csv +import gzip +import hashlib +import json +import math +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Iterator + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from scipy.stats import binom, poisson + +from transcriptml.progress import ProgressReporter, log_progress +from transcriptml.rbpnet.experiment import ProcessedECLIPDataset +from transcriptml.rbpnet.windows import REGION_TYPES, summarize_regions + +SELECTION_STRATEGIES = ("original_rbpnet", "yeo_2026", "peak_gray_negative") + + +@dataclass(frozen=True) +class SelectionConfig: + """Configuration for selecting eligible experimental loci. + + Defaults for ``original_rbpnet`` reproduce the published Horlacher et al. + candidate rules. Defaults for the other two strategies are transparent + starting points and should be reviewed for each assay. + """ + + processed_dir: Path + windows: Path + output_prefix: Path + strategy: str + overwrite: bool = False + progress: bool = True + batch_size: int = 10_000 + # Published v1 selector. + original_min_pvalue: float = 0.01 + original_min_count: int = 8 + original_min_height: int = 2 + original_advance: int = 50 + # Broad measured-window selector. + min_total_count: int = 8 + min_sminput_count: int = 1 + min_ip_count: int = 1 + min_sminput_tpm: float = 0.0 + replicate_mode: str = "combined" + # Peak / gray / confident-negative selector. + peak_fdr: float = 0.05 + peak_min_log2_ratio: float = 1.0 + negative_fdr: float = 0.05 + negative_max_log2_ratio: float = -0.5 + stitch_gap: int = 0 + + +@dataclass(frozen=True) +class SelectionManifest: + """Loaded selection table and its versioned provenance metadata.""" + + path: Path + table: pa.Table + metadata: dict + + @property + def rows(self) -> list[dict]: + return self.table.to_pylist() + + +def _resolve_parquet(path: Path) -> Path: + if path.suffix == ".parquet": + return path + candidate = Path(str(path) + ".parquet") + if candidate.is_file(): + return candidate + raise ValueError("selection currently requires the scanner's Parquet table; pass its .parquet path or prefix") + + +def _scan_metadata(path: Path) -> dict: + metadata = pq.read_schema(path).metadata or {} + raw = metadata.get(b"transcriptml_rbpnet_window_scan") or metadata.get(b"rbpnet_window_scan") + if raw is None: + raise ValueError(f"window Parquet lacks TranscriptML/RBPNet scan metadata: {path}") + return json.loads(raw.decode()) + + +def _iter_window_rows(path: Path, *, batch_size: int) -> Iterator[dict]: + parquet = pq.ParquetFile(path) + for batch in parquet.iter_batches(batch_size=batch_size): + yield from batch.to_pylist() + + +def _stable_example_id( + strategy: str, + transcript_id: str, + start: int, + end: int, + state: str, + replicate_id: str, +) -> str: + payload = "\x1f".join( + ("selection-v1", strategy, transcript_id, str(start), str(end), state, replicate_id) + ) + return "rbp_" + hashlib.sha256(payload.encode()).hexdigest()[:20] + + +def _manifest_schema(ds: ProcessedECLIPDataset, metadata: dict[bytes, bytes]) -> pa.Schema: + fields = [ + pa.field("example_id", pa.string()), + pa.field("gene_id", pa.string()), + pa.field("transcript_id", pa.string()), + pa.field("chromosome", pa.string()), + pa.field("strand", pa.string()), + pa.field("transcript_anchor", pa.int64()), + pa.field("selection_start", pa.int64()), + pa.field("selection_end", pa.int64()), + pa.field("selection_length", pa.int64()), + pa.field("region_type", pa.string()), + ] + for region_type in REGION_TYPES: + fields.extend([ + pa.field(f"region_{region_type}_nt", pa.int64()), + pa.field(f"region_{region_type}_fraction", pa.float64()), + ]) + fields.extend([ + pa.field("genomic_blocks", pa.string()), + pa.field("selection_strategy", pa.string()), + pa.field("selection_state", pa.string()), + pa.field("replicate_id", pa.string()), + pa.field("source_window_count", pa.int64()), + pa.field("sminput_tpm", pa.float64()), + ]) + fields.extend(pa.field(f"{sample.name}_count", pa.int64()) for sample in ds.samples) + fields.extend([ + pa.field("ip_pooled_count", pa.int64()), + ]) + fields.extend(pa.field(f"{sample.name}_cpm", pa.float64()) for sample in ds.samples) + fields.extend([ + pa.field("ip_pooled_cpm", pa.float64()), + pa.field("total_ip_sminput_count", pa.int64()), + pa.field("log2_ip_pooled_vs_sminput", pa.float64()), + ]) + fields.extend(pa.field(f"max_{sample.name}_5pend", pa.int64()) for sample in ds.samples) + fields.extend([ + pa.field("max_ip_pooled_5pend", pa.int64()), + pa.field("selection_pvalue", pa.float64()), + pa.field("selection_qvalue", pa.float64()), + pa.field("source_min_enrichment_pvalue", pa.float64()), + pa.field("source_min_enrichment_qvalue", pa.float64()), + pa.field("source_min_depletion_pvalue", pa.float64()), + pa.field("source_min_depletion_qvalue", pa.float64()), + pa.field("group_gene_id", pa.string()), + pa.field("group_transcript_id", pa.string()), + pa.field("group_chromosome", pa.string()), + ]) + return pa.schema(fields, metadata=metadata) + + +def _format_blocks(ds: ProcessedECLIPDataset, transcript_id: str, start: int, end: int) -> str: + return ";".join( + f"{block.chromosome}:{block.start}-{block.end}" + for block in ds.get_genomic_blocks(transcript_id, start, end) + ) + + +def _base_manifest_row( + ds: ProcessedECLIPDataset, + source: dict, + *, + strategy: str, + state: str, + replicate_id: str = "", + source_window_count: int = 1, + anchor: int | None = None, + selection_pvalue: float = math.nan, + selection_qvalue: float = math.nan, + enrichment_pvalue: float = math.nan, + enrichment_qvalue: float = math.nan, + depletion_pvalue: float = math.nan, + depletion_qvalue: float = math.nan, +) -> dict: + start = int(source["tx_start"]) + end = int(source["tx_end"]) + tx = ds.get_transcript(source["transcript_id"]) + anchor = start + (end - start) // 2 if anchor is None else int(anchor) + row = { + "example_id": _stable_example_id(strategy, tx.transcript_id, start, end, state, replicate_id), + "gene_id": tx.gene_id, + "transcript_id": tx.transcript_id, + "chromosome": tx.chromosome, + "strand": tx.strand, + "transcript_anchor": anchor, + "selection_start": start, + "selection_end": end, + "selection_length": end - start, + "region_type": source["region_type"], + "genomic_blocks": source["genomic_blocks"], + "selection_strategy": strategy, + "selection_state": state, + "replicate_id": replicate_id, + "source_window_count": source_window_count, + "sminput_tpm": float(source["sminput_tpm"]), + "ip_pooled_count": int(source["ip_pooled_count"]), + "ip_pooled_cpm": float(source["ip_pooled_cpm"]), + "total_ip_sminput_count": int(source["total_ip_sminput_count"]), + "log2_ip_pooled_vs_sminput": float(source["log2_ip_pooled_vs_sminput"]), + "max_ip_pooled_5pend": int(source["max_ip_pooled_5pend"]), + "selection_pvalue": float(selection_pvalue), + "selection_qvalue": float(selection_qvalue), + "source_min_enrichment_pvalue": float(enrichment_pvalue), + "source_min_enrichment_qvalue": float(enrichment_qvalue), + "source_min_depletion_pvalue": float(depletion_pvalue), + "source_min_depletion_qvalue": float(depletion_qvalue), + "group_gene_id": tx.gene_id, + "group_transcript_id": tx.transcript_id, + "group_chromosome": tx.chromosome, + } + for region_type in REGION_TYPES: + row[f"region_{region_type}_nt"] = int(source[f"region_{region_type}_nt"]) + row[f"region_{region_type}_fraction"] = float(source[f"region_{region_type}_fraction"]) + for sample in ds.samples: + row[f"{sample.name}_count"] = int(source[f"{sample.name}_count"]) + row[f"{sample.name}_cpm"] = float(source[f"{sample.name}_cpm"]) + row[f"max_{sample.name}_5pend"] = int(source[f"max_{sample.name}_5pend"]) + return row + + +def _interval_source( + ds: ProcessedECLIPDataset, + transcript_id: str, + start: int, + end: int, + *, + pseudocount: float, +) -> dict: + tx = ds.get_transcript(transcript_id) + profiles = ds.get_profiles(transcript_id, start, end) + pooled = ds.get_pooled_ip_profile(transcript_id, start, end) + counts = profiles.sum(axis=1, dtype=np.uint64) + pooled_count = int(pooled.sum(dtype=np.uint64)) + denominators = {sample.name: int(sample.effective_library_size) for sample in ds.samples} + pooled_denominator = ds.pooled_ip_effective_library_size + cpms = { + sample.name: float(counts[i]) / denominators[sample.name] * 1_000_000.0 + for i, sample in enumerate(ds.samples) + } + pooled_cpm = pooled_count / pooled_denominator * 1_000_000.0 + sminput = ds.sminput_sample.name + sminput_index = ds.sample_names.index(sminput) + region_type, region_counts, region_fractions = summarize_regions(tx.regions, start, end) + source = { + "transcript_id": transcript_id, + "tx_start": start, + "tx_end": end, + "region_type": region_type, + "genomic_blocks": _format_blocks(ds, transcript_id, start, end), + "sminput_tpm": tx.sminput_tpm, + "ip_pooled_count": pooled_count, + "ip_pooled_cpm": pooled_cpm, + "total_ip_sminput_count": pooled_count + int(counts[sminput_index]), + "log2_ip_pooled_vs_sminput": math.log2( + (pooled_cpm + pseudocount) / (cpms[sminput] + pseudocount) + ), + "max_ip_pooled_5pend": int(pooled.max(initial=0)), + } + for region in REGION_TYPES: + source[f"region_{region}_nt"] = region_counts[region] + source[f"region_{region}_fraction"] = region_fractions[region] + for i, sample in enumerate(ds.samples): + source[f"{sample.name}_count"] = int(counts[i]) + source[f"{sample.name}_cpm"] = cpms[sample.name] + source[f"max_{sample.name}_5pend"] = int(profiles[i].max(initial=0)) + return source + + +def _bh_adjust(pvalues: np.ndarray, tested: np.ndarray) -> np.ndarray: + qvalues = np.ones(pvalues.shape, dtype=np.float64) + indices = np.nonzero(tested)[0] + if indices.size == 0: + return qvalues + order = indices[np.argsort(pvalues[indices], kind="stable")] + ranked = pvalues[order] * len(order) / np.arange(1, len(order) + 1) + ranked = np.minimum.accumulate(ranked[::-1])[::-1] + qvalues[order] = np.minimum(ranked, 1.0) + return qvalues + + +def _original_rows( + config: SelectionConfig, + ds: ProcessedECLIPDataset, + windows_path: Path, + scan_metadata: dict, +) -> Iterator[dict]: + if int(scan_metadata.get("window_size", -1)) != 100 or int(scan_metadata.get("stride", -1)) != 1: + raise ValueError( + "original_rbpnet requires a 100-nt, stride-1 scan; rerun scan-windows " + "with --window-size 100 --stride 1" + ) + if not bool(scan_metadata.get("omit_incomplete_terminal_windows", False)): + raise ValueError("original_rbpnet requires incomplete terminal windows to be omitted") + current_tx = None + mu = 0.0 + next_start = 0 + reporter = ProgressReporter( + "rbpnet select-regions: test v1 windows", + total=pq.ParquetFile(windows_path).metadata.num_rows, + unit="windows", + enabled=config.progress, + ) + try: + for row in _iter_window_rows(windows_path, batch_size=config.batch_size): + reporter.update() + tx_id = row["transcript_id"] + if tx_id != current_tx: + tx = ds.get_transcript(tx_id) + transcript_count = int(ds.get_pooled_ip_profile(tx_id, 0, tx.length).sum(dtype=np.uint64)) + mu = transcript_count / tx.length * int(row["window_length"]) + current_tx = tx_id + next_start = 0 + start = int(row["tx_start"]) + if start < next_start: + continue + count = int(row["ip_pooled_count"]) + height = int(row["max_ip_pooled_5pend"]) + pvalue = float(poisson.sf(count - 1, mu)) + if ( + pvalue < config.original_min_pvalue + and count >= config.original_min_count + and height >= config.original_min_height + ): + yield _base_manifest_row( + ds, + row, + strategy="original_rbpnet", + state="candidate", + selection_pvalue=pvalue, + ) + next_start = start + config.original_advance + finally: + reporter.close() + + +def _yeo_rows( + config: SelectionConfig, + ds: ProcessedECLIPDataset, + windows_path: Path, +) -> Iterator[dict]: + input_name = ds.sminput_sample.name + reporter = ProgressReporter( + "rbpnet select-regions: filter measured windows", + total=pq.ParquetFile(windows_path).metadata.num_rows, + unit="windows", + enabled=config.progress, + ) + try: + for row in _iter_window_rows(windows_path, batch_size=config.batch_size): + reporter.update() + if float(row["sminput_tpm"]) < config.min_sminput_tpm: + continue + input_count = int(row[f"{input_name}_count"]) + if config.replicate_mode == "combined": + ip_count = int(row["ip_pooled_count"]) + if ( + input_count + ip_count >= config.min_total_count + and input_count >= config.min_sminput_count + and ip_count >= config.min_ip_count + ): + yield _base_manifest_row( + ds, + row, + strategy="yeo_2026", + state="measured", + replicate_id="", + ) + else: + for sample in ds.ip_samples: + ip_count = int(row[f"{sample.name}_count"]) + if ( + input_count + ip_count >= config.min_total_count + and input_count >= config.min_sminput_count + and ip_count >= config.min_ip_count + ): + yield _base_manifest_row( + ds, + row, + strategy="yeo_2026", + state="measured", + replicate_id=sample.name, + ) + finally: + reporter.close() + + +def _peak_statistics( + config: SelectionConfig, + ds: ProcessedECLIPDataset, + windows_path: Path, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + input_name = ds.sminput_sample.name + log_progress("rbpnet select-regions: calculate window statistics", enabled=config.progress) + table = pq.read_table( + windows_path, + columns=[ + f"{input_name}_count", "ip_pooled_count", "total_ip_sminput_count", + "sminput_tpm", + ], + ) + input_counts = table[f"{input_name}_count"].to_numpy(zero_copy_only=False).astype(np.int64) + ip_counts = table["ip_pooled_count"].to_numpy(zero_copy_only=False).astype(np.int64) + totals = table["total_ip_sminput_count"].to_numpy(zero_copy_only=False).astype(np.int64) + tpm = table["sminput_tpm"].to_numpy(zero_copy_only=False).astype(np.float64) + adequate = ( + (totals >= config.min_total_count) + & (input_counts >= config.min_sminput_count) + & (ip_counts >= config.min_ip_count) + & (tpm >= config.min_sminput_tpm) + ) + input_size = int(ds.sminput_sample.effective_library_size) + ip_size = ds.pooled_ip_effective_library_size + null_ip_probability = ip_size / (ip_size + input_size) + enrichment_p = np.ones(len(totals), dtype=np.float64) + depletion_p = np.ones(len(totals), dtype=np.float64) + enrichment_p[adequate] = binom.sf( + ip_counts[adequate] - 1, totals[adequate], null_ip_probability + ) + depletion_p[adequate] = binom.cdf( + ip_counts[adequate], totals[adequate], null_ip_probability + ) + log_progress( + f"rbpnet select-regions: {int(adequate.sum()):,}/{len(adequate):,} windows adequately measured", + enabled=config.progress, + ) + return ( + adequate, + enrichment_p, + _bh_adjust(enrichment_p, adequate), + depletion_p, + _bh_adjust(depletion_p, adequate), + ) + + +def _peak_gray_negative_rows( + config: SelectionConfig, + ds: ProcessedECLIPDataset, + windows_path: Path, + scan_metadata: dict, +) -> Iterator[dict]: + adequate, enrichment_p, enrichment_q, depletion_p, depletion_q = _peak_statistics( + config, ds, windows_path + ) + current: dict | None = None + reporter = ProgressReporter( + "rbpnet select-regions: classify and stitch windows", + total=len(adequate), + unit="windows", + enabled=config.progress, + ) + + def flush_current() -> dict | None: + nonlocal current + if current is None: + return None + source = _interval_source( + ds, + current["transcript_id"], + current["start"], + current["end"], + pseudocount=float(scan_metadata["pseudocount_cpm"]), + ) + if current["state"] == "peak": + pooled = ds.get_pooled_ip_profile( + current["transcript_id"], current["start"], current["end"] + ) + maximum = pooled.max(initial=0) + candidates = np.flatnonzero(pooled == maximum) + center = (len(pooled) - 1) / 2 + local_anchor = min(candidates.tolist(), key=lambda x: (abs(x - center), x)) + anchor = current["start"] + int(local_anchor) + selection_pvalue = current["min_enrichment_p"] + selection_qvalue = current["min_enrichment_q"] + else: + anchor = current["start"] + (current["end"] - current["start"]) // 2 + if current["state"] == "confident_negative": + selection_pvalue = current["min_depletion_p"] + selection_qvalue = current["min_depletion_q"] + else: + selection_pvalue = math.nan + selection_qvalue = math.nan + result = _base_manifest_row( + ds, + source, + strategy="peak_gray_negative", + state=current["state"], + source_window_count=current["source_window_count"], + anchor=anchor, + selection_pvalue=selection_pvalue, + selection_qvalue=selection_qvalue, + enrichment_pvalue=current["min_enrichment_p"], + enrichment_qvalue=current["min_enrichment_q"], + depletion_pvalue=current["min_depletion_p"], + depletion_qvalue=current["min_depletion_q"], + ) + current = None + return result + + try: + for index, row in enumerate(_iter_window_rows(windows_path, batch_size=config.batch_size)): + reporter.update() + if not adequate[index]: + flushed = flush_current() + if flushed is not None: + yield flushed + continue + ratio = float(row["log2_ip_pooled_vs_sminput"]) + if enrichment_q[index] <= config.peak_fdr and ratio >= config.peak_min_log2_ratio: + state = "peak" + elif depletion_q[index] <= config.negative_fdr and ratio <= config.negative_max_log2_ratio: + state = "confident_negative" + else: + state = "gray" + compatible = ( + current is not None + and current["transcript_id"] == row["transcript_id"] + and current["state"] == state + and current["region_type"] == row["region_type"] + and int(row["tx_start"]) <= current["end"] + config.stitch_gap + ) + if not compatible: + flushed = flush_current() + if flushed is not None: + yield flushed + current = { + "transcript_id": row["transcript_id"], + "start": int(row["tx_start"]), + "end": int(row["tx_end"]), + "state": state, + "region_type": row["region_type"], + "source_window_count": 1, + "min_enrichment_p": float(enrichment_p[index]), + "min_enrichment_q": float(enrichment_q[index]), + "min_depletion_p": float(depletion_p[index]), + "min_depletion_q": float(depletion_q[index]), + } + else: + current["end"] = max(current["end"], int(row["tx_end"])) + current["source_window_count"] += 1 + current["min_enrichment_p"] = min(current["min_enrichment_p"], float(enrichment_p[index])) + current["min_enrichment_q"] = min(current["min_enrichment_q"], float(enrichment_q[index])) + current["min_depletion_p"] = min(current["min_depletion_p"], float(depletion_p[index])) + current["min_depletion_q"] = min(current["min_depletion_q"], float(depletion_q[index])) + flushed = flush_current() + if flushed is not None: + yield flushed + finally: + reporter.close() + + +def _validate_config(config: SelectionConfig) -> None: + if config.strategy not in SELECTION_STRATEGIES: + raise ValueError(f"unknown selection strategy {config.strategy!r}") + if config.batch_size <= 0 or config.original_advance <= 0: + raise ValueError("batch_size and original_advance must be positive") + for name in ("original_min_pvalue", "peak_fdr", "negative_fdr"): + value = float(getattr(config, name)) + if not 0 < value <= 1: + raise ValueError(f"{name} must be in (0, 1]") + if min(config.original_min_count, config.original_min_height, config.min_total_count, + config.min_sminput_count, config.min_ip_count, config.stitch_gap) < 0: + raise ValueError("count thresholds and stitch_gap must be non-negative") + if config.min_sminput_tpm < 0: + raise ValueError("min_sminput_tpm must be non-negative") + if config.replicate_mode not in {"combined", "per_ip"}: + raise ValueError("replicate_mode must be combined or per_ip") + + +def _validate_scan_dataset( + ds: ProcessedECLIPDataset, + windows_path: Path, + scan_metadata: dict, +) -> None: + expected_sizes = { + sample.name: int(sample.effective_library_size) + for sample in ds.samples + } + observed_sizes = { + str(name): int(value) + for name, value in scan_metadata.get("effective_library_sizes", {}).items() + } + if observed_sizes != expected_sizes: + raise ValueError( + "window scan effective library sizes do not match the processed experiment" + ) + if int(scan_metadata.get("ip_pooled_effective_library_size", -1)) != ds.pooled_ip_effective_library_size: + raise ValueError("window scan pooled-IP library size does not match the processed experiment") + required_columns = { + "transcript_id", "tx_start", "tx_end", "window_length", "region_type", + "sminput_tpm", "ip_pooled_count", "ip_pooled_cpm", + "total_ip_sminput_count", "log2_ip_pooled_vs_sminput", + "max_ip_pooled_5pend", "genomic_blocks", + } + for sample in ds.samples: + required_columns.update({ + f"{sample.name}_count", f"{sample.name}_cpm", f"max_{sample.name}_5pend" + }) + missing = sorted(required_columns - set(pq.read_schema(windows_path).names)) + if missing: + raise ValueError(f"window table lacks columns required by this experiment: {', '.join(missing)}") + + +def select_regions(config: SelectionConfig) -> dict: + """Select biological loci and write a versioned lightweight manifest.""" + + _validate_config(config) + windows_path = _resolve_parquet(config.windows) + scan_metadata = _scan_metadata(windows_path) + prefix = str(config.output_prefix) + if prefix.endswith((".parquet", ".tsv", ".tsv.gz", ".selection.json")): + raise ValueError("output_prefix must not include a table or metadata suffix") + parquet_path = Path(prefix + ".parquet") + tsv_path = Path(prefix + ".tsv.gz") + sidecar_path = Path(prefix + ".selection.json") + parquet_path.parent.mkdir(parents=True, exist_ok=True) + conflicts = [path for path in (parquet_path, tsv_path, sidecar_path) if path.exists()] + if conflicts and not config.overwrite: + raise FileExistsError(f"selection output already exists ({conflicts[0]}); pass --overwrite") + + log_progress(f"rbpnet select-regions: {config.strategy}", enabled=config.progress) + with ProcessedECLIPDataset(config.processed_dir) as ds: + if any(sample.effective_library_size is None or sample.effective_library_size <= 0 for sample in ds.samples): + raise ValueError("all samples need positive effective_library_size values for selection") + _validate_scan_dataset(ds, windows_path, scan_metadata) + provenance = { + "format": "transcriptml-rbpnet-selection", + "format_version": "1", + "strategy": config.strategy, + "source_processed_dir": str(config.processed_dir.resolve()), + "source_windows": str(windows_path.resolve()), + "window_scan": scan_metadata, + "configuration": { + key: (str(value) if isinstance(value, Path) else value) + for key, value in config.__dict__.items() + if key not in {"processed_dir", "windows", "output_prefix", "progress"} + }, + "statistical_notes": ( + "original_rbpnet uses a one-sided Poisson test against the transcript-level pooled-IP rate" + if config.strategy == "original_rbpnet" + else "peak_gray_negative uses exact conditional binomial tails and BH correction over adequately measured windows" + if config.strategy == "peak_gray_negative" + else "yeo_2026 applies coverage thresholds only and performs no peak test" + ), + } + schema = _manifest_schema( + ds, + {b"transcriptml_rbpnet_selection": json.dumps(provenance, sort_keys=True).encode()}, + ) + if config.strategy == "original_rbpnet": + rows: Iterable[dict] = _original_rows(config, ds, windows_path, scan_metadata) + elif config.strategy == "yeo_2026": + rows = _yeo_rows(config, ds, windows_path) + else: + rows = _peak_gray_negative_rows(config, ds, windows_path, scan_metadata) + + selected = 0 + state_counts: Counter[str] = Counter() + transcript_ids: set[str] = set() + batch: list[dict] = [] + reporter = ProgressReporter( + "rbpnet select-regions: write manifest", + total=None, + unit="examples", + enabled=config.progress, + ) + with gzip.open(tsv_path, "wt", newline="") as tsv_handle, pq.ParquetWriter( + parquet_path, schema, compression="zstd" + ) as parquet_writer: + tsv_writer = csv.DictWriter( + tsv_handle, fieldnames=schema.names, delimiter="\t", lineterminator="\n" + ) + tsv_writer.writeheader() + + def flush() -> None: + if not batch: + return + tsv_writer.writerows(batch) + parquet_writer.write_table(pa.Table.from_pylist(batch, schema=schema)) + batch.clear() + + for row in rows: + batch.append(row) + selected += 1 + state_counts[row["selection_state"]] += 1 + transcript_ids.add(row["transcript_id"]) + reporter.update() + if len(batch) >= config.batch_size: + flush() + flush() + reporter.close() + summary = { + **provenance, + "n_examples": selected, + "n_transcripts": len(transcript_ids), + "state_counts": dict(sorted(state_counts.items())), + "parquet": str(parquet_path), + "tsv": str(tsv_path), + } + sidecar_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") + log_progress( + f"rbpnet select-regions: wrote {selected:,} examples", + enabled=config.progress, + ) + return summary + + +def load_selection_manifest(path: str | Path) -> SelectionManifest: + """Load and validate a version-1 Parquet selection manifest.""" + + manifest_path = _resolve_parquet(Path(path)) + table = pq.read_table(manifest_path) + raw = (table.schema.metadata or {}).get(b"transcriptml_rbpnet_selection") + if raw is None: + raise ValueError(f"selection manifest metadata is missing: {manifest_path}") + metadata = json.loads(raw.decode()) + if metadata.get("format") != "transcriptml-rbpnet-selection" or str(metadata.get("format_version")) != "1": + raise ValueError("unsupported RBPNet selection manifest format/version") + required = { + "example_id", "gene_id", "transcript_id", "chromosome", "strand", + "transcript_anchor", "selection_start", "selection_end", "selection_strategy", + "selection_state", "group_gene_id", "group_transcript_id", "group_chromosome", + } + missing = sorted(required - set(table.column_names)) + if missing: + raise ValueError(f"selection manifest lacks columns: {', '.join(missing)}") + ids = table["example_id"].to_pylist() + if len(ids) != len(set(ids)): + raise ValueError("selection manifest contains duplicate example_id values") + return SelectionManifest(manifest_path, table, metadata) diff --git a/src/transcriptml/rbpnet/serialization.py b/src/transcriptml/rbpnet/serialization.py new file mode 100644 index 0000000..da696dd --- /dev/null +++ b/src/transcriptml/rbpnet/serialization.py @@ -0,0 +1,127 @@ +"""Interoperable metadata and coordinate-table serializers.""" + +from __future__ import annotations + +import csv +import gzip +import json +from pathlib import Path + +import numpy as np + +from transcriptml.rbpnet._progress import track +from transcriptml.rbpnet.coordinates import Transcript + + +def calculate_tpm(raw_counts: np.ndarray, transcripts: list[Transcript]) -> np.ndarray: + """Calculate length-normalized TPM from retained transcript event counts.""" + + lengths_kb = np.asarray([tx.length / 1000.0 for tx in transcripts], dtype=np.float64) + rates = raw_counts.astype(np.float64) / lengths_kb + denominator = rates.sum() + return rates / denominator * 1_000_000.0 if denominator else np.zeros_like(rates) + + +def write_metadata( + path: str | Path, + transcripts: list[Transcript], + sample_names: list[str], + sample_counts: list[np.ndarray], + sminput_index: int, + *, + progress: bool = True, +) -> np.ndarray: + """Write the transcript metadata table and return SMInput TPM values.""" + + tpm = calculate_tpm(sample_counts[sminput_index], transcripts) + fields = [ + "transcript_id", "gene_id", "gene_name", "transcript_name", "transcript_type", + "chrom", "strand", "transcript_length", "signal_offset", "sm_input_raw_count", + "sm_input_tpm", + ] + [f"{name}_raw_5p_count" for name in sample_names] + ["region_annotations"] + with Path(path).open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n") + writer.writeheader() + for index, tx in enumerate(track( + transcripts, + "rbpnet preprocess: write transcript metadata", + total=len(transcripts), + unit="transcripts", + enabled=progress, + )): + row = { + "transcript_id": tx.transcript_id, + "gene_id": tx.gene_id, + "gene_name": tx.gene_name, + "transcript_name": tx.transcript_name, + "transcript_type": tx.transcript_type, + "chrom": tx.chrom, + "strand": tx.strand, + "transcript_length": tx.length, + "signal_offset": tx.offset, + "sm_input_raw_count": int(sample_counts[sminput_index][index]), + "sm_input_tpm": f"{tpm[index]:.8g}", + "region_annotations": json.dumps( + [{"start": r.start, "end": r.end, "type": r.label} for r in tx.regions], + separators=(",", ":"), + ), + } + for sample_name, counts in zip(sample_names, sample_counts): + row[f"{sample_name}_raw_5p_count"] = int(counts[index]) + writer.writerow(row) + return tpm + + +def write_exons(path: str | Path, transcripts: list[Transcript], *, progress: bool = True) -> None: + """Write compressed transcript/genome exon mappings.""" + + fields = [ + "transcript_id", "exon_index_5to3", "exon_number", "exon_id", "tx_start", "tx_end", + "chrom", "genomic_start", "genomic_end", "strand", + ] + with gzip.open(path, "wt", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n") + writer.writeheader() + for tx in track( + transcripts, + "rbpnet preprocess: write exon mappings", + total=len(transcripts), + unit="transcripts", + enabled=progress, + ): + for index, exon in enumerate(tx.exons, 1): + writer.writerow({ + "transcript_id": tx.transcript_id, + "exon_index_5to3": index, + "exon_number": exon.exon_number, + "exon_id": exon.exon_id, + "tx_start": exon.tx_start, + "tx_end": exon.tx_end, + "chrom": exon.chrom, + "genomic_start": exon.start, + "genomic_end": exon.end, + "strand": exon.strand, + }) + + +def write_regions(path: str | Path, transcripts: list[Transcript], *, progress: bool = True) -> None: + """Write compressed transcript region annotations.""" + + with gzip.open(path, "wt", newline="") as handle: + writer = csv.writer(handle, delimiter="\t", lineterminator="\n") + writer.writerow(["transcript_id", "tx_start", "tx_end", "region_type"]) + for tx in track( + transcripts, + "rbpnet preprocess: write region annotations", + total=len(transcripts), + unit="transcripts", + enabled=progress, + ): + for region in tx.regions: + writer.writerow([tx.transcript_id, region.start, region.end, region.label]) + + +def write_json(path: str | Path, value: dict) -> None: + """Write deterministic indented JSON with a trailing newline.""" + + Path(path).write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") diff --git a/src/transcriptml/rbpnet/signals.py b/src/transcriptml/rbpnet/signals.py new file mode 100644 index 0000000..40c84cc --- /dev/null +++ b/src/transcriptml/rbpnet/signals.py @@ -0,0 +1,360 @@ +"""Strand-aware BAM crosslink extraction into a compact HDF5 matrix.""" + +from __future__ import annotations + +import sqlite3 +import tempfile +from collections import Counter +from pathlib import Path + +import h5py +import numpy as np +import pysam + +from transcriptml.rbpnet._progress import ProgressReporter, track +from transcriptml.rbpnet.coordinates import Transcript + + +def read1_rna_strand(is_reverse: bool, orientation: str) -> str | None: + """Infer RNA strand from the read1 alignment and library convention.""" + + read_strand = "-" if is_reverse else "+" + if orientation == "opposite": + return "+" if read_strand == "-" else "-" + if orientation == "same": + return read_strand + if orientation == "unstranded": + return None + raise ValueError(f"unknown read1/RNA orientation: {orientation}") + + +def five_prime_reference_position(read) -> int | None: + """Return the 5-prime aligned reference base, excluding soft clips.""" + + if read.reference_start is None or read.reference_end is None: + return None + return read.reference_end - 1 if read.is_reverse else read.reference_start + + +def _alignment_is_compatible( + read, + exons: tuple[tuple[int, int], ...], + introns: frozenset[tuple[int, int]], +) -> bool: + """Check that the complete CIGAR is compatible with one mature transcript.""" + + if read.reference_start is None or not read.cigartuples: + return False + reference_pos = read.reference_start + has_aligned_bases = False + + def within_one_exon(start: int, end: int) -> bool: + return any(exon_start <= start and end <= exon_end for exon_start, exon_end in exons) + + for operation, length in read.cigartuples: + if length <= 0: + return False + if operation in {pysam.CMATCH, pysam.CEQUAL, pysam.CDIFF}: + if not within_one_exon(reference_pos, reference_pos + length): + return False + reference_pos += length + has_aligned_bases = True + elif operation == pysam.CDEL: + if not within_one_exon(reference_pos, reference_pos + length): + return False + reference_pos += length + elif operation == pysam.CREF_SKIP: + if (reference_pos, reference_pos + length) not in introns: + return False + reference_pos += length + elif operation in {pysam.CINS, pysam.CSOFT_CLIP, pysam.CHARD_CLIP, pysam.CPAD}: + continue + else: + return False + return has_aligned_bases + + +def alignment_is_transcript_compatible(read, transcript: Transcript) -> bool: + """Return whether an alignment follows selected exons and exact junctions.""" + + exons = tuple(sorted((exon.start, exon.end) for exon in transcript.exons)) + introns = frozenset((left[1], right[0]) for left, right in zip(exons, exons[1:])) + return _alignment_is_compatible(read, exons, introns) + + +class ExonBinIndex: + """Small-memory genomic point index for candidate transcript exons.""" + + def __init__(self, transcripts: list[Transcript], bin_size: int = 16_384): + self.transcripts = transcripts + self.bin_size = bin_size + self._bins: dict[tuple[str, str, int], list[tuple[int, int, int, int]]] = {} + self._exons: list[tuple[tuple[int, int], ...]] = [] + self._introns: list[frozenset[tuple[int, int]]] = [] + for tx_index, tx in enumerate(transcripts): + genomic_exons = tuple(sorted((exon.start, exon.end) for exon in tx.exons)) + self._exons.append(genomic_exons) + self._introns.append(frozenset( + (left[1], right[0]) for left, right in zip(genomic_exons, genomic_exons[1:]) + )) + for exon in tx.exons: + record = (exon.start, exon.end, tx_index, exon.tx_start) + for bin_id in range(exon.start // bin_size, (exon.end - 1) // bin_size + 1): + self._bins.setdefault((exon.chrom, exon.strand, bin_id), []).append(record) + + def query(self, chrom: str, pos: int, strand: str | None) -> list[tuple[int, int]]: + """Return ``(transcript_index, transcript_position)`` exon hits.""" + + strands = ("+", "-") if strand is None else (strand,) + matches: list[tuple[int, int]] = [] + for query_strand in strands: + for start, end, tx_index, tx_start in self._bins.get( + (chrom, query_strand, pos // self.bin_size), [] + ): + if start <= pos < end: + delta = pos - start if query_strand == "+" else end - 1 - pos + matches.append((tx_index, tx_start + delta)) + return matches + + def alignment_is_compatible(self, read, transcript_index: int) -> bool: + """Check a read against precomputed exon and intron intervals.""" + + return _alignment_is_compatible( + read, self._exons[transcript_index], self._introns[transcript_index] + ) + + +def ensure_bam_index(path: str | Path) -> None: + """Validate coordinate sorting and create a missing BAM index.""" + + path = Path(path) + if not path.is_file(): + raise FileNotFoundError(f"BAM not found: {path}") + with pysam.AlignmentFile(str(path), "rb") as bam: + sort_order = bam.header.to_dict().get("HD", {}).get("SO") + if sort_order != "coordinate": + raise ValueError(f"BAM must be coordinate sorted (SO:coordinate): {path}") + try: + bam.check_index() + return + except ValueError: + pass + try: + pysam.index(str(path)) + except Exception as exc: + raise RuntimeError(f"could not create BAM index for {path}: {exc}") from exc + + +def validate_bam_contigs(path: str | Path, transcripts: list[Transcript]) -> None: + """Require every retained annotation contig to exist in a BAM header.""" + + with pysam.AlignmentFile(str(path), "rb") as bam: + missing = sorted({tx.chrom for tx in transcripts} - set(bam.references)) + if missing: + raise ValueError( + f"retained GTF chromosome(s) absent from BAM {path}: {', '.join(missing[:5])}; " + f"BAM examples: {', '.join(bam.references[:5])}" + ) + + +def create_signal_store( + path: str | Path, + transcripts: list[Transcript], + sample_names: list[str], + sample_roles: list[str], +) -> h5py.File: + """Create the canonical concatenated transcript-space HDF5 store.""" + + total_length = sum(tx.length for tx in transcripts) + if total_length == 0: + raise ValueError("annotation has zero mature-transcript bases") + store = h5py.File(path, "w") + store.attrs["format"] = "transcriptml-rbpnet-signals" + store.attrs["format_version"] = "1" + store.attrs["coordinate_system"] = "0-based half-open transcript coordinates" + strings = h5py.string_dtype("utf-8") + store.create_dataset( + "transcript_ids", + data=np.asarray([tx.transcript_id for tx in transcripts], dtype=object), + dtype=strings, + ) + store.create_dataset("transcript_offsets", data=np.asarray([tx.offset for tx in transcripts], dtype=np.int64)) + store.create_dataset("transcript_lengths", data=np.asarray([tx.length for tx in transcripts], dtype=np.int64)) + store.create_dataset("sample_names", data=np.asarray(sample_names, dtype=object), dtype=strings) + store.create_dataset("sample_roles", data=np.asarray(sample_roles, dtype=object), dtype=strings) + chunk = min(total_length, 1_048_576) + store.create_dataset( + "counts", + shape=(len(sample_names), total_length), + dtype=np.uint32, + chunks=(1, chunk), + compression="gzip", + compression_opts=4, + shuffle=True, + fillvalue=0, + ) + return store + + +def _flush_counts(connection: sqlite3.Connection, counts: Counter[int]) -> None: + if not counts: + return + connection.executemany( + "INSERT INTO counts(position, count) VALUES (?, ?) " + "ON CONFLICT(position) DO UPDATE SET count=count+excluded.count", + counts.items(), + ) + connection.commit() + counts.clear() + + +def extract_bam_to_store( + bam_path: str | Path, + row: int, + dataset: h5py.Dataset, + transcripts: list[Transcript], + exon_index: ExonBinIndex, + orientation: str, + min_mapq: int, + exclude_duplicates: bool, + temp_dir: str | Path | None = None, + *, + progress: bool = True, +) -> tuple[dict, np.ndarray]: + """Stream one BAM, disk-aggregate sparse events, and fill one HDF5 row.""" + + ensure_bam_index(bam_path) + validate_bam_contigs(bam_path, transcripts) + qc: Counter[str] = Counter() + transcript_counts = np.zeros(len(transcripts), dtype=np.int64) + transcript_hit = np.zeros(len(transcripts), dtype=bool) + with tempfile.TemporaryDirectory(prefix="transcriptml_rbpnet_", dir=temp_dir) as work: + connection = sqlite3.connect(str(Path(work) / "counts.sqlite")) + connection.execute("CREATE TABLE counts(position INTEGER PRIMARY KEY, count INTEGER NOT NULL)") + batch: Counter[int] = Counter() + with pysam.AlignmentFile(str(bam_path), "rb") as bam: + try: + total_records = bam.mapped + bam.unmapped + except (AttributeError, ValueError): + total_records = None + reads = track( + bam.fetch(until_eof=True), + f"rbpnet preprocess: read {Path(bam_path).name}", + total=total_records, + unit="records", + enabled=progress, + ) + for read in reads: + qc["records_seen"] += 1 + if not read.is_read1: + qc["not_read1"] += 1 + continue + qc["read1_seen"] += 1 + if read.is_unmapped: + qc["unmapped"] += 1 + continue + if read.is_secondary: + qc["secondary"] += 1 + continue + if read.is_supplementary: + qc["supplementary"] += 1 + continue + if read.is_qcfail: + qc["qc_fail"] += 1 + continue + if exclude_duplicates and read.is_duplicate: + qc["duplicate"] += 1 + continue + if read.mapping_quality < min_mapq: + qc["low_mapq"] += 1 + continue + pos = five_prime_reference_position(read) + if pos is None: + qc["invalid_crosslink_position"] += 1 + continue + qc["passing_filters"] += 1 + strand = read1_rna_strand(read.is_reverse, orientation) + matches = exon_index.query(read.reference_name, pos, strand) + if not matches: + qc["no_compatible_transcript"] += 1 + continue + compatible = [m for m in matches if exon_index.alignment_is_compatible(read, m[0])] + if not compatible: + qc["transcript_incompatible"] += 1 + continue + if len(compatible) != 1 or len({tx_index for tx_index, _ in compatible}) != 1: + qc["ambiguous_transcript"] += 1 + continue + tx_index, tx_pos = compatible[0] + flat_position = transcripts[tx_index].offset + tx_pos + batch[flat_position] += 1 + transcript_counts[tx_index] += 1 + transcript_hit[tx_index] = True + qc["retained"] += 1 + if len(batch) >= 100_000: + _flush_counts(connection, batch) + _flush_counts(connection, batch) + unique_total = int(connection.execute("SELECT COUNT(*) FROM counts").fetchone()[0]) + cursor = connection.execute("SELECT position, count FROM counts ORDER BY position") + unique_positions = 0 + reporter = ProgressReporter( + f"rbpnet preprocess: write {Path(bam_path).name} signal", + total=unique_total, + unit="positions", + enabled=progress, + ) + while True: + rows = cursor.fetchmany(100_000) + if not rows: + break + unique_positions += len(rows) + positions = np.fromiter((item[0] for item in rows), dtype=np.int64, count=len(rows)) + values64 = np.fromiter((item[1] for item in rows), dtype=np.uint64, count=len(rows)) + if values64.max(initial=0) > np.iinfo(np.uint32).max: + raise OverflowError(f"a crosslink-position count in {bam_path} exceeds uint32") + dataset[row, positions] = values64.astype(np.uint32) + reporter.update(len(rows)) + reporter.close() + connection.close() + qc["unique_crosslink_positions"] = unique_positions + qc["transcripts_with_signal"] = int(transcript_hit.sum()) + qc["transcripts_total"] = len(transcripts) + fields = ( + "records_seen", "read1_seen", "not_read1", "unmapped", "secondary", "supplementary", + "qc_fail", "duplicate", "low_mapq", "invalid_crosslink_position", "passing_filters", + "no_compatible_transcript", "transcript_incompatible", "ambiguous_transcript", "retained", + "unique_crosslink_positions", "transcripts_with_signal", "transcripts_total", + ) + return {field: int(qc[field]) for field in fields}, transcript_counts + + +def write_ip_pooled(store: h5py.File, ip_rows: list[int], *, progress: bool = True) -> None: + """Sum IP rows chunkwise into the derived pooled-IP HDF5 track.""" + + counts = store["counts"] + total = counts.shape[1] + chunk = counts.chunks[1] + pooled = store.create_dataset( + "ip_pooled", + shape=(total,), + dtype=np.uint32, + chunks=(chunk,), + compression="gzip", + compression_opts=4, + shuffle=True, + fillvalue=0, + ) + starts = range(0, total, chunk) + for start in track( + starts, + "rbpnet preprocess: pool IP tracks", + total=len(starts), + unit="chunks", + enabled=progress, + ): + end = min(start + chunk, total) + values = counts[ip_rows, start:end].astype(np.uint64).sum(axis=0) + if values.max(initial=0) > np.iinfo(np.uint32).max: + raise OverflowError("pooled IP count exceeds uint32") + pooled[start:end] = values.astype(np.uint32) + pooled.attrs["source_rows"] = np.asarray(ip_rows, dtype=np.int64) diff --git a/src/transcriptml/rbpnet/windows.py b/src/transcriptml/rbpnet/windows.py new file mode 100644 index 0000000..7127170 --- /dev/null +++ b/src/transcriptml/rbpnet/windows.py @@ -0,0 +1,307 @@ +"""Descriptive configurable scanning over canonical transcript-space signals.""" + +from __future__ import annotations + +import csv +import gzip +import json +import math +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Iterator + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from transcriptml.progress import ProgressReporter, log_progress +from transcriptml.rbpnet.experiment import ProcessedECLIPDataset, RegionRecord + +REGION_TYPES = ("5putr", "cds", "3putr", "noncoding_exon") + + +@dataclass(frozen=True) +class WindowScanConfig: + """Configuration for a descriptive transcript-window scan.""" + + processed_dir: Path + output_prefix: Path + window_size: int = 100 + stride: int = 50 + min_sminput_tpm: float = 0.0 + pseudocount: float = 1.0 + omit_incomplete_terminal_windows: bool = True + overwrite: bool = False + batch_size: int = 10_000 + progress: bool = True + + +def generate_window_bounds( + transcript_length: int, + window_size: int, + stride: int, + omit_incomplete_terminal_windows: bool = True, +) -> Iterator[tuple[int, int]]: + """Yield deterministic zero-based, half-open transcript windows.""" + + if transcript_length < 0: + raise ValueError("transcript length must be non-negative") + if window_size <= 0 or stride <= 0: + raise ValueError("window size and stride must be positive") + for start in range(0, transcript_length, stride): + end = start + window_size + if end > transcript_length: + if omit_incomplete_terminal_windows: + break + end = transcript_length + if end > start: + yield start, end + + +def calculate_gc_fraction(sequence: str) -> float: + """Calculate GC bases divided by total length; ambiguous bases are non-GC.""" + + if not sequence: + return 0.0 + upper = sequence.upper() + return (upper.count("G") + upper.count("C")) / len(upper) + + +def summarize_regions( + regions: Iterable[RegionRecord], start: int, end: int +) -> tuple[str, dict[str, int], dict[str, float]]: + """Summarize exact region overlap and label boundary-crossing windows mixed.""" + + length = end - start + if length <= 0: + raise ValueError("window must have positive length") + counts = {region_type: 0 for region_type in REGION_TYPES} + for region in regions: + if region.region_type not in counts: + raise ValueError(f"unsupported region type in processed metadata: {region.region_type}") + counts[region.region_type] += max(0, min(end, region.end) - max(start, region.start)) + covered = sum(counts.values()) + if covered != length: + raise ValueError(f"region annotations cover {covered} of {length} bases for window {start}-{end}") + present = [region_type for region_type, count in counts.items() if count] + region_type = present[0] if len(present) == 1 else "mixed" + return region_type, counts, {key: value / length for key, value in counts.items()} + + +def _window_schema(sample_names: tuple[str, ...], metadata: dict[bytes, bytes]) -> pa.Schema: + fields = [ + pa.field("gene_id", pa.string()), + pa.field("transcript_id", pa.string()), + pa.field("chromosome", pa.string()), + pa.field("strand", pa.string()), + pa.field("tx_start", pa.int64()), + pa.field("tx_end", pa.int64()), + pa.field("window_length", pa.int64()), + pa.field("region_type", pa.string()), + ] + for region_type in REGION_TYPES: + fields.append(pa.field(f"region_{region_type}_nt", pa.int64())) + fields.append(pa.field(f"region_{region_type}_fraction", pa.float64())) + fields.extend([ + pa.field("gc_fraction", pa.float64()), + pa.field("sminput_tpm", pa.float64()), + pa.field("genomic_blocks", pa.string()), + ]) + fields.extend(pa.field(f"{sample}_count", pa.int64()) for sample in sample_names) + fields.append(pa.field("ip_pooled_count", pa.int64())) + fields.extend(pa.field(f"{sample}_cpm", pa.float64()) for sample in sample_names) + fields.append(pa.field("ip_pooled_cpm", pa.float64())) + fields.extend([ + pa.field("total_ip_sminput_count", pa.int64()), + pa.field("log2_ip_pooled_vs_sminput", pa.float64()), + ]) + fields.extend(pa.field(f"max_{sample}_5pend", pa.int64()) for sample in sample_names) + fields.append(pa.field("max_ip_pooled_5pend", pa.int64())) + return pa.schema(fields, metadata=metadata) + + +def _format_blocks(ds: ProcessedECLIPDataset, tx_id: str, start: int, end: int) -> str: + return ";".join( + f"{block.chromosome}:{block.start}-{block.end}" + for block in ds.get_genomic_blocks(tx_id, start, end) + ) + + +def _validate_config(config: WindowScanConfig) -> None: + if config.window_size <= 0 or config.stride <= 0: + raise ValueError("window_size and stride must be positive") + if config.min_sminput_tpm < 0: + raise ValueError("min_sminput_tpm must be non-negative") + if config.pseudocount <= 0: + raise ValueError("pseudocount must be positive") + if config.batch_size <= 0: + raise ValueError("batch_size must be positive") + + +def scan_windows(config: WindowScanConfig) -> dict: + """Write equivalent gzipped TSV and Parquet descriptive window tables.""" + + _validate_config(config) + prefix_text = str(config.output_prefix) + if prefix_text.endswith((".tsv", ".tsv.gz", ".parquet", ".scan.json")): + raise ValueError("output_prefix must not include a table or metadata suffix") + tsv_path = Path(prefix_text + ".tsv.gz") + parquet_path = Path(prefix_text + ".parquet") + metadata_path = Path(prefix_text + ".scan.json") + tsv_path.parent.mkdir(parents=True, exist_ok=True) + conflicts = [path for path in (tsv_path, parquet_path, metadata_path) if path.exists()] + if conflicts and not config.overwrite: + raise FileExistsError(f"window output already exists ({conflicts[0]}); pass --overwrite to replace it") + + log_progress(f"rbpnet scan-windows: open {config.processed_dir}", enabled=config.progress) + with ProcessedECLIPDataset(config.processed_dir) as ds: + if not ds.ip_samples: + raise ValueError("processed dataset contains no IP samples") + if "ip_pooled" in ds.sample_names: + raise ValueError("sample name 'ip_pooled' is reserved for the derived pooled signal") + missing_sizes = [sample.name for sample in ds.samples if sample.effective_library_size is None] + if missing_sizes: + raise ValueError( + "manifest lacks effective_library_size for sample(s) " + f"{', '.join(missing_sizes)}; rerun preprocessing with the current package" + ) + zero_sizes = [sample.name for sample in ds.samples if int(sample.effective_library_size) <= 0] + if zero_sizes: + raise ValueError(f"effective_library_size must be positive for CPM: {', '.join(zero_sizes)}") + denominators = {sample.name: int(sample.effective_library_size) for sample in ds.samples} + pooled_denominator = sum(denominators[sample.name] for sample in ds.ip_samples) + derived = ds.manifest.get("derived_signals", {}).get("ip_pooled", {}) + if "effective_library_size" in derived and int(derived["effective_library_size"]) != pooled_denominator: + raise ValueError("manifest pooled-IP denominator disagrees with summed IP denominators") + + scan_metadata = { + "format": "transcriptml-rbpnet-window-scan", + "format_version": "1", + "source_processed_dir": str(config.processed_dir.resolve()), + "window_size": config.window_size, + "stride": config.stride, + "min_sminput_tpm": config.min_sminput_tpm, + "pseudocount_cpm": config.pseudocount, + "omit_incomplete_terminal_windows": config.omit_incomplete_terminal_windows, + "effective_library_sizes": denominators, + "ip_pooled_effective_library_size": pooled_denominator, + "log_ratio_formula": "log2((ip_pooled_cpm+pseudocount)/(sminput_cpm+pseudocount))", + } + arrow_metadata = { + b"transcriptml_rbpnet_window_scan": json.dumps(scan_metadata, sort_keys=True).encode() + } + schema = _window_schema(ds.sample_names, arrow_metadata) + summary = { + **scan_metadata, + "transcripts_total": len(ds.transcripts), + "transcripts_passing_sminput_tpm": 0, + "transcripts_scanned": 0, + "windows": 0, + "region_type_windows": Counter(), + "tsv": str(tsv_path), + "parquet": str(parquet_path), + } + batch: list[dict] = [] + with gzip.open(tsv_path, "wt", newline="") as tsv_handle, pq.ParquetWriter( + parquet_path, schema, compression="zstd" + ) as parquet_writer: + tsv_writer = csv.DictWriter( + tsv_handle, fieldnames=schema.names, delimiter="\t", lineterminator="\n" + ) + tsv_writer.writeheader() + + def flush() -> None: + if not batch: + return + tsv_writer.writerows(batch) + parquet_writer.write_table(pa.Table.from_pylist(batch, schema=schema)) + batch.clear() + + reporter = ProgressReporter( + "rbpnet scan-windows: scan transcripts", + total=len(ds.transcripts), + unit="transcripts", + enabled=config.progress, + ) + for tx in ds.transcripts: + if tx.sminput_tpm < config.min_sminput_tpm: + reporter.update() + continue + summary["transcripts_passing_sminput_tpm"] += 1 + emitted = False + sequence = ds.get_sequence(tx.transcript_id, 0, tx.length) + profiles = ds.get_profiles(tx.transcript_id, 0, tx.length) + pooled_profile = ds.get_pooled_ip_profile(tx.transcript_id, 0, tx.length) + # Prefix sums make count aggregation O(1) per window, including + # stride-1 scans used by the published v1 selector. + profile_prefix = np.pad( + profiles.astype(np.uint64).cumsum(axis=1), ((0, 0), (1, 0)) + ) + pooled_prefix = np.pad(pooled_profile.astype(np.uint64).cumsum(), (1, 0)) + gc = np.fromiter((base.upper() in {"G", "C"} for base in sequence), dtype=np.uint8) + gc_prefix = np.pad(gc.astype(np.uint64).cumsum(), (1, 0)) + for start, end in generate_window_bounds( + tx.length, + config.window_size, + config.stride, + config.omit_incomplete_terminal_windows, + ): + emitted = True + window_profiles = profiles[:, start:end] + window_pooled = pooled_profile[start:end] + counts = profile_prefix[:, end] - profile_prefix[:, start] + pooled_count = int(pooled_prefix[end] - pooled_prefix[start]) + region_type, region_counts, region_fractions = summarize_regions(tx.regions, start, end) + cpms = { + sample.name: float(counts[index]) / denominators[sample.name] * 1_000_000.0 + for index, sample in enumerate(ds.samples) + } + pooled_cpm = pooled_count / pooled_denominator * 1_000_000.0 + sminput_name = ds.sminput_sample.name + sminput_index = ds.sample_names.index(sminput_name) + row = { + "gene_id": tx.gene_id, + "transcript_id": tx.transcript_id, + "chromosome": tx.chromosome, + "strand": tx.strand, + "tx_start": start, + "tx_end": end, + "window_length": end - start, + "region_type": region_type, + "gc_fraction": float(gc_prefix[end] - gc_prefix[start]) / (end - start), + "sminput_tpm": tx.sminput_tpm, + "genomic_blocks": _format_blocks(ds, tx.transcript_id, start, end), + "ip_pooled_count": pooled_count, + "ip_pooled_cpm": pooled_cpm, + "total_ip_sminput_count": pooled_count + int(counts[sminput_index]), + "log2_ip_pooled_vs_sminput": math.log2( + (pooled_cpm + config.pseudocount) + / (cpms[sminput_name] + config.pseudocount) + ), + "max_ip_pooled_5pend": int(window_pooled.max(initial=0)), + } + for label in REGION_TYPES: + row[f"region_{label}_nt"] = region_counts[label] + row[f"region_{label}_fraction"] = region_fractions[label] + for index, sample in enumerate(ds.samples): + row[f"{sample.name}_count"] = int(counts[index]) + row[f"{sample.name}_cpm"] = cpms[sample.name] + row[f"max_{sample.name}_5pend"] = int(window_profiles[index].max(initial=0)) + batch.append(row) + summary["windows"] += 1 + summary["region_type_windows"][region_type] += 1 + if len(batch) >= config.batch_size: + flush() + if emitted: + summary["transcripts_scanned"] += 1 + reporter.update(extra=f"{summary['windows']:,} windows") + reporter.close(extra=f"{summary['windows']:,} windows") + flush() + summary["region_type_windows"] = dict(sorted(summary["region_type_windows"].items())) + metadata_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") + log_progress( + f"rbpnet scan-windows: wrote {summary['windows']:,} windows", + enabled=config.progress, + ) + return summary diff --git a/tests/test_data.py b/tests/test_data.py index b12bd2f..d69a86e 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -68,6 +68,21 @@ def test_bundle_roundtrip(tmp_path): np.testing.assert_array_equal(loaded.y, bundle.y) +def test_bundle_named_array_roundtrip_and_mmap(tmp_path): + profiles = np.arange(12, dtype=np.uint32).reshape(2, 6) + bundle = DatasetBundle( + X=np.zeros((2, 4, 6), dtype=np.uint8), + ids=["a", "b"], + arrays={"sminput_profiles": profiles}, + ) + save_bundle(bundle, tmp_path) + loaded = load_bundle(tmp_path, mmap_mode="r") + assert isinstance(loaded.X, np.memmap) + assert isinstance(loaded.arrays["sminput_profiles"], np.memmap) + np.testing.assert_array_equal(loaded.arrays["sminput_profiles"], profiles) + assert loaded.config["named_arrays"]["sminput_profiles"]["dtype"] == "uint32" + + def test_builders_reject_unknown_split_labels(tmp_path): table = tmp_path / "mpra.csv" table.write_text("id,seq,y,split\nx,ACGU,1.0,holdout\n", encoding="utf-8") diff --git a/tests/test_rbpnet.py b/tests/test_rbpnet.py new file mode 100644 index 0000000..92a2b4a --- /dev/null +++ b/tests/test_rbpnet.py @@ -0,0 +1,415 @@ +import csv +import gzip +import json +from pathlib import Path + +import h5py +import numpy as np +import pyarrow.parquet as pq +import pysam +import pytest + +from transcriptml.data.encoding import encode_rna_sequence +from transcriptml.rbpnet.bundle import RBPNetBundleConfig, load_rbpnet_bundle, make_rbpnet_bundle +from transcriptml.rbpnet.coordinates import Exon, Region, Transcript, annotate_regions +from transcriptml.rbpnet.experiment import ProcessedECLIPDataset +from transcriptml.rbpnet.fasta import transcript_sequence +from transcriptml.rbpnet.preprocessing import PipelineConfig, Sample, preprocess_eclip +from transcriptml.rbpnet.selection import SelectionConfig, load_selection_manifest, select_regions +from transcriptml.rbpnet.serialization import calculate_tpm +from transcriptml.rbpnet.signals import ( + ExonBinIndex, + alignment_is_transcript_compatible, + create_signal_store, + extract_bam_to_store, + five_prime_reference_position, + read1_rna_strand, +) +from transcriptml.rbpnet.windows import WindowScanConfig, generate_window_bounds, scan_windows + + +def _tx(strand="+"): + tx = Transcript( + f"tx_{strand}", f"gene_{strand}", "G", "T", "protein_coding", "chr1", strand, + [Exon("chr1", 100, 105, strand), Exon("chr1", 200, 204, strand)], + ) + tx.finalize() + return tx + + +def test_transcript_coordinates_regions_and_minus_sequence(tmp_path): + plus = _tx("+") + minus = _tx("-") + assert [(e.tx_start, e.tx_end) for e in plus.exons] == [(0, 5), (5, 9)] + assert plus.genome_to_transcript("chr1", 200) == 5 + assert plus.transcript_to_genome(5) == ("chr1", 200, "+") + assert [(e.start, e.tx_start, e.tx_end) for e in minus.exons] == [(200, 0, 4), (100, 4, 9)] + assert minus.genome_to_transcript("chr1", 203) == 0 + assert minus.transcript_to_genome(4) == ("chr1", 104, "-") + minus.feature_intervals = {"CDS": [(102, 105), (200, 202)]} + assert annotate_regions(minus) == [ + Region(0, 2, "5putr"), Region(2, 7, "cds"), Region(7, 9, "3putr") + ] + + fasta_path = tmp_path / "genome.fa" + fasta_path.write_text(">chr1\n" + "A" * 100 + "ACGTA" + "N" * 95 + "CCGG" + "A" * 20 + "\n") + pysam.faidx(str(fasta_path)) + with pysam.FastaFile(str(fasta_path)) as fasta: + assert transcript_sequence(fasta, plus) == "ACGTACCGG" + assert transcript_sequence(fasta, minus) == "CCGGTACGT" + + +def _alignment(start, cigar, reverse=False): + read = pysam.AlignedSegment() + read.query_name = "compatibility" + read.flag = 16 if reverse else 0 + read.reference_start = start + read.cigartuples = cigar + return read + + +def _junction_tx(strand="+"): + tx = Transcript( + f"junction_{strand}", f"g_{strand}", "", "", "protein_coding", "chr1", strand, + [Exon("chr1", 100, 110, strand), Exon("chr1", 200, 210, strand)], + ) + tx.finalize() + return tx + + +def test_alignment_compatibility_and_library_orientation(): + assert five_prime_reference_position(_alignment(102, ((0, 6),))) == 102 + reverse = _alignment(102, ((0, 6),), reverse=True) + assert five_prime_reference_position(reverse) == 107 + assert read1_rna_strand(True, "opposite") == "+" + assert read1_rna_strand(False, "same") == "+" + assert read1_rna_strand(False, "unstranded") is None + contained = _alignment(102, ((pysam.CMATCH, 6),)) + junction = _alignment(105, ((pysam.CMATCH, 5), (pysam.CREF_SKIP, 90), (pysam.CMATCH, 5))) + intronic = _alignment(105, ((pysam.CMATCH, 100),)) + wrong_junction = _alignment( + 104, ((pysam.CMATCH, 5), (pysam.CREF_SKIP, 91), (pysam.CMATCH, 5)) + ) + assert alignment_is_transcript_compatible(contained, _junction_tx()) + assert alignment_is_transcript_compatible(junction, _junction_tx()) + assert alignment_is_transcript_compatible(junction, _junction_tx("-")) + assert not alignment_is_transcript_compatible(intronic, _junction_tx()) + assert not alignment_is_transcript_compatible(wrong_junction, _junction_tx()) + + +def test_bam_assignment_reports_transcript_incompatibility(tmp_path): + tx = _junction_tx() + bam_path = tmp_path / "reads.bam" + header = {"HD": {"VN": "1.6", "SO": "coordinate"}, "SQ": [{"SN": "chr1", "LN": 1000}]} + with pysam.AlignmentFile(bam_path, "wb", header=header) as bam: + for name, start, cigar, length in [ + ("contained", 101, ((pysam.CMATCH, 4),), 4), + ("intronic", 105, ((pysam.CMATCH, 100),), 100), + ]: + read = pysam.AlignedSegment() + read.query_name = name + read.query_sequence = "A" * length + read.flag = 81 + read.reference_id = 0 + read.reference_start = start + read.mapping_quality = 60 + read.cigar = cigar + read.query_qualities = pysam.qualitystring_to_array("I" * length) + bam.write(read) + pysam.index(str(bam_path)) + with create_signal_store(tmp_path / "signals.h5", [tx], ["ip"], ["ip"]) as store: + qc, counts = extract_bam_to_store( + bam_path, 0, store["counts"], [tx], ExonBinIndex([tx]), + "opposite", 1, True, tmp_path, progress=False, + ) + assert qc["retained"] == 1 + assert qc["transcript_incompatible"] == 1 + assert counts.tolist() == [1] + + +def test_preprocess_pipeline_manifest_tpm_effective_sizes_and_missing_contig(tmp_path): + fasta = tmp_path / "genome.fa" + fasta.write_text(">chr1\n" + "ACGT" * 20 + "\n") + gtf = tmp_path / "annotation.gtf" + gtf.write_text( + 'chr1\ttest\ttranscript\t1\t20\t.\t+\t.\tgene_id "g1"; transcript_id "t1"; transcript_type "lncRNA";\n' + 'chr1\ttest\texon\t1\t20\t.\t+\t.\tgene_id "g1"; transcript_id "t1"; exon_number 1;\n' + 'chrPatch\ttest\ttranscript\t1\t20\t.\t+\t.\tgene_id "g2"; transcript_id "t2"; transcript_type "lncRNA";\n' + 'chrPatch\ttest\texon\t1\t20\t.\t+\t.\tgene_id "g2"; transcript_id "t2"; exon_number 1;\n' + ) + bam_path = tmp_path / "reads.bam" + header = {"HD": {"VN": "1.6", "SO": "coordinate"}, "SQ": [{"SN": "chr1", "LN": 80}]} + with pysam.AlignmentFile(bam_path, "wb", header=header) as bam: + read = pysam.AlignedSegment() + read.query_name = "retained_read1" + read.query_sequence = "AAAA" + read.flag = 81 + read.reference_id = 0 + read.reference_start = 4 + read.mapping_quality = 60 + read.cigar = ((pysam.CMATCH, 4),) + read.query_qualities = pysam.qualitystring_to_array("IIII") + bam.write(read) + pysam.index(str(bam_path)) + output = tmp_path / "processed" + qc = preprocess_eclip(PipelineConfig( + genome_fasta=fasta, + gtf=gtf, + sminput=Sample("sminput", bam_path, "sminput"), + ips=(Sample("ip1", bam_path, "ip"), Sample("ip2", bam_path, "ip")), + output_dir=output, + progress=False, + )) + assert qc["annotation"]["transcripts"] == 1 + assert qc["annotation"]["transcripts_skipped_missing_fasta_contig"] == 1 + manifest = json.loads((output / "manifest.json").read_text()) + assert [sample["effective_library_size"] for sample in manifest["samples"]] == [1, 1, 1] + assert manifest["derived_signals"]["ip_pooled"]["effective_library_size"] == 2 + with (output / "transcripts.tsv").open() as handle: + row = next(csv.DictReader(handle, delimiter="\t")) + assert int(row["sm_input_raw_count"]) == 1 + assert float(row["sm_input_tpm"]) == pytest.approx(1_000_000) + with ProcessedECLIPDataset(output) as ds: + assert ds.get_sequence("t1", 0, 20) == ("ACGT" * 5) + assert int(ds.get_pooled_ip_profile("t1", 0, 20).sum()) == 2 + + +def _write_processed_fixture(root: Path, *, length=12, profiles=None): + root.mkdir(parents=True) + sequence = ("ACGTGCGTAAAA" * ((length + 11) // 12))[:length] + (root / "transcripts.fa").write_text(f">tx1\n{sequence}\n") + pysam.faidx(str(root / "transcripts.fa")) + if length == 12: + regions = [ + {"start": 0, "end": 4, "type": "5putr"}, + {"start": 4, "end": 9, "type": "cds"}, + {"start": 9, "end": 12, "type": "3putr"}, + ] + exons = [(0, 5, 100, 105), (5, 12, 200, 207)] + else: + regions = [{"start": 0, "end": length, "type": "noncoding_exon"}] + exons = [(0, length, 100, 100 + length)] + with (root / "transcripts.tsv").open("w", newline="") as handle: + writer = csv.DictWriter( + handle, + fieldnames=[ + "transcript_id", "gene_id", "chrom", "strand", "transcript_length", + "signal_offset", "sm_input_tpm", "region_annotations", + ], + delimiter="\t", lineterminator="\n", + ) + writer.writeheader() + writer.writerow({ + "transcript_id": "tx1", "gene_id": "gene1", "chrom": "chr1", "strand": "+", + "transcript_length": length, "signal_offset": 0, "sm_input_tpm": 5.0, + "region_annotations": json.dumps(regions, separators=(",", ":")), + }) + with gzip.open(root / "exons.tsv.gz", "wt", newline="") as handle: + writer = csv.DictWriter( + handle, + fieldnames=[ + "transcript_id", "tx_start", "tx_end", "chrom", "genomic_start", + "genomic_end", "strand", + ], + delimiter="\t", lineterminator="\n", + ) + writer.writeheader() + for tx_start, tx_end, genomic_start, genomic_end in exons: + writer.writerow({ + "transcript_id": "tx1", "tx_start": tx_start, "tx_end": tx_end, + "chrom": "chr1", "genomic_start": genomic_start, "genomic_end": genomic_end, + "strand": "+", + }) + if profiles is None: + profiles = np.asarray([ + [0, 1, 0, 2, 0, 0, 3, 0, 0, 0, 1, 0], + [1, 0, 0, 1, 0, 4, 0, 0, 0, 0, 0, 0], + [0, 2, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1], + ], dtype=np.uint32) + with h5py.File(root / "signals.h5", "w") as h5: + strings = h5py.string_dtype("utf-8") + h5.create_dataset("counts", data=profiles) + h5.create_dataset("ip_pooled", data=profiles[1:].sum(axis=0, dtype=np.uint32)) + h5.create_dataset("sample_names", data=np.asarray(["sminput", "ipA", "ipB"], dtype=object), dtype=strings) + h5.create_dataset("sample_roles", data=np.asarray(["sminput", "ip", "ip"], dtype=object), dtype=strings) + h5.create_dataset("transcript_ids", data=np.asarray(["tx1"], dtype=object), dtype=strings) + h5.create_dataset("transcript_offsets", data=np.asarray([0], dtype=np.int64)) + h5.create_dataset("transcript_lengths", data=np.asarray([length], dtype=np.int64)) + manifest = { + "format": "transcriptml-rbpnet-experiment", "format_version": "1", + "files": { + "metadata": "transcripts.tsv", "exon_mapping": "exons.tsv.gz", + "sequences": "transcripts.fa", "signals": "signals.h5", + }, + "samples": [ + {"name": "sminput", "role": "sminput", "effective_library_size": 7}, + {"name": "ipA", "role": "ip", "effective_library_size": 6}, + {"name": "ipB", "role": "ip", "effective_library_size": 6}, + ], + "derived_signals": {"ip_pooled": {"effective_library_size": 12}}, + } + (root / "manifest.json").write_text(json.dumps(manifest)) + return sequence, profiles + + +def test_tpm_reader_window_counts_mixed_regions_and_normalization(tmp_path): + tx1 = Transcript("a", "g1", "", "", "lncRNA", "chr1", "+", [Exon("chr1", 0, 100, "+")]) + tx2 = Transcript("b", "g2", "", "", "lncRNA", "chr1", "+", [Exon("chr1", 0, 200, "+")]) + tx1.finalize() + tx2.finalize() + np.testing.assert_allclose(calculate_tpm(np.array([10, 10]), [tx1, tx2]), [2 / 3 * 1e6, 1 / 3 * 1e6]) + + root = tmp_path / "processed" + sequence, profiles = _write_processed_fixture(root) + with ProcessedECLIPDataset(root) as ds: + assert ds.get_sequence("tx1", 4, 8) == sequence[4:8] + np.testing.assert_array_equal(ds.get_profile("tx1", 4, 8, "ipA"), profiles[1, 4:8]) + assert [(b.start, b.end) for b in ds.get_genomic_blocks("tx1", 4, 8)] == [(104, 105), (200, 203)] + assert list(generate_window_bounds(10, 4, 4, True)) == [(0, 4), (4, 8)] + assert list(generate_window_bounds(10, 4, 4, False)) == [(0, 4), (4, 8), (8, 10)] + + prefix = tmp_path / "windows" + summary = scan_windows(WindowScanConfig( + processed_dir=root, output_prefix=prefix, window_size=4, stride=4, progress=False, + )) + assert summary["region_type_windows"] == {"5putr": 1, "cds": 1, "mixed": 1} + rows = pq.read_table(str(prefix) + ".parquet").to_pylist() + junction = rows[1] + np.testing.assert_array_equal( + [junction[f"{name}_count"] for name in ("sminput", "ipA", "ipB")], + profiles[:, 4:8].sum(axis=1), + ) + assert junction["ip_pooled_count"] == int(profiles[1:, 4:8].sum()) + assert junction["sminput_cpm"] == pytest.approx(3 / 7 * 1e6) + assert junction["ip_pooled_cpm"] == pytest.approx(6 / 12 * 1e6) + assert junction["max_ipA_5pend"] == 4 + assert rows[2]["region_type"] == "mixed" + assert rows[2]["region_cds_nt"] == 1 and rows[2]["region_3putr_nt"] == 3 + + +def test_selection_strategies_ids_serialization_and_stitching(tmp_path): + root = tmp_path / "processed" + _write_processed_fixture(root) + windows = tmp_path / "windows" + scan_windows(WindowScanConfig( + processed_dir=root, output_prefix=windows, window_size=4, stride=2, progress=False, + )) + yeo = tmp_path / "yeo" + summary = select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=yeo, strategy="yeo_2026", + min_total_count=1, min_sminput_count=0, min_ip_count=0, progress=False, + )) + loaded = load_selection_manifest(yeo) + assert summary["n_examples"] == len(loaded.rows) > 0 + assert len({row["example_id"] for row in loaded.rows}) == len(loaded.rows) + yeo2 = tmp_path / "yeo2" + select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=yeo2, strategy="yeo_2026", + min_total_count=1, min_sminput_count=0, min_ip_count=0, progress=False, + )) + assert [r["example_id"] for r in loaded.rows] == [ + r["example_id"] for r in load_selection_manifest(yeo2).rows + ] + per_ip = tmp_path / "yeo_per_ip" + select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=per_ip, strategy="yeo_2026", + min_total_count=0, min_sminput_count=0, min_ip_count=0, + replicate_mode="per_ip", progress=False, + )) + assert {row["replicate_id"] for row in load_selection_manifest(per_ip).rows} == {"ipA", "ipB"} + + classified = tmp_path / "classified" + classified_summary = select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=classified, + strategy="peak_gray_negative", min_total_count=1, min_sminput_count=0, + min_ip_count=0, peak_fdr=1.0, negative_fdr=1.0, + peak_min_log2_ratio=100.0, negative_max_log2_ratio=-100.0, + progress=False, + )) + assert classified_summary["n_examples"] > 0 + states = set(classified_summary["state_counts"]) + assert states <= {"peak", "gray", "confident_negative"} + assert any(row["source_window_count"] > 1 for row in load_selection_manifest(classified).rows) + + +def test_original_rbpnet_poisson_selection_and_50nt_advance(tmp_path): + length = 500 + profiles = np.zeros((3, length), dtype=np.uint32) + profiles[0, ::50] = 1 + profiles[1, 200] = 10 + profiles[2, 202] = 10 + root = tmp_path / "processed" + _write_processed_fixture(root, length=length, profiles=profiles) + windows = tmp_path / "windows_v1" + scan_windows(WindowScanConfig( + processed_dir=root, output_prefix=windows, window_size=100, stride=1, progress=False, + )) + selected = tmp_path / "original" + summary = select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=selected, + strategy="original_rbpnet", progress=False, + )) + rows = load_selection_manifest(selected).rows + assert summary["n_examples"] > 0 + assert all(row["selection_length"] == 100 for row in rows) + starts = [row["selection_start"] for row in rows] + assert all(right - left >= 50 for left, right in zip(starts, starts[1:])) + assert all(row["selection_pvalue"] < 0.01 for row in rows) + + +def test_materialized_bundle_exact_arrays_jitter_padding_and_mmap(tmp_path): + root = tmp_path / "processed" + sequence, profiles = _write_processed_fixture(root) + windows = tmp_path / "windows" + scan_windows(WindowScanConfig( + processed_dir=root, output_prefix=windows, window_size=4, stride=4, progress=False, + )) + selection = tmp_path / "selection" + select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=selection, strategy="yeo_2026", + min_total_count=0, min_sminput_count=0, min_ip_count=0, progress=False, + )) + out = tmp_path / "bundle" + built = make_rbpnet_bundle(RBPNetBundleConfig( + processed_dir=root, selection_manifest=selection, output_dir=out, + input_length=4, profile_length=4, max_jitter=2, + transcript_end_policy="pad", progress=False, + )) + assert built.X.dtype == np.uint8 and built.X.shape[1:] == (4, 8) + assert built.arrays["sminput_profiles"].dtype == np.uint32 + assert built.arrays["ip_profiles"].shape == (len(built.ids), 2, 8) + assert built.arrays["profile_ip_totals"].dtype == np.uint64 + first = built.metadata[0] + start, end = first["sequence_context_start"], first["sequence_context_end"] + src_start, src_end = max(0, start), min(len(sequence), end) + destination = src_start - start + expected = np.zeros((4, 8), dtype=np.uint8) + expected[:, destination:destination + src_end - src_start] = encode_rna_sequence( + sequence[src_start:src_end] + ) + np.testing.assert_array_equal(built.X[0], expected) + pstart, pend = first["profile_context_start"], first["profile_context_end"] + psrc_start, psrc_end = max(0, pstart), min(len(sequence), pend) + pdestination = psrc_start - pstart + expected_profiles = np.zeros((3, 8), dtype=np.uint32) + expected_profiles[:, pdestination:pdestination + psrc_end - psrc_start] = profiles[:, psrc_start:psrc_end] + np.testing.assert_array_equal(built.arrays["sminput_profiles"][0], expected_profiles[0]) + np.testing.assert_array_equal(built.arrays["ip_profiles"][0], expected_profiles[1:]) + np.testing.assert_array_equal( + built.arrays["profile_ip_totals"][0], expected_profiles[1:].sum(axis=1, dtype=np.uint64) + ) + loaded = load_rbpnet_bundle(out, mmap_mode="r") + assert isinstance(loaded.X, np.memmap) + assert isinstance(loaded.arrays["ip_profiles"], np.memmap) + assert loaded.config["max_jitter"] == 2 + assert loaded.config["sample_metadata"]["ip_axis_order"] == ["ipA", "ipB"] + assert (out / "examples.parquet").is_file() + + dropped_out = tmp_path / "bundle_drop" + dropped = make_rbpnet_bundle(RBPNetBundleConfig( + processed_dir=root, selection_manifest=selection, output_dir=dropped_out, + input_length=4, profile_length=4, max_jitter=1, + transcript_end_policy="drop", progress=False, + )) + assert len(dropped.ids) < len(built.ids) + assert all(m["sequence_left_pad"] == m["sequence_right_pad"] == 0 for m in dropped.metadata) From 0fd311aa290fab613c1e10dfa9a43ae10a4faf25 Mon Sep 17 00:00:00 2001 From: isvock Date: Tue, 11 Aug 2026 16:28:54 -0700 Subject: [PATCH 03/12] Refine RBPNet preprocessing and bundle creation --- README.md | 2 +- docs/api.rst | 2 +- docs/rbpnet.md | 201 +++++++++--- src/transcriptml/rbpnet/__init__.py | 2 +- src/transcriptml/rbpnet/annotation.py | 16 +- src/transcriptml/rbpnet/bundle.py | 292 +++++++++++++---- src/transcriptml/rbpnet/cli.py | 49 ++- src/transcriptml/rbpnet/coordinates.py | 108 ++++++- src/transcriptml/rbpnet/experiment.py | 88 ++++- src/transcriptml/rbpnet/fasta.py | 25 +- src/transcriptml/rbpnet/preprocessing.py | 47 ++- src/transcriptml/rbpnet/selection.py | 97 ++++-- src/transcriptml/rbpnet/serialization.py | 9 +- src/transcriptml/rbpnet/signals.py | 77 ++++- src/transcriptml/rbpnet/windows.py | 3 +- tests/test_rbpnet.py | 394 ++++++++++++++++++++++- 16 files changed, 1211 insertions(+), 201 deletions(-) diff --git a/README.md b/README.md index d5cb359..b38de91 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ TranscriptML currently supports three main workflows: supports targets such as RNA stability, translation, protein output, etc. - **RBPNet/eCLIP data** converts FASTA/GTF/BAM inputs into a canonical - transcript-space experiment, descriptive windows, explicit selection + mature-transcript or full-gene coordinate experiment, descriptive windows, explicit selection manifests, and memory-mappable model-ready NumPy bundles. The RBPNet model and trainer are not implemented yet. diff --git a/docs/api.rst b/docs/api.rst index b83f4be..27c4518 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -63,7 +63,7 @@ RBPNet/eCLIP data :member-order: bysource .. automodule:: transcriptml.rbpnet.bundle - :members: RBPNetBundleConfig, make_rbpnet_bundle, load_rbpnet_bundle + :members: RBPNetBundleConfig, jitter_crop_offset, make_rbpnet_bundle, load_rbpnet_bundle :member-order: bysource Models diff --git a/docs/rbpnet.md b/docs/rbpnet.md index eb01ef6..64b28c3 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -13,7 +13,7 @@ FASTA + one-transcript-per-gene GTF + IP BAM(s) + SMInput BAM preprocess | v - canonical transcript-space HDF5 + FASTA + metadata + canonical locus-coordinate HDF5 + FASTA + metadata | v descriptive window scan @@ -40,11 +40,18 @@ python -m pip install -e ".[rbpnet]" ## 1. Canonical preprocessing The GTF must select one transcript per gene. GTF coordinates are converted to -zero-based, half-open intervals. Exons are spliced in transcript 5′→3′ order; -minus-strand exon sequence is reverse-complemented and minus-strand transcript -position zero is therefore the highest-genomic-coordinate mature RNA base. -Protein-coding transcripts are partitioned into `5putr`, `cds`, and `3putr`. -Transcripts without a CDS are `noncoding_exon` throughout. +zero-based, half-open intervals. `--coordinate-space mature_transcript` is the +default and preserves the original behavior: exons are spliced in transcript +5′→3′ order. `--coordinate-space gene` instead retains the contiguous genomic +span from the first through last selected exon, including introns. Both spaces +are represented in annotated RNA 5′→3′ orientation. Thus, position zero on a +minus-strand locus is its highest-genomic-coordinate base and gene-space +sequence is reverse-complemented. + +Protein-coding exon sequence is partitioned into `5putr`, `cds`, and `3putr`; +exons without a CDS are `noncoding_exon`. Gene space additionally labels the +gaps between exons as `intron`. The labels exhaustively partition each stored +locus, so a boundary-crossing window is reported as `mixed`. ```bash transcriptml rbpnet preprocess \ @@ -53,9 +60,21 @@ transcriptml rbpnet preprocess \ --ip-bam ip1=../RBPNet2/Data/chr21_test/ip1_chr21.bam \ --ip-bam ip2=../RBPNet2/Data/chr21_test/ip2_chr21.bam \ --sminput-bam sminput=../RBPNet2/Data/chr21_test/sminput_chr21.bam \ + --coordinate-space mature_transcript \ --output-dir processed/chr21 ``` +Use `--coordinate-space gene --output-dir processed/chr21_gene` to build the +corresponding full-gene experiment. Coordinate space is recorded in +`manifest.json`, `signals.h5`, scan metadata, selection provenance, and bundle +metadata. Downstream stages reject mismatched inputs. + +For backward compatibility, table columns such as `transcript_id`, `tx_start`, +`tx_end`, and `transcript_anchor` retain their established names. In a +gene-space experiment, their numeric values are gene-space coordinates in the +selected transcript's 5′→3′ orientation; `coordinate_space` removes the +ambiguity. + Only read1 alignments are considered. The crosslink-position signal is the aligned 5′ reference base: `reference_start` for forward alignments and `reference_end - 1` for reverse alignments. In the eCLIP libraries used during @@ -66,12 +85,21 @@ confused with papers that name the crosslink-bearing FASTQ mate R2 before BAM construction. Unmapped, secondary, supplementary, QC-failed, duplicate (by default), and -low-MAPQ alignments are filtered. An event is assigned only when its 5′ base -maps to one strand-compatible selected transcript and the complete alignment is -compatible with that mature transcript. Aligned/deleted reference segments -must remain in exons and each CIGAR `N` must exactly match a selected adjacent -exon junction. Intronic/pre-mRNA alignments are reported separately as -`transcript_incompatible` in `qc.json`. +low-MAPQ alignments are filtered. Assignment then depends on coordinate space: + +- In mature-transcript space, the 5′ base must map uniquely to a + strand-compatible selected exon. Aligned/deleted reference segments must + remain in exons and each CIGAR `N` must exactly match an adjacent selected + exon junction. Intronic/pre-mRNA alignments are not retained. +- In gene space, the 5′ base must map uniquely to a strand-compatible selected + gene span. Every reference-consuming CIGAR operation must remain inside that + span. Intronic alignments are retained, and `N` operations may represent any + splice within the selected gene span; they need not match the one selected + mature isoform. Reads overlapping multiple same-strand gene spans remain + ambiguous rather than being assigned arbitrarily. + +The same read filtering and unique-gene rule apply in both modes. Rejections +are itemized in `qc.json`. Missing `.fai`/`.bai` indexes are created when the source files and their directories are writable. GTF transcripts on contigs absent from the analysis @@ -84,14 +112,16 @@ The canonical directory contains: | --- | --- | | `manifest.json` | Format version, exact sample order/roles, input file stat records, configuration, and effective library sizes | | `qc.json` | Annotation, sequence, read-filter, assignment, and per-sample totals | -| `transcripts.tsv` | Gene/transcript metadata, length, strand, raw counts, SMInput TPM, and compact region annotations | +| `transcripts.tsv` | Gene/transcript metadata, coordinate space, genomic span, length, strand, raw counts, SMInput TPM, and compact region annotations | | `exons.tsv.gz` | Transcript interval ↔ genomic exon block mappings | | `regions.tsv.gz` | Long-form transcript region intervals | -| `transcripts.fa` + `.fai` | Indexed mature-transcript sequences | +| `transcripts.fa` + `.fai` | Indexed 5′→3′ selected-locus sequences | | `signals.h5` | Canonical base-resolution retained read1 5′ counts | -SMInput TPM is calculated from retained transcript counts divided by mature -transcript length, followed by normalization of those rates to one million. +SMInput TPM is calculated from retained locus counts divided by stored locus +length, followed by normalization of those rates to one million. Therefore, +gene-space TPM uses gene-space retained events and full gene-span length; it is +not numerically interchangeable with mature-transcript TPM. Each sample's `effective_library_size` is exactly the number of retained read1 5′ events used to construct its HDF5 track. That field is the CPM denominator. Pooled-IP CPM uses the sum of IP counts divided by the sum of IP effective @@ -102,12 +132,12 @@ library sizes. HDF5 is a hierarchical binary container: datasets behave like typed, multidimensional arrays stored inside a file, can be compressed and chunked, and can be sliced without loading the entire array. `signals.h5` uses a compact -concatenated-transcript layout: +concatenated-locus layout: | Dataset | Shape/dtype | Meaning | | --- | --- | --- | -| `counts` | `(S, total_transcript_bases)`, `uint32` | One row per manifest sample | -| `ip_pooled` | `(total_transcript_bases,)`, `uint32` | Chunkwise sum of all IP rows | +| `counts` | `(S, total_locus_bases)`, `uint32` | One row per manifest sample | +| `ip_pooled` | `(total_locus_bases,)`, `uint32` | Chunkwise sum of all IP rows | | `sample_names` / `sample_roles` | `(S,)`, UTF-8 | HDF5 row identity | | `transcript_ids` | `(T,)`, UTF-8 | Transcript order | | `transcript_offsets` / `transcript_lengths` | `(T,)`, `int64` | Slice boundaries in concatenated space | @@ -121,11 +151,14 @@ from transcriptml.rbpnet import ProcessedECLIPDataset with ProcessedECLIPDataset("processed/chr21") as ds: print(ds.transcripts) print(ds.samples) + print(ds.coordinate_space) seq = ds.get_sequence("ENST...", 100, 400) input_profile = ds.get_profile("ENST...", 100, 400, sample="sminput") ip1_profile = ds.get_profile("ENST...", 100, 400, sample="ip1") pooled = ds.get_pooled_ip_profile("ENST...", 100, 400) blocks = ds.get_genomic_blocks("ENST...", 100, 400) + genomic = ds.coordinate_to_genome("ENST...", 100) + coordinate = ds.genome_to_coordinate("ENST...", "chr21", genomic[1]) ``` ## 2. Descriptive window scanning @@ -176,17 +209,52 @@ transcriptml rbpnet select-regions \ --processed-dir processed/chr21 \ --windows processed/chr21_windows_v1.parquet \ --strategy original_rbpnet \ + --poisson-null ip_locus_density \ --output-prefix processed/chr21_original ``` The default preset follows [Horlacher et al. 2023](https://doi.org/10.1186/s13059-023-03015-7) -and its [reference implementation](https://github.com/mhorlacher/rbpnet): a one-sided Poisson test -against the transcript-level pooled-IP rate (`p < 0.01`), at least 8 pooled -counts, a maximum positional count of at least 2, and a 50-nt advance after an -accepted candidate. The selected interval remains 100 nt; a later 300-nt -bundle context is independent. All thresholds are CLI-configurable. +and its [reference implementation](https://github.com/mhorlacher/rbpnet): a +one-sided Poisson test against the whole-locus pooled-IP rate. For a window of +length `W` in a locus of length `T`, the published/default null is: + +```text +mu = pooled_IP_locus_total / T * W +``` + +A candidate requires uncorrected `p < 0.01`, at least 8 pooled counts, and a +maximum positional count of at least 2. After accepting one, scanning advances +50 nt. No multiple-testing correction is applied: this is the intentionally +lenient published candidate generator, not a calibrated peak caller. The +selected interval remains 100 nt; later model context is independent. + +An explicitly experimental alternative uses the matched input: -### Broad measured windows (`yeo_2026`) +```bash +transcriptml rbpnet select-regions \ + --processed-dir processed/chr21 \ + --windows processed/chr21_windows_v1.parquet \ + --strategy original_rbpnet \ + --poisson-null sminput \ + --sminput-poisson-pseudocount 1 \ + --output-prefix processed/chr21_original_sminput_null +``` + +Its exact expectation is: + +```text +mu = (SMInput_window_count + 1) + * (pooled_IP_effective_library_size / SMInput_effective_library_size) +``` + +The add-one term is applied in SMInput count space and then exposure-scaled to +the pooled-IP library. This asks whether the observed pooled-IP count exceeds +the background expected at the IP sequencing depth. It retains the same +uncorrected `p < 0.01`, count/height requirements, and 50-nt advance. The +manifest records `poisson_null`, the pseudocount, formula, and each selected +window's `selection_null_mean`. This mode is not attributed to Horlacher et al. + +### Broad measured windows (`broad_coverage`) This strategy applies coverage thresholds without a peak test: @@ -194,16 +262,21 @@ This strategy applies coverage thresholds without a peak test: transcriptml rbpnet select-regions \ --processed-dir processed/chr21 \ --windows processed/chr21_windows_100nt.parquet \ - --strategy yeo_2026 \ - --min-total-count 8 --min-sminput-count 1 --min-ip-count 1 \ + --strategy broad_coverage \ + --min-total-count 6 --min-sminput-count 0 --min-ip-count 0 \ + --replicate-mode per_ip \ --output-prefix processed/chr21_measured ``` -`--replicate-mode combined` retains one row while preserving every replicate -column. `per_ip` emits a replicate-identified row for each IP satisfying the -thresholds. This is a configurable implementation of the broad inclusion -philosophy; its defaults are not an assertion that one coverage cutoff is -universally optimal. +The defaults implement `IP + SMInput > 5`, including cases where either track +is zero. The default `--replicate-mode per_ip` applies that criterion independently as +`IP_replicate_j + SMInput >= 6` and emits a replicate-identified row. +`combined` instead applies it to pooled IP and emits one row while preserving +every replicate column. This selector is inspired by the broad-coverage +training philosophy of Yeo et al.; it works on arbitrary TranscriptML scan +tables and is **not** an exact reproduction of their non-overlapping, +annotation-aware Skipper window generation. A locked exact-window preset can +be added separately later. ### Peak / gray / confident negative @@ -212,19 +285,25 @@ transcriptml rbpnet select-regions \ --processed-dir processed/chr21 \ --windows processed/chr21_windows_100nt.parquet \ --strategy peak_gray_negative \ - --min-total-count 8 --min-sminput-count 1 --min-ip-count 1 \ + --min-total-count 8 --min-sminput-count 0 --min-ip-count 0 \ --peak-fdr 0.05 --peak-min-log2-ratio 1 \ --negative-fdr 0.05 --negative-max-log2-ratio -0.5 \ --stitch-gap 0 \ --output-prefix processed/chr21_peak_gray_negative ``` +By default, combined pooled-IP + SMInput count of at least 8 is sufficient to +enter classification; neither individual track must be nonzero. Optional +`--min-ip-count` and `--min-sminput-count` knobs remain available. Thus both +`IP=0, SMInput=N` and `IP=N, SMInput=0` are valid, informative cases. Adequately measured windows are tested by conditioning on pooled-IP + SMInput -counts. The null IP probability is determined by their effective library -sizes. One-sided exact binomial enrichment/depletion p-values are -Benjamini–Hochberg corrected. Peaks require enrichment plus a minimum log2 -effect; confident negatives require depletion plus a maximum log2 effect; the -remaining adequate windows are gray. Low-information windows are omitted. +counts. The null IP probability is determined by effective library sizes. +One-sided exact binomial enrichment/depletion p-values naturally handle both +extremes and are Benjamini–Hochberg corrected. The log2 effect uses the +scanner's explicit CPM-scale pseudocount, so it remains finite at zero. Peaks +require enrichment plus a minimum log2 effect; confident negatives require +depletion plus a maximum log2 effect; the remaining adequate windows are gray. +Low-total-information windows are omitted. Overlapping/nearby windows are stitched only when transcript, state, and region type agree. Peak anchors are pooled-signal maxima; negative and gray anchors are interval midpoints. These defaults are transparent starting choices, not a @@ -233,7 +312,7 @@ definitive CLIP peak caller. ### Selection manifest v1 Every row has a deterministic content-derived `example_id`; gene, transcript, -chromosome, strand, anchor, and half-open selection interval; region overlap; +chromosome, strand, coordinate space, anchor, and half-open selection interval; region overlap; strategy/state and optional replicate identity; sample/pooled signal summaries; statistical fields; and explicit gene/transcript/chromosome grouping columns. Parquet metadata and the sidecar preserve the complete selection and scan @@ -252,7 +331,7 @@ transcriptml rbpnet make-bundle \ --input-length 300 \ --profile-length 300 \ --max-jitter 32 \ - --transcript-end-policy pad + --transcript-end-policy shift_to_fit ``` For requested input length `L`, profile length `P`, and maximum future jitter @@ -276,16 +355,40 @@ representation. The ordinary TranscriptML sidecars (`ids.txt`, full selected-example metadata is copied to `examples.parquet`. All arrays can be loaded with NumPy `mmap_mode="r"`. -With `--transcript-end-policy drop`, examples whose full materialized sequence -or profile context crosses a transcript end are removed and counted. With -`pad`, fixed widths are preserved using all-zero sequence/profile padding and -the validity masks distinguish padding from valid ambiguous sequence or true -zero signal. +Boundary handling is configurable: + +- `shift_to_fit` is the default. It first centers the full stored width around + the anchor, then shifts the interval right or left until it lies entirely in + the locus. The anchor need not remain centered, but every stored position is + real sequence/signal. A locus shorter than either requested stored width is + genuinely insufficient; its examples are dropped and counted as + `n_dropped_short_loci`, never silently padded. +- `drop` preserves the old behavior of discarding any example whose centered + stored sequence or profile crosses a boundary. +- `pad` preserves fixed widths with all-zero sequence/profile padding. Validity + masks distinguish padding from ambiguous sequence or true zero signal. + +Each example's metadata records `locus_length`, biological anchor, +`sequence_materialized_start/end`, `profile_materialized_start/end`, anchor +offsets, padding amounts, and the crop offsets at both jitter extremes. + +For `shift_to_fit`, a future jitter shift `s` in `[-J,+J]` must be resolved in +biological coordinates—not universally as `J+s`: + +```text +desired_crop_start = anchor - crop_length//2 + s +actual_crop_start = clip(desired_crop_start, 0, locus_length-crop_length) +crop_offset = actual_crop_start - materialized_start +``` -A future jitter shift `s` in `[-J, +J]` takes each requested crop starting at -`J + s` in its materialized array. Bundle construction does not perform random -augmentation. Original v1 can use `J=0`; a jitter-ready workflow can use -`J=32`. Candidate-scan advance and training-time jitter are unrelated. +Near a boundary, multiple requested shifts can collapse to the same closest +legal crop. The shifted stored interval of width `L+2J` contains every such +legal `L`-nt crop when the locus is long enough. Sequence and profile starts +are recorded separately because their requested lengths may differ. The helper +`transcriptml.rbpnet.bundle.jitter_crop_offset` implements this formula. +Legacy `pad` bundles retain their centered `J+s` crop convention. Bundle +construction does not itself perform random augmentation. Candidate-scan +advance and training-time jitter remain unrelated. ```python from transcriptml.rbpnet.bundle import load_rbpnet_bundle diff --git a/src/transcriptml/rbpnet/__init__.py b/src/transcriptml/rbpnet/__init__.py index 390c299..e7e2f73 100644 --- a/src/transcriptml/rbpnet/__init__.py +++ b/src/transcriptml/rbpnet/__init__.py @@ -1,4 +1,4 @@ -"""RBPNet/eCLIP transcript-space data preparation. +"""RBPNet/eCLIP transcript-oriented locus data preparation. The public API deliberately stops at model-ready data. Neural-network architectures, losses, and training are not part of this module. diff --git a/src/transcriptml/rbpnet/annotation.py b/src/transcriptml/rbpnet/annotation.py index 37ad1ee..da25b2b 100644 --- a/src/transcriptml/rbpnet/annotation.py +++ b/src/transcriptml/rbpnet/annotation.py @@ -8,7 +8,7 @@ from pathlib import Path from transcriptml.rbpnet._progress import track -from transcriptml.rbpnet.coordinates import Exon, Transcript, annotate_regions +from transcriptml.rbpnet.coordinates import COORDINATE_SPACES, Exon, Transcript, annotate_regions _ATTR_RE = re.compile(r'([^\s;]+)\s+(?:"([^"]*)"|([^;\s]+))') @@ -23,10 +23,19 @@ def _open_text(path: Path): return gzip.open(path, "rt") if path.suffix == ".gz" else path.open(encoding="utf-8") -def parse_gtf(path: str | Path, *, progress: bool = True) -> list[Transcript]: +def parse_gtf( + path: str | Path, + *, + coordinate_space: str = "mature_transcript", + progress: bool = True, +) -> list[Transcript]: """Parse and validate a one-transcript-per-gene annotation.""" path = Path(path) + if coordinate_space not in COORDINATE_SPACES: + raise ValueError( + f"coordinate_space must be one of {', '.join(COORDINATE_SPACES)}" + ) if not path.is_file(): raise FileNotFoundError(f"GTF not found: {path}") transcripts: dict[str, Transcript] = {} @@ -60,6 +69,9 @@ def parse_gtf(path: str | Path, *, progress: bool = True) -> list[Transcript]: transcript_type=attrs.get("transcript_type", attrs.get("gene_type", "")), chrom=chrom, strand=strand, + coordinate_space=coordinate_space, + genomic_start=start, + genomic_end=end, ) elif feature == "exon": exon_rows.setdefault(tx_id, []).append( diff --git a/src/transcriptml/rbpnet/bundle.py b/src/transcriptml/rbpnet/bundle.py index 7034fa7..884278b 100644 --- a/src/transcriptml/rbpnet/bundle.py +++ b/src/transcriptml/rbpnet/bundle.py @@ -29,7 +29,7 @@ class RBPNetBundleConfig: input_length: int = 300 profile_length: int = 300 max_jitter: int = 0 - transcript_end_policy: str = "drop" + transcript_end_policy: str = "shift_to_fit" overwrite: bool = False progress: bool = True @@ -40,6 +40,46 @@ def _materialized_interval(anchor: int, length: int, jitter: int) -> tuple[int, return start, start + width +def _shifted_materialized_interval( + anchor: int, + length: int, + jitter: int, + locus_length: int, +) -> tuple[int, int] | None: + """Return a centered-then-clipped real interval, or ``None`` if too short.""" + + width = length + 2 * jitter + if locus_length < width: + return None + centered_start, _ = _materialized_interval(anchor, length, jitter) + start = min(max(centered_start, 0), locus_length - width) + return start, start + width + + +def jitter_crop_offset( + *, + anchor: int, + materialized_start: int, + locus_length: int, + crop_length: int, + jitter_shift: int, +) -> int: + """Derive a legal future crop offset from explicit biological coordinates. + + This is the coordinate contract used by ``shift_to_fit`` bundles. Requested + shifts near a boundary can map to the same closest legal crop. + """ + + if crop_length <= 0 or locus_length < crop_length: + raise ValueError("locus must be at least as long as the requested crop") + desired_start = anchor - crop_length // 2 + jitter_shift + actual_start = min(max(desired_start, 0), locus_length - crop_length) + offset = actual_start - materialized_start + if offset < 0: + raise ValueError("materialized interval does not contain the requested legal crop") + return offset + + def _source_and_destination(start: int, end: int, transcript_length: int) -> tuple[int, int, int, int]: source_start = max(0, start) source_end = min(transcript_length, end) @@ -77,8 +117,8 @@ def _validate_config(config: RBPNetBundleConfig) -> None: raise ValueError("input_length and profile_length must be positive") if config.max_jitter < 0: raise ValueError("max_jitter must be non-negative") - if config.transcript_end_policy not in {"drop", "pad"}: - raise ValueError("transcript_end_policy must be drop or pad") + if config.transcript_end_policy not in {"drop", "pad", "shift_to_fit"}: + raise ValueError("transcript_end_policy must be drop, pad, or shift_to_fit") def _sorted_rows(manifest: SelectionManifest) -> list[dict]: @@ -91,6 +131,11 @@ def _validate_manifest_dataset( ds: ProcessedECLIPDataset, ) -> None: scan = manifest.metadata.get("window_scan", {}) + manifest_space = manifest.metadata.get( + "coordinate_space", scan.get("coordinate_space", "mature_transcript") + ) + if manifest_space != ds.coordinate_space: + raise ValueError("selection manifest coordinate space does not match processed experiment") observed = { str(name): int(value) for name, value in scan.get("effective_library_sizes", {}).items() @@ -141,6 +186,7 @@ def make_rbpnet_bundle(config: RBPNetBundleConfig) -> DatasetBundle: raise ValueError("processed dataset contains no IP samples") kept: list[dict] = [] dropped = 0 + dropped_short_locus = 0 for row in rows: tx = ds.get_transcript(row["transcript_id"]) anchor = int(row["transcript_anchor"]) @@ -158,8 +204,26 @@ def make_rbpnet_bundle(config: RBPNetBundleConfig) -> DatasetBundle: replicate_id = str(row["replicate_id"]) if replicate_id and replicate_id not in {sample.name for sample in ip_samples}: raise ValueError(f"selection manifest has unknown IP replicate_id {replicate_id!r}") - seq_start, seq_end = _materialized_interval(anchor, config.input_length, config.max_jitter) - profile_start, profile_end = _materialized_interval(anchor, config.profile_length, config.max_jitter) + if config.transcript_end_policy == "shift_to_fit": + seq_interval = _shifted_materialized_interval( + anchor, config.input_length, config.max_jitter, tx.length + ) + profile_interval = _shifted_materialized_interval( + anchor, config.profile_length, config.max_jitter, tx.length + ) + if seq_interval is None or profile_interval is None: + dropped += 1 + dropped_short_locus += 1 + continue + seq_start, seq_end = seq_interval + profile_start, profile_end = profile_interval + else: + seq_start, seq_end = _materialized_interval( + anchor, config.input_length, config.max_jitter + ) + profile_start, profile_end = _materialized_interval( + anchor, config.profile_length, config.max_jitter + ) in_bounds = ( seq_start >= 0 and seq_end <= tx.length and profile_start >= 0 and profile_end <= tx.length @@ -169,15 +233,23 @@ def make_rbpnet_bundle(config: RBPNetBundleConfig) -> DatasetBundle: continue row = dict(row) row.update({ + "coordinate_space": ds.coordinate_space, + "locus_length": tx.length, "sequence_context_start": seq_start, "sequence_context_end": seq_end, + "sequence_materialized_start": seq_start, + "sequence_materialized_end": seq_end, "profile_context_start": profile_start, "profile_context_end": profile_end, + "profile_materialized_start": profile_start, + "profile_materialized_end": profile_end, + "sequence_anchor_offset": anchor - seq_start, + "profile_anchor_offset": anchor - profile_start, }) kept.append(row) if not kept: raise ValueError( - "no examples remain after transcript-end handling; use --transcript-end-policy pad " + "no examples remain after locus-end handling; use --transcript-end-policy pad " "or reduce context/jitter lengths" ) @@ -229,7 +301,7 @@ def make_rbpnet_bundle(config: RBPNetBundleConfig) -> DatasetBundle: "selection_sminput_counts": selection_sminput_counts, "selection_ip_counts": selection_ip_counts, } - metadata: list[dict] = [] + metadata_by_index: list[dict | None] = [None] * n_examples reporter = ProgressReporter( "rbpnet make-bundle: materialize examples", total=n_examples, @@ -237,71 +309,143 @@ def make_rbpnet_bundle(config: RBPNetBundleConfig) -> DatasetBundle: enabled=config.progress, ) input_name = ds.sminput_sample.name + input_index = ds.sample_names.index(input_name) + ip_sample_indices = [ds.sample_names.index(sample.name) for sample in ip_samples] + rows_by_transcript: dict[str, list[tuple[int, dict]]] = {} for index, row in enumerate(kept): - tx = ds.get_transcript(row["transcript_id"]) - seq_start = int(row["sequence_context_start"]) - seq_end = int(row["sequence_context_end"]) - src_start, src_end, dst_start, dst_end = _source_and_destination( - seq_start, seq_end, tx.length - ) - if src_end > src_start: - encoded = encode_rna_sequence(ds.get_sequence(tx.transcript_id, src_start, src_end)) - X[index, :, dst_start:dst_end] = encoded - sequence_valid_mask[index, dst_start:dst_end] = 1 - - profile_start = int(row["profile_context_start"]) - profile_end = int(row["profile_context_end"]) - psrc_start, psrc_end, pdst_start, pdst_end = _source_and_destination( - profile_start, profile_end, tx.length - ) - if psrc_end > psrc_start: - sminput = ds.get_profile(tx.transcript_id, psrc_start, psrc_end, input_name) - sminput_profiles[index, pdst_start:pdst_end] = sminput - for ip_index, sample in enumerate(ip_samples): - ip_profiles[index, ip_index, pdst_start:pdst_end] = ds.get_profile( - tx.transcript_id, psrc_start, psrc_end, sample.name - ) - profile_valid_mask[index, pdst_start:pdst_end] = 1 - profile_sminput_totals[index] = sminput_profiles[index].sum(dtype=np.uint64) - profile_ip_totals[index] = ip_profiles[index].sum(axis=1, dtype=np.uint64) - selection_sminput_counts[index] = int(row[f"{input_name}_count"]) - selection_ip_counts[index] = np.asarray( - [int(row[f"{sample.name}_count"]) for sample in ip_samples], dtype=np.uint64 - ) - metadata.append({ - "example_id": row["example_id"], - "gene_id": row["gene_id"], - "transcript_id": row["transcript_id"], - "chromosome": row["chromosome"], - "strand": row["strand"], - "transcript_anchor": int(row["transcript_anchor"]), - "selection_start": int(row["selection_start"]), - "selection_end": int(row["selection_end"]), - "region_type": row["region_type"], - "selection_strategy": row["selection_strategy"], - "selection_state": row["selection_state"], - "replicate_id": row["replicate_id"], - "group_gene_id": row["group_gene_id"], - "group_transcript_id": row["group_transcript_id"], - "group_chromosome": row["group_chromosome"], - "sequence_context_start": seq_start, - "sequence_context_end": seq_end, - "sequence_left_pad": max(0, -seq_start), - "sequence_right_pad": max(0, seq_end - tx.length), - "profile_context_start": profile_start, - "profile_context_end": profile_end, - "profile_left_pad": max(0, -profile_start), - "profile_right_pad": max(0, profile_end - tx.length), - }) - reporter.update() + rows_by_transcript.setdefault(row["transcript_id"], []).append((index, row)) + for transcript_id, indexed_rows in rows_by_transcript.items(): + tx = ds.get_transcript(transcript_id) + # Read each locus only once. This avoids repeatedly decompressing + # the same HDF5 chunks when stable example-ID order interleaves + # windows from many transcripts. + locus_sequence = ds.get_sequence(transcript_id, 0, tx.length) + locus_profiles = ds.get_profiles(transcript_id, 0, tx.length) + for index, row in indexed_rows: + seq_start = int(row["sequence_context_start"]) + seq_end = int(row["sequence_context_end"]) + src_start, src_end, dst_start, dst_end = _source_and_destination( + seq_start, seq_end, tx.length + ) + if src_end > src_start: + encoded = encode_rna_sequence(locus_sequence[src_start:src_end]) + X[index, :, dst_start:dst_end] = encoded + sequence_valid_mask[index, dst_start:dst_end] = 1 + + profile_start = int(row["profile_context_start"]) + profile_end = int(row["profile_context_end"]) + psrc_start, psrc_end, pdst_start, pdst_end = _source_and_destination( + profile_start, profile_end, tx.length + ) + if psrc_end > psrc_start: + sminput_profiles[index, pdst_start:pdst_end] = locus_profiles[ + input_index, psrc_start:psrc_end + ] + ip_profiles[index, :, pdst_start:pdst_end] = locus_profiles[ + ip_sample_indices, psrc_start:psrc_end + ] + profile_valid_mask[index, pdst_start:pdst_end] = 1 + profile_sminput_totals[index] = sminput_profiles[index].sum(dtype=np.uint64) + profile_ip_totals[index] = ip_profiles[index].sum(axis=1, dtype=np.uint64) + selection_sminput_counts[index] = int(row[f"{input_name}_count"]) + selection_ip_counts[index] = np.asarray( + [int(row[f"{sample.name}_count"]) for sample in ip_samples], dtype=np.uint64 + ) + metadata_by_index[index] = { + "example_id": row["example_id"], + "gene_id": row["gene_id"], + "transcript_id": row["transcript_id"], + "chromosome": row["chromosome"], + "strand": row["strand"], + "coordinate_space": ds.coordinate_space, + "locus_length": tx.length, + "transcript_anchor": int(row["transcript_anchor"]), + "selection_start": int(row["selection_start"]), + "selection_end": int(row["selection_end"]), + "region_type": row["region_type"], + "selection_strategy": row["selection_strategy"], + "selection_state": row["selection_state"], + "replicate_id": row["replicate_id"], + "group_gene_id": row["group_gene_id"], + "group_transcript_id": row["group_transcript_id"], + "group_chromosome": row["group_chromosome"], + "sequence_context_start": seq_start, + "sequence_context_end": seq_end, + "sequence_materialized_start": seq_start, + "sequence_materialized_end": seq_end, + "sequence_anchor_offset": int(row["sequence_anchor_offset"]), + "sequence_left_pad": max(0, -seq_start), + "sequence_right_pad": max(0, seq_end - tx.length), + "profile_context_start": profile_start, + "profile_context_end": profile_end, + "profile_materialized_start": profile_start, + "profile_materialized_end": profile_end, + "profile_anchor_offset": int(row["profile_anchor_offset"]), + "profile_left_pad": max(0, -profile_start), + "profile_right_pad": max(0, profile_end - tx.length), + "sequence_crop_offset_at_minus_max_jitter": ( + jitter_crop_offset( + anchor=int(row["transcript_anchor"]), + materialized_start=seq_start, + locus_length=tx.length, + crop_length=config.input_length, + jitter_shift=-config.max_jitter, + ) + if tx.length >= config.input_length + and config.transcript_end_policy != "pad" + else None + ), + "sequence_crop_offset_at_plus_max_jitter": ( + jitter_crop_offset( + anchor=int(row["transcript_anchor"]), + materialized_start=seq_start, + locus_length=tx.length, + crop_length=config.input_length, + jitter_shift=config.max_jitter, + ) + if tx.length >= config.input_length + and config.transcript_end_policy != "pad" + else None + ), + "profile_crop_offset_at_minus_max_jitter": ( + jitter_crop_offset( + anchor=int(row["transcript_anchor"]), + materialized_start=profile_start, + locus_length=tx.length, + crop_length=config.profile_length, + jitter_shift=-config.max_jitter, + ) + if tx.length >= config.profile_length + and config.transcript_end_policy != "pad" + else None + ), + "profile_crop_offset_at_plus_max_jitter": ( + jitter_crop_offset( + anchor=int(row["transcript_anchor"]), + materialized_start=profile_start, + locus_length=tx.length, + crop_length=config.profile_length, + jitter_shift=config.max_jitter, + ) + if tx.length >= config.profile_length + and config.transcript_end_policy != "pad" + else None + ), + } + reporter.update() reporter.close() + if any(item is None for item in metadata_by_index): + raise AssertionError("internal error: missing materialized example metadata") + metadata = [item for item in metadata_by_index if item is not None] X.flush() for array in arrays.values(): array.flush() # Keep a self-contained scalable copy of the selected-example contract, # sorted in the exact same stable-ID order as the arrays. - example_table = pa.Table.from_pylist(kept, schema=manifest.table.schema) + example_table = pa.Table.from_pylist(kept).replace_schema_metadata( + manifest.table.schema.metadata + ) pq.write_table(example_table, config.output_dir / "examples.parquet", compression="zstd") config_payload = { "builder": "rbpnet", @@ -311,18 +455,34 @@ def make_rbpnet_bundle(config: RBPNetBundleConfig) -> DatasetBundle: "source_selection_manifest": str(manifest.path.resolve()), "source_selection_sha256": _sha256(manifest.path), "selection": manifest.metadata, + "coordinate_space": ds.coordinate_space, "input_length": config.input_length, "profile_length": config.profile_length, "max_jitter": config.max_jitter, "materialized_sequence_length": sequence_width, "materialized_profile_length": profile_width, "jitter_contract": ( - "future shift s in [-max_jitter,+max_jitter] takes sequence/profile crop " - "starting at max_jitter+s from their respective materialized arrays" + { + "crop_offset": "max_jitter + jitter_shift", + "jitter_shift_range": [-config.max_jitter, config.max_jitter], + "note": "pad preserves the legacy centered padded materialization contract", + } + if config.transcript_end_policy == "pad" + else { + "desired_crop_start": "anchor - crop_length//2 + jitter_shift", + "actual_crop_start": "clip(desired_crop_start, 0, locus_length-crop_length)", + "crop_offset": "actual_crop_start - materialized_start", + "jitter_shift_range": [-config.max_jitter, config.max_jitter], + "note": ( + "boundary clipping may map multiple requested shifts to the same legal crop; " + "use per-example materialized starts and anchors, not max_jitter+jitter_shift" + ), + } ), "transcript_end_policy": config.transcript_end_policy, "n_selected_manifest_rows": len(rows), "n_dropped_at_transcript_ends": dropped, + "n_dropped_short_loci": dropped_short_locus, "sample_metadata": { "sminput": { "name": ds.sminput_sample.name, diff --git a/src/transcriptml/rbpnet/cli.py b/src/transcriptml/rbpnet/cli.py index 7dc16b6..6784503 100644 --- a/src/transcriptml/rbpnet/cli.py +++ b/src/transcriptml/rbpnet/cli.py @@ -17,7 +17,7 @@ def add_rbpnet_parser(subparsers) -> None: commands = root.add_subparsers(dest="rbpnet_command", required=True) preprocess = commands.add_parser( - "preprocess", help="Build a canonical mature-transcript eCLIP experiment" + "preprocess", help="Build a canonical mature-transcript or full-gene eCLIP experiment" ) preprocess.add_argument("--genome-fasta", required=True, type=Path) preprocess.add_argument("--gtf", required=True, type=Path, help="one-transcript-per-gene GTF") @@ -27,6 +27,12 @@ def add_rbpnet_parser(subparsers) -> None: ) preprocess.add_argument("--sminput-bam", required=True, metavar="[LABEL=]PATH") preprocess.add_argument("--output-dir", required=True, type=Path) + preprocess.add_argument( + "--coordinate-space", + choices=("mature_transcript", "gene"), + default="mature_transcript", + help="canonical locus coordinates (default: mature_transcript)", + ) preprocess.add_argument( "--read1-rna-strand", choices=("opposite", "same", "unstranded"), default="opposite", help="RNA strand relative to read1 alignment (default: opposite for eCLIP)", @@ -64,17 +70,38 @@ def add_rbpnet_parser(subparsers) -> None: select.add_argument("--output-prefix", required=True, type=Path) select.add_argument( "--strategy", required=True, - choices=("original_rbpnet", "yeo_2026", "peak_gray_negative"), + choices=("original_rbpnet", "broad_coverage", "peak_gray_negative"), ) select.add_argument("--original-min-pvalue", type=float, default=0.01) select.add_argument("--original-min-count", type=int, default=8) select.add_argument("--original-min-height", type=int, default=2) select.add_argument("--original-advance", type=int, default=50) - select.add_argument("--min-total-count", type=int, default=8) - select.add_argument("--min-sminput-count", type=int, default=1) - select.add_argument("--min-ip-count", type=int, default=1) + select.add_argument( + "--poisson-null", choices=("ip_locus_density", "sminput"), + default="ip_locus_density", + help="v1 Poisson null; sminput is an experimental library-scaled alternative", + ) + select.add_argument( + "--sminput-poisson-pseudocount", type=float, default=1.0, + help="additive SMInput window-count pseudocount for the experimental null (default: 1)", + ) + select.add_argument( + "--min-total-count", type=int, default=None, + help="minimum IP+SMInput count (default: 6 for broad_coverage, 8 for peak_gray_negative)", + ) + select.add_argument( + "--min-sminput-count", type=int, default=0, + help="optional per-window SMInput minimum (default: 0)", + ) + select.add_argument( + "--min-ip-count", type=int, default=0, + help="optional pooled/per-replicate IP minimum (default: 0)", + ) select.add_argument("--min-sminput-tpm", type=float, default=0.0) - select.add_argument("--replicate-mode", choices=("combined", "per_ip"), default="combined") + select.add_argument( + "--replicate-mode", choices=("combined", "per_ip"), default="per_ip", + help="broad_coverage eligibility mode (default: per_ip)", + ) select.add_argument("--peak-fdr", type=float, default=0.05) select.add_argument("--peak-min-log2-ratio", type=float, default=1.0) select.add_argument("--negative-fdr", type=float, default=0.05) @@ -92,7 +119,12 @@ def add_rbpnet_parser(subparsers) -> None: bundle.add_argument("--input-length", type=int, default=300) bundle.add_argument("--profile-length", type=int, default=300) bundle.add_argument("--max-jitter", type=int, default=0) - bundle.add_argument("--transcript-end-policy", choices=("drop", "pad"), default="drop") + bundle.add_argument( + "--transcript-end-policy", + choices=("drop", "pad", "shift_to_fit"), + default="shift_to_fit", + help="locus-boundary handling (default: shift_to_fit)", + ) bundle.add_argument("--overwrite", action="store_true") bundle.add_argument("--no-progress", action="store_true") @@ -129,6 +161,7 @@ def run_rbpnet_command(args: argparse.Namespace, parser: argparse.ArgumentParser sminput=_sample(args.sminput_bam, "sminput"), ips=tuple(_sample(value, "ip") for value in args.ip_bam), output_dir=args.output_dir, + coordinate_space=args.coordinate_space, read1_rna_strand=args.read1_rna_strand, min_mapq=args.min_mapq, exclude_duplicates=not args.include_duplicates, @@ -166,6 +199,8 @@ def run_rbpnet_command(args: argparse.Namespace, parser: argparse.ArgumentParser original_min_count=args.original_min_count, original_min_height=args.original_min_height, original_advance=args.original_advance, + poisson_null=args.poisson_null, + sminput_poisson_pseudocount=args.sminput_poisson_pseudocount, min_total_count=args.min_total_count, min_sminput_count=args.min_sminput_count, min_ip_count=args.min_ip_count, diff --git a/src/transcriptml/rbpnet/coordinates.py b/src/transcriptml/rbpnet/coordinates.py index 39faf06..e32f984 100644 --- a/src/transcriptml/rbpnet/coordinates.py +++ b/src/transcriptml/rbpnet/coordinates.py @@ -1,4 +1,4 @@ -"""Spliced transcript coordinate models and interval conversion.""" +"""Transcript-oriented mature-transcript and full-gene coordinates.""" from __future__ import annotations @@ -6,6 +6,9 @@ from typing import Iterable +COORDINATE_SPACES = ("mature_transcript", "gene") + + @dataclass class Exon: """Half-open genomic exon with an assigned transcript interval.""" @@ -37,7 +40,7 @@ class Region: @dataclass class Transcript: - """One mature transcript in transcript 5-prime to 3-prime orientation.""" + """One selected transcript/gene in RNA 5-prime to 3-prime orientation.""" transcript_id: str gene_id: str @@ -50,6 +53,9 @@ class Transcript: feature_intervals: dict[str, list[tuple[int, int]]] = field(default_factory=dict) regions: list[Region] = field(default_factory=list) offset: int = 0 + coordinate_space: str = "mature_transcript" + genomic_start: int | None = None + genomic_end: int | None = None def finalize(self) -> None: """Validate exons, order them 5-prime to 3-prime, and assign coordinates.""" @@ -58,20 +64,48 @@ def finalize(self) -> None: raise ValueError(f"transcript {self.transcript_id} has no exons") if any(e.chrom != self.chrom or e.strand != self.strand for e in self.exons): raise ValueError(f"inconsistent chromosome/strand in {self.transcript_id}") + if self.coordinate_space not in COORDINATE_SPACES: + raise ValueError(f"invalid coordinate space: {self.coordinate_space!r}") genomic = sorted(self.exons, key=lambda e: (e.start, e.end)) for left, right in zip(genomic, genomic[1:]): if left.end > right.start: raise ValueError(f"overlapping exons in {self.transcript_id}") + exon_start = genomic[0].start + exon_end = genomic[-1].end + if self.genomic_start is None: + self.genomic_start = exon_start + if self.genomic_end is None: + self.genomic_end = exon_end + # The selected gene/transcript span is intentionally bounded by its + # first and last exon. GTF transcript rows should agree, but using the + # actual annotated sequence prevents terminal non-exonic padding. + if self.genomic_start > exon_start or self.genomic_end < exon_end: + raise ValueError(f"transcript span does not contain all exons in {self.transcript_id}") + self.genomic_start = exon_start + self.genomic_end = exon_end ordered = genomic if self.strand == "+" else list(reversed(genomic)) - cursor = 0 - for exon in ordered: - exon.tx_start = cursor - cursor += exon.end - exon.start - exon.tx_end = cursor + if self.coordinate_space == "mature_transcript": + cursor = 0 + for exon in ordered: + exon.tx_start = cursor + cursor += exon.end - exon.start + exon.tx_end = cursor + else: + for exon in ordered: + if self.strand == "+": + exon.tx_start = exon.start - self.genomic_start + exon.tx_end = exon.end - self.genomic_start + else: + exon.tx_start = self.genomic_end - exon.end + exon.tx_end = self.genomic_end - exon.start self.exons = ordered @property def length(self) -> int: + if self.coordinate_space == "gene": + if self.genomic_start is None or self.genomic_end is None: + raise ValueError(f"transcript {self.transcript_id} has not been finalized") + return self.genomic_end - self.genomic_start return sum(e.end - e.start for e in self.exons) def genome_to_transcript(self, chrom: str, pos: int) -> int | None: @@ -79,6 +113,11 @@ def genome_to_transcript(self, chrom: str, pos: int) -> int | None: if chrom != self.chrom: return None + if self.coordinate_space == "gene": + assert self.genomic_start is not None and self.genomic_end is not None + if self.genomic_start <= pos < self.genomic_end: + return pos - self.genomic_start if self.strand == "+" else self.genomic_end - 1 - pos + return None for exon in self.exons: if exon.start <= pos < exon.end: delta = pos - exon.start if self.strand == "+" else exon.end - 1 - pos @@ -90,6 +129,10 @@ def transcript_to_genome(self, pos: int) -> tuple[str, int, str]: if pos < 0 or pos >= self.length: raise IndexError(f"transcript position {pos} outside [0,{self.length})") + if self.coordinate_space == "gene": + assert self.genomic_start is not None and self.genomic_end is not None + genomic = self.genomic_start + pos if self.strand == "+" else self.genomic_end - 1 - pos + return self.chrom, genomic, self.strand for exon in self.exons: if exon.tx_start <= pos < exon.tx_end: delta = pos - exon.tx_start @@ -100,6 +143,14 @@ def transcript_to_genome(self, pos: int) -> tuple[str, int, str]: def genomic_interval_to_transcript(self, start: int, end: int) -> list[tuple[int, int]]: """Map a half-open genomic interval to covered transcript pieces.""" + if self.coordinate_space == "gene": + assert self.genomic_start is not None and self.genomic_end is not None + lo, hi = max(start, self.genomic_start), min(end, self.genomic_end) + if lo >= hi: + return [] + if self.strand == "+": + return [(lo - self.genomic_start, hi - self.genomic_start)] + return [(self.genomic_end - hi, self.genomic_end - lo)] pieces: list[tuple[int, int]] = [] for exon in self.exons: lo, hi = max(start, exon.start), min(end, exon.end) @@ -124,8 +175,47 @@ def merge_intervals(intervals: Iterable[tuple[int, int]]) -> list[tuple[int, int return [(start, end) for start, end in merged] +def _gene_regions(tx: Transcript, cds: list[tuple[int, int]]) -> list[Region]: + """Partition a full-gene locus into exon-derived labels and introns.""" + + coding = bool(cds) + if not coding and tx.transcript_type == "protein_coding": + raise ValueError(f"protein-coding transcript {tx.transcript_id} has no CDS annotation") + cds_start = min((start for start, _ in cds), default=0) + cds_end = max((end for _, end in cds), default=0) + labeled_exons: list[Region] = [] + for exon in sorted(tx.exons, key=lambda item: item.tx_start): + boundaries = [exon.tx_start, exon.tx_end] + if coding: + boundaries.extend( + value for value in (cds_start, cds_end) if exon.tx_start < value < exon.tx_end + ) + boundaries = sorted(set(boundaries)) + for start, end in zip(boundaries, boundaries[1:]): + if not coding: + label = "noncoding_exon" + elif end <= cds_start: + label = "5putr" + elif start >= cds_end: + label = "3putr" + else: + label = "cds" + labeled_exons.append(Region(start, end, label)) + + result: list[Region] = [] + cursor = 0 + for region in labeled_exons: + if cursor < region.start: + result.append(Region(cursor, region.start, "intron")) + result.append(region) + cursor = region.end + if cursor < tx.length: + result.append(Region(cursor, tx.length, "intron")) + return result + + def annotate_regions(tx: Transcript) -> list[Region]: - """Partition a mature transcript into UTR/CDS or noncoding-exon sequence.""" + """Partition the selected coordinate space into exhaustive region labels.""" cds_genomic = tx.feature_intervals.get("CDS", []) + tx.feature_intervals.get("stop_codon", []) cds = merge_intervals( @@ -133,6 +223,8 @@ def annotate_regions(tx: Transcript) -> list[Region]: for interval in cds_genomic for piece in tx.genomic_interval_to_transcript(*interval) ) + if tx.coordinate_space == "gene": + return _gene_regions(tx, cds) if not cds: if tx.transcript_type == "protein_coding": raise ValueError(f"protein-coding transcript {tx.transcript_id} has no CDS annotation") diff --git a/src/transcriptml/rbpnet/experiment.py b/src/transcriptml/rbpnet/experiment.py index 85412ff..fcfdc51 100644 --- a/src/transcriptml/rbpnet/experiment.py +++ b/src/transcriptml/rbpnet/experiment.py @@ -30,6 +30,9 @@ class TranscriptRecord: signal_offset: int sminput_tpm: float regions: tuple[RegionRecord, ...] + coordinate_space: str = "mature_transcript" + genomic_start: int = 0 + genomic_end: int = 0 @dataclass(frozen=True) @@ -72,6 +75,9 @@ def __init__(self, processed_dir: str | Path): raise ValueError( f"unsupported processed experiment format_version {self.manifest.get('format_version')!r}" ) + self.coordinate_space = self.manifest.get("coordinate_space", "mature_transcript") + if self.coordinate_space not in {"mature_transcript", "gene"}: + raise ValueError(f"unsupported coordinate_space {self.coordinate_space!r}") files = self.manifest.get("files", {}) required = {"metadata", "exon_mapping", "sequences", "signals"} missing_keys = sorted(required - set(files)) @@ -121,6 +127,9 @@ def _read_transcripts(self) -> tuple[TranscriptRecord, ...]: transcript_id=row["transcript_id"], chromosome=row["chrom"], strand=row["strand"], + coordinate_space=row.get("coordinate_space", self.coordinate_space), + genomic_start=int(row.get("genomic_start") or 0), + genomic_end=int(row.get("genomic_end") or 0), length=int(row["transcript_length"]), signal_offset=int(row["signal_offset"]), sminput_tpm=float(row["sm_input_tpm"]), @@ -128,6 +137,18 @@ def _read_transcripts(self) -> tuple[TranscriptRecord, ...]: ) if record.length <= 0: raise ValueError(f"transcript {record.transcript_id} has non-positive length") + if record.coordinate_space != self.coordinate_space: + raise ValueError( + f"metadata coordinate space differs for {record.transcript_id}" + ) + if record.coordinate_space == "gene" and ( + record.genomic_start < 0 + or record.genomic_end - record.genomic_start != record.length + ): + raise ValueError( + f"invalid gene span for {record.transcript_id}: " + f"{record.genomic_start}-{record.genomic_end}" + ) if ( not regions or regions[0].start != 0 @@ -184,6 +205,11 @@ def _validate_store(self) -> None: raise ValueError("counts shape disagrees with sample and transcript metadata") if store["ip_pooled"].shape != (total_length,): raise ValueError("ip_pooled shape disagrees with transcript metadata") + stored_space = store.attrs.get("coordinate_space") + if isinstance(stored_space, bytes): + stored_space = stored_space.decode() + if stored_space is not None and str(stored_space) != self.coordinate_space: + raise ValueError("coordinate space differs between manifest and signals.h5") with pysam.FastaFile(str(self._fasta_path)) as fasta: if tuple(fasta.references) != tuple(tx.transcript_id for tx in self.transcripts): raise ValueError("transcript order differs between metadata and transcript FASTA") @@ -273,9 +299,17 @@ def _load_exons(self) -> None: def get_genomic_blocks( self, transcript_id: str, start: int, end: int ) -> tuple[GenomicBlock, ...]: - """Map one transcript interval to compact genomic exon blocks.""" + """Map one locus interval to compact ascending genomic blocks.""" tx, _ = self._slice(transcript_id, start, end) + if self.coordinate_space == "gene": + if tx.strand == "+": + genomic_start = tx.genomic_start + start + genomic_end = tx.genomic_start + end + else: + genomic_start = tx.genomic_end - end + genomic_end = tx.genomic_end - start + return (GenomicBlock(tx.chromosome, genomic_start, genomic_end),) if self._exons_by_transcript is None: self._load_exons() assert self._exons_by_transcript is not None @@ -296,6 +330,58 @@ def get_genomic_blocks( raise ValueError(f"exon mapping does not cover {transcript_id}:{start}-{end}") return tuple(blocks) + def coordinate_to_genome(self, transcript_id: str, pos: int) -> tuple[str, int, str]: + """Map one selected-coordinate-space base to a genomic base.""" + + tx, _ = self._slice(transcript_id, pos, pos + 1) + if self.coordinate_space == "gene": + genomic = ( + tx.genomic_start + pos + if tx.strand == "+" + else tx.genomic_end - 1 - pos + ) + return tx.chromosome, genomic, tx.strand + if self._exons_by_transcript is None: + self._load_exons() + assert self._exons_by_transcript is not None + for exon in self._exons_by_transcript.get(transcript_id, []): + if exon["tx_start"] <= pos < exon["tx_end"]: + offset = pos - exon["tx_start"] + genomic = ( + exon["genomic_start"] + offset + if tx.strand == "+" + else exon["genomic_end"] - 1 - offset + ) + return tx.chromosome, genomic, tx.strand + raise ValueError(f"exon mapping does not cover {transcript_id}:{pos}") + + def genome_to_coordinate(self, transcript_id: str, chromosome: str, pos: int) -> int | None: + """Map one genomic base into the selected coordinate space, if represented.""" + + tx = self.get_transcript(transcript_id) + if chromosome != tx.chromosome: + return None + if self.coordinate_space == "gene": + if not tx.genomic_start <= pos < tx.genomic_end: + return None + return ( + pos - tx.genomic_start + if tx.strand == "+" + else tx.genomic_end - 1 - pos + ) + if self._exons_by_transcript is None: + self._load_exons() + assert self._exons_by_transcript is not None + for exon in self._exons_by_transcript.get(transcript_id, []): + if exon["genomic_start"] <= pos < exon["genomic_end"]: + offset = ( + pos - exon["genomic_start"] + if tx.strand == "+" + else exon["genomic_end"] - 1 - pos + ) + return exon["tx_start"] + offset + return None + def close(self) -> None: if self._h5 is not None: self._h5.close() diff --git a/src/transcriptml/rbpnet/fasta.py b/src/transcriptml/rbpnet/fasta.py index 06f3554..adbfcc6 100644 --- a/src/transcriptml/rbpnet/fasta.py +++ b/src/transcriptml/rbpnet/fasta.py @@ -1,4 +1,4 @@ -"""FASTA validation and mature-transcript sequence extraction.""" +"""FASTA validation and transcript-oriented locus sequence extraction.""" from __future__ import annotations @@ -63,13 +63,20 @@ def retain_fasta_transcripts( def transcript_sequence(fasta: pysam.FastaFile, tx: Transcript) -> str: - """Assemble one mature transcript in transcript 5-prime to 3-prime order.""" - - chunks = [] - for exon in tx.exons: - chunk = fasta.fetch(exon.chrom, exon.start, exon.end) - chunks.append(chunk if tx.strand == "+" else reverse_complement(chunk)) - sequence = "".join(chunks).upper() + """Extract one selected locus in annotated RNA 5-prime to 3-prime order.""" + + if tx.coordinate_space == "gene": + assert tx.genomic_start is not None and tx.genomic_end is not None + sequence = fasta.fetch(tx.chrom, tx.genomic_start, tx.genomic_end) + if tx.strand == "-": + sequence = reverse_complement(sequence) + else: + chunks = [] + for exon in tx.exons: + chunk = fasta.fetch(exon.chrom, exon.start, exon.end) + chunks.append(chunk if tx.strand == "+" else reverse_complement(chunk)) + sequence = "".join(chunks) + sequence = sequence.upper() if len(sequence) != tx.length: raise RuntimeError(f"sequence length mismatch for {tx.transcript_id}") return sequence @@ -82,7 +89,7 @@ def write_transcript_fasta( *, progress: bool = True, ) -> dict: - """Write indexed mature-transcript FASTA and return sequence QC.""" + """Write indexed transcript-oriented locus FASTA and return sequence QC.""" path = Path(path) ensure_fasta_index(genome_fasta) diff --git a/src/transcriptml/rbpnet/preprocessing.py b/src/transcriptml/rbpnet/preprocessing.py index 0a0c127..d437f8c 100644 --- a/src/transcriptml/rbpnet/preprocessing.py +++ b/src/transcriptml/rbpnet/preprocessing.py @@ -16,6 +16,7 @@ from transcriptml import __version__ from transcriptml.progress import log_progress from transcriptml.rbpnet.annotation import parse_gtf +from transcriptml.rbpnet.coordinates import COORDINATE_SPACES from transcriptml.rbpnet.fasta import retain_fasta_transcripts, write_transcript_fasta from transcriptml.rbpnet.serialization import write_exons, write_json, write_metadata, write_regions from transcriptml.rbpnet.signals import ExonBinIndex, create_signal_store, extract_bam_to_store, write_ip_pooled @@ -34,13 +35,14 @@ class Sample: @dataclass(frozen=True) class PipelineConfig: - """Configuration for canonical transcript-space eCLIP preprocessing.""" + """Configuration for canonical transcript-oriented eCLIP preprocessing.""" genome_fasta: Path gtf: Path sminput: Sample ips: tuple[Sample, ...] output_dir: Path + coordinate_space: str = "mature_transcript" read1_rna_strand: str = "opposite" min_mapq: int = 1 exclude_duplicates: bool = True @@ -72,9 +74,9 @@ def _input_record(path: Path) -> dict: def preprocess_eclip(config: PipelineConfig) -> dict: - """Create a reusable canonical transcript-space eCLIP experiment. + """Create a reusable canonical transcript-oriented eCLIP experiment. - The HDF5 track is concatenated transcript space and stays lazy on read. + The HDF5 track is concatenated selected-locus space and stays lazy on read. This stage deliberately performs no peak calling or region selection. """ @@ -85,6 +87,10 @@ def preprocess_eclip(config: PipelineConfig) -> dict: raise ValueError("sample roles must be one 'sminput' followed by one or more 'ip' samples") if config.read1_rna_strand not in {"opposite", "same", "unstranded"}: raise ValueError("read1_rna_strand must be opposite, same, or unstranded") + if config.coordinate_space not in COORDINATE_SPACES: + raise ValueError( + f"coordinate_space must be one of {', '.join(COORDINATE_SPACES)}" + ) if config.min_mapq < 0: raise ValueError("min_mapq must be non-negative") sample_names = [config.sminput.name] + [sample.name for sample in config.ips] @@ -93,7 +99,11 @@ def preprocess_eclip(config: PipelineConfig) -> dict: log_progress(f"rbpnet preprocess: prepare {config.output_dir}", enabled=config.progress) _prepare_output(config.output_dir, config.overwrite) - all_transcripts = parse_gtf(config.gtf, progress=config.progress) + all_transcripts = parse_gtf( + config.gtf, + coordinate_space=config.coordinate_space, + progress=config.progress, + ) transcripts, contig_filter_qc = retain_fasta_transcripts(config.genome_fasta, all_transcripts) skipped = contig_filter_qc["transcripts_skipped_missing_fasta_contig"] if skipped: @@ -169,17 +179,23 @@ def preprocess_eclip(config: PipelineConfig) -> dict: for tx in transcripts: for region in tx.regions: region_counts[region.label] = region_counts.get(region.label, 0) + region.end - region.start + assignment = ( + "unique strand-compatible selected gene at read1 5-prime aligned base, with all " + "reference-consuming CIGAR operations contained in the full gene span; intronic " + "alignments and non-annotated splice junctions inside that span are permitted" + if config.coordinate_space == "gene" + else "unique strand-compatible mature transcript at read1 5-prime aligned base, " + "with all aligned CIGAR segments compatible with selected exons and junctions" + ) qc = { "format_version": "1", "pipeline_version": __version__, "configuration": { + "coordinate_space": config.coordinate_space, "read1_rna_strand": config.read1_rna_strand, "min_mapq": config.min_mapq, "exclude_duplicates": config.exclude_duplicates, - "assignment": ( - "unique strand-compatible mature transcript at read1 5-prime aligned base, " - "with all aligned CIGAR segments compatible with selected exons and junctions" - ), + "assignment": assignment, }, "annotation": { **contig_filter_qc, @@ -187,6 +203,15 @@ def preprocess_eclip(config: PipelineConfig) -> dict: "genes": len({tx.gene_id for tx in transcripts}), "exons": sum(len(tx.exons) for tx in transcripts), "transcriptome_bases": sum(tx.length for tx in transcripts), + "coordinate_space": config.coordinate_space, + "coordinate_space_bases": sum(tx.length for tx in transcripts), + "exonic_bases": sum( + sum(exon.end - exon.start for exon in tx.exons) for tx in transcripts + ), + "intronic_bases": sum( + sum(region.end - region.start for region in tx.regions if region.label == "intron") + for tx in transcripts + ), "chromosomes": sorted({tx.chrom for tx in transcripts}), "region_bases": dict(sorted(region_counts.items())), "sminput_transcripts_nonzero": int(np.count_nonzero(sample_counts[0])), @@ -201,8 +226,10 @@ def preprocess_eclip(config: PipelineConfig) -> dict: "format": "transcriptml-rbpnet-experiment", "format_version": "1", "created_at": datetime.now(timezone.utc).isoformat(), + "coordinate_space": config.coordinate_space, "coordinate_system": ( - "all intervals are 0-based, half-open; sequences/tracks are transcript 5-prime to 3-prime" + "all intervals are 0-based, half-open in the selected locus coordinate space; " + "sequences and tracks run annotated RNA 5-prime to 3-prime" ), "inputs": { "genome_fasta": _input_record(config.genome_fasta), @@ -231,7 +258,7 @@ def preprocess_eclip(config: PipelineConfig) -> dict: "cpm_formula": "window_count / effective_library_size * 1e6", "sample_denominator_field": "samples[].effective_library_size", "effective_library_size_definition": ( - "retained read1 5-prime events used to construct the transcript-space signal track" + "retained read1 5-prime events used to construct the selected-coordinate-space signal track" ), "pooled_ip_denominator": "sum of effective_library_size over IP source samples", }, diff --git a/src/transcriptml/rbpnet/selection.py b/src/transcriptml/rbpnet/selection.py index 7009eb7..7dd3034 100644 --- a/src/transcriptml/rbpnet/selection.py +++ b/src/transcriptml/rbpnet/selection.py @@ -8,7 +8,7 @@ import json import math from collections import Counter -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Iterable, Iterator @@ -21,7 +21,8 @@ from transcriptml.rbpnet.experiment import ProcessedECLIPDataset from transcriptml.rbpnet.windows import REGION_TYPES, summarize_regions -SELECTION_STRATEGIES = ("original_rbpnet", "yeo_2026", "peak_gray_negative") +SELECTION_STRATEGIES = ("original_rbpnet", "broad_coverage", "peak_gray_negative") +POISSON_NULLS = ("ip_locus_density", "sminput") @dataclass(frozen=True) @@ -29,8 +30,9 @@ class SelectionConfig: """Configuration for selecting eligible experimental loci. Defaults for ``original_rbpnet`` reproduce the published Horlacher et al. - candidate rules. Defaults for the other two strategies are transparent - starting points and should be reviewed for each assay. + candidate rules. ``broad_coverage`` defaults to replicate-wise + IP+SMInput >= 6. Peak/gray/negative thresholds remain configurable + scientific starting points. """ processed_dir: Path @@ -45,12 +47,14 @@ class SelectionConfig: original_min_count: int = 8 original_min_height: int = 2 original_advance: int = 50 + poisson_null: str = "ip_locus_density" + sminput_poisson_pseudocount: float = 1.0 # Broad measured-window selector. - min_total_count: int = 8 - min_sminput_count: int = 1 - min_ip_count: int = 1 + min_total_count: int | None = None + min_sminput_count: int = 0 + min_ip_count: int = 0 min_sminput_tpm: float = 0.0 - replicate_mode: str = "combined" + replicate_mode: str = "per_ip" # Peak / gray / confident-negative selector. peak_fdr: float = 0.05 peak_min_log2_ratio: float = 1.0 @@ -116,6 +120,7 @@ def _manifest_schema(ds: ProcessedECLIPDataset, metadata: dict[bytes, bytes]) -> pa.field("transcript_id", pa.string()), pa.field("chromosome", pa.string()), pa.field("strand", pa.string()), + pa.field("coordinate_space", pa.string()), pa.field("transcript_anchor", pa.int64()), pa.field("selection_start", pa.int64()), pa.field("selection_end", pa.int64()), @@ -148,6 +153,7 @@ def _manifest_schema(ds: ProcessedECLIPDataset, metadata: dict[bytes, bytes]) -> fields.extend(pa.field(f"max_{sample.name}_5pend", pa.int64()) for sample in ds.samples) fields.extend([ pa.field("max_ip_pooled_5pend", pa.int64()), + pa.field("selection_null_mean", pa.float64()), pa.field("selection_pvalue", pa.float64()), pa.field("selection_qvalue", pa.float64()), pa.field("source_min_enrichment_pvalue", pa.float64()), @@ -177,6 +183,7 @@ def _base_manifest_row( replicate_id: str = "", source_window_count: int = 1, anchor: int | None = None, + selection_null_mean: float = math.nan, selection_pvalue: float = math.nan, selection_qvalue: float = math.nan, enrichment_pvalue: float = math.nan, @@ -194,6 +201,7 @@ def _base_manifest_row( "transcript_id": tx.transcript_id, "chromosome": tx.chromosome, "strand": tx.strand, + "coordinate_space": ds.coordinate_space, "transcript_anchor": anchor, "selection_start": start, "selection_end": end, @@ -210,6 +218,7 @@ def _base_manifest_row( "total_ip_sminput_count": int(source["total_ip_sminput_count"]), "log2_ip_pooled_vs_sminput": float(source["log2_ip_pooled_vs_sminput"]), "max_ip_pooled_5pend": int(source["max_ip_pooled_5pend"]), + "selection_null_mean": float(selection_null_mean), "selection_pvalue": float(selection_pvalue), "selection_qvalue": float(selection_qvalue), "source_min_enrichment_pvalue": float(enrichment_pvalue), @@ -221,8 +230,10 @@ def _base_manifest_row( "group_chromosome": tx.chromosome, } for region_type in REGION_TYPES: - row[f"region_{region_type}_nt"] = int(source[f"region_{region_type}_nt"]) - row[f"region_{region_type}_fraction"] = float(source[f"region_{region_type}_fraction"]) + row[f"region_{region_type}_nt"] = int(source.get(f"region_{region_type}_nt", 0)) + row[f"region_{region_type}_fraction"] = float( + source.get(f"region_{region_type}_fraction", 0.0) + ) for sample in ds.samples: row[f"{sample.name}_count"] = int(source[f"{sample.name}_count"]) row[f"{sample.name}_cpm"] = float(source[f"{sample.name}_cpm"]) @@ -318,8 +329,11 @@ def _original_rows( tx_id = row["transcript_id"] if tx_id != current_tx: tx = ds.get_transcript(tx_id) - transcript_count = int(ds.get_pooled_ip_profile(tx_id, 0, tx.length).sum(dtype=np.uint64)) - mu = transcript_count / tx.length * int(row["window_length"]) + if config.poisson_null == "ip_locus_density": + transcript_count = int( + ds.get_pooled_ip_profile(tx_id, 0, tx.length).sum(dtype=np.uint64) + ) + mu = transcript_count / tx.length * int(row["window_length"]) current_tx = tx_id next_start = 0 start = int(row["tx_start"]) @@ -327,7 +341,18 @@ def _original_rows( continue count = int(row["ip_pooled_count"]) height = int(row["max_ip_pooled_5pend"]) - pvalue = float(poisson.sf(count - 1, mu)) + if config.poisson_null == "sminput": + input_count = int(row[f"{ds.sminput_sample.name}_count"]) + exposure_ratio = ( + ds.pooled_ip_effective_library_size + / int(ds.sminput_sample.effective_library_size) + ) + row_mu = ( + input_count + config.sminput_poisson_pseudocount + ) * exposure_ratio + else: + row_mu = mu + pvalue = float(poisson.sf(count - 1, row_mu)) if ( pvalue < config.original_min_pvalue and count >= config.original_min_count @@ -338,6 +363,7 @@ def _original_rows( row, strategy="original_rbpnet", state="candidate", + selection_null_mean=row_mu, selection_pvalue=pvalue, ) next_start = start + config.original_advance @@ -345,7 +371,7 @@ def _original_rows( reporter.close() -def _yeo_rows( +def _broad_coverage_rows( config: SelectionConfig, ds: ProcessedECLIPDataset, windows_path: Path, @@ -373,7 +399,7 @@ def _yeo_rows( yield _base_manifest_row( ds, row, - strategy="yeo_2026", + strategy="broad_coverage", state="measured", replicate_id="", ) @@ -388,7 +414,7 @@ def _yeo_rows( yield _base_manifest_row( ds, row, - strategy="yeo_2026", + strategy="broad_coverage", state="measured", replicate_id=sample.name, ) @@ -569,11 +595,16 @@ def _validate_config(config: SelectionConfig) -> None: value = float(getattr(config, name)) if not 0 < value <= 1: raise ValueError(f"{name} must be in (0, 1]") + assert config.min_total_count is not None if min(config.original_min_count, config.original_min_height, config.min_total_count, config.min_sminput_count, config.min_ip_count, config.stitch_gap) < 0: raise ValueError("count thresholds and stitch_gap must be non-negative") if config.min_sminput_tpm < 0: raise ValueError("min_sminput_tpm must be non-negative") + if config.poisson_null not in POISSON_NULLS: + raise ValueError(f"poisson_null must be one of {', '.join(POISSON_NULLS)}") + if config.sminput_poisson_pseudocount <= 0: + raise ValueError("sminput_poisson_pseudocount must be positive") if config.replicate_mode not in {"combined", "per_ip"}: raise ValueError("replicate_mode must be combined or per_ip") @@ -597,6 +628,9 @@ def _validate_scan_dataset( ) if int(scan_metadata.get("ip_pooled_effective_library_size", -1)) != ds.pooled_ip_effective_library_size: raise ValueError("window scan pooled-IP library size does not match the processed experiment") + scan_coordinate_space = scan_metadata.get("coordinate_space", "mature_transcript") + if scan_coordinate_space != ds.coordinate_space: + raise ValueError("window scan coordinate space does not match the processed experiment") required_columns = { "transcript_id", "tx_start", "tx_end", "window_length", "region_type", "sminput_tpm", "ip_pooled_count", "ip_pooled_cpm", @@ -615,6 +649,11 @@ def _validate_scan_dataset( def select_regions(config: SelectionConfig) -> dict: """Select biological loci and write a versioned lightweight manifest.""" + if config.min_total_count is None: + config = replace( + config, + min_total_count=6 if config.strategy == "broad_coverage" else 8, + ) _validate_config(config) windows_path = _resolve_parquet(config.windows) scan_metadata = _scan_metadata(windows_path) @@ -639,6 +678,7 @@ def select_regions(config: SelectionConfig) -> dict: "format_version": "1", "strategy": config.strategy, "source_processed_dir": str(config.processed_dir.resolve()), + "coordinate_space": ds.coordinate_space, "source_windows": str(windows_path.resolve()), "window_scan": scan_metadata, "configuration": { @@ -646,12 +686,29 @@ def select_regions(config: SelectionConfig) -> dict: for key, value in config.__dict__.items() if key not in {"processed_dir", "windows", "output_prefix", "progress"} }, + "poisson_null_formula": ( + "pooled_IP_transcript_or_gene_total / locus_length * window_length" + if config.strategy == "original_rbpnet" + and config.poisson_null == "ip_locus_density" + else "(SMInput_window_count + sminput_poisson_pseudocount) * " + "(pooled_IP_effective_library_size / SMInput_effective_library_size)" + if config.strategy == "original_rbpnet" + else None + ), "statistical_notes": ( - "original_rbpnet uses a one-sided Poisson test against the transcript-level pooled-IP rate" + ( + "original_rbpnet uses a one-sided uncorrected Poisson test; " + "poisson_null=ip_locus_density is the published pooled-IP locus-density null" + if config.poisson_null == "ip_locus_density" + else "original_rbpnet uses a one-sided uncorrected experimental SMInput null: " + "mu=(SMInput_window_count+sminput_poisson_pseudocount)*" + "(pooled_IP_effective_library_size/SMInput_effective_library_size)" + ) if config.strategy == "original_rbpnet" else "peak_gray_negative uses exact conditional binomial tails and BH correction over adequately measured windows" if config.strategy == "peak_gray_negative" - else "yeo_2026 applies coverage thresholds only and performs no peak test" + else "broad_coverage applies coverage thresholds only and performs no peak test; " + "it is Yeo-inspired but is not an exact Skipper window-generation preset" ), } schema = _manifest_schema( @@ -660,8 +717,8 @@ def select_regions(config: SelectionConfig) -> dict: ) if config.strategy == "original_rbpnet": rows: Iterable[dict] = _original_rows(config, ds, windows_path, scan_metadata) - elif config.strategy == "yeo_2026": - rows = _yeo_rows(config, ds, windows_path) + elif config.strategy == "broad_coverage": + rows = _broad_coverage_rows(config, ds, windows_path) else: rows = _peak_gray_negative_rows(config, ds, windows_path, scan_metadata) diff --git a/src/transcriptml/rbpnet/serialization.py b/src/transcriptml/rbpnet/serialization.py index da696dd..87e391d 100644 --- a/src/transcriptml/rbpnet/serialization.py +++ b/src/transcriptml/rbpnet/serialization.py @@ -14,7 +14,7 @@ def calculate_tpm(raw_counts: np.ndarray, transcripts: list[Transcript]) -> np.ndarray: - """Calculate length-normalized TPM from retained transcript event counts.""" + """Calculate length-normalized TPM in the selected coordinate space.""" lengths_kb = np.asarray([tx.length / 1000.0 for tx in transcripts], dtype=np.float64) rates = raw_counts.astype(np.float64) / lengths_kb @@ -36,8 +36,8 @@ def write_metadata( tpm = calculate_tpm(sample_counts[sminput_index], transcripts) fields = [ "transcript_id", "gene_id", "gene_name", "transcript_name", "transcript_type", - "chrom", "strand", "transcript_length", "signal_offset", "sm_input_raw_count", - "sm_input_tpm", + "chrom", "strand", "coordinate_space", "genomic_start", "genomic_end", + "transcript_length", "signal_offset", "sm_input_raw_count", "sm_input_tpm", ] + [f"{name}_raw_5p_count" for name in sample_names] + ["region_annotations"] with Path(path).open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n") @@ -57,6 +57,9 @@ def write_metadata( "transcript_type": tx.transcript_type, "chrom": tx.chrom, "strand": tx.strand, + "coordinate_space": tx.coordinate_space, + "genomic_start": tx.genomic_start, + "genomic_end": tx.genomic_end, "transcript_length": tx.length, "signal_offset": tx.offset, "sm_input_raw_count": int(sample_counts[sminput_index][index]), diff --git a/src/transcriptml/rbpnet/signals.py b/src/transcriptml/rbpnet/signals.py index 40c84cc..8564c7a 100644 --- a/src/transcriptml/rbpnet/signals.py +++ b/src/transcriptml/rbpnet/signals.py @@ -74,8 +74,45 @@ def within_one_exon(start: int, end: int) -> bool: return has_aligned_bases +def _alignment_is_gene_compatible(read, start: int, end: int) -> bool: + """Check that every reference-consuming CIGAR segment stays in one gene span. + + Full-gene coordinates contain introns, so ``N`` operations need not match + the selected mature-transcript junctions. They do, however, have to remain + completely within the selected gene locus. + """ + + if read.reference_start is None or not read.cigartuples: + return False + reference_pos = read.reference_start + has_aligned_bases = False + for operation, length in read.cigartuples: + if length <= 0: + return False + if operation in { + pysam.CMATCH, pysam.CEQUAL, pysam.CDIFF, pysam.CDEL, pysam.CREF_SKIP, + }: + next_pos = reference_pos + length + if reference_pos < start or next_pos > end: + return False + if operation in {pysam.CMATCH, pysam.CEQUAL, pysam.CDIFF}: + has_aligned_bases = True + reference_pos = next_pos + elif operation in {pysam.CINS, pysam.CSOFT_CLIP, pysam.CHARD_CLIP, pysam.CPAD}: + continue + else: + return False + return has_aligned_bases + + def alignment_is_transcript_compatible(read, transcript: Transcript) -> bool: - """Return whether an alignment follows selected exons and exact junctions.""" + """Return whether an alignment is compatible with the selected coordinates.""" + + if transcript.coordinate_space == "gene": + assert transcript.genomic_start is not None and transcript.genomic_end is not None + return _alignment_is_gene_compatible( + read, transcript.genomic_start, transcript.genomic_end + ) exons = tuple(sorted((exon.start, exon.end) for exon in transcript.exons)) introns = frozenset((left[1], right[0]) for left, right in zip(exons, exons[1:])) @@ -83,7 +120,7 @@ def alignment_is_transcript_compatible(read, transcript: Transcript) -> bool: class ExonBinIndex: - """Small-memory genomic point index for candidate transcript exons.""" + """Small-memory genomic point index for candidate coordinate-space loci.""" def __init__(self, transcripts: list[Transcript], bin_size: int = 16_384): self.transcripts = transcripts @@ -91,19 +128,27 @@ def __init__(self, transcripts: list[Transcript], bin_size: int = 16_384): self._bins: dict[tuple[str, str, int], list[tuple[int, int, int, int]]] = {} self._exons: list[tuple[tuple[int, int], ...]] = [] self._introns: list[frozenset[tuple[int, int]]] = [] + self._gene_spans: list[tuple[int, int] | None] = [] for tx_index, tx in enumerate(transcripts): genomic_exons = tuple(sorted((exon.start, exon.end) for exon in tx.exons)) self._exons.append(genomic_exons) self._introns.append(frozenset( (left[1], right[0]) for left, right in zip(genomic_exons, genomic_exons[1:]) )) - for exon in tx.exons: - record = (exon.start, exon.end, tx_index, exon.tx_start) - for bin_id in range(exon.start // bin_size, (exon.end - 1) // bin_size + 1): - self._bins.setdefault((exon.chrom, exon.strand, bin_id), []).append(record) + if tx.coordinate_space == "gene": + assert tx.genomic_start is not None and tx.genomic_end is not None + self._gene_spans.append((tx.genomic_start, tx.genomic_end)) + intervals = [(tx.genomic_start, tx.genomic_end, 0)] + else: + self._gene_spans.append(None) + intervals = [(exon.start, exon.end, exon.tx_start) for exon in tx.exons] + for start, end, tx_start in intervals: + record = (start, end, tx_index, tx_start) + for bin_id in range(start // bin_size, (end - 1) // bin_size + 1): + self._bins.setdefault((tx.chrom, tx.strand, bin_id), []).append(record) def query(self, chrom: str, pos: int, strand: str | None) -> list[tuple[int, int]]: - """Return ``(transcript_index, transcript_position)`` exon hits.""" + """Return ``(transcript_index, coordinate_position)`` locus hits.""" strands = ("+", "-") if strand is None else (strand,) matches: list[tuple[int, int]] = [] @@ -117,8 +162,11 @@ def query(self, chrom: str, pos: int, strand: str | None) -> list[tuple[int, int return matches def alignment_is_compatible(self, read, transcript_index: int) -> bool: - """Check a read against precomputed exon and intron intervals.""" + """Check a read against the selected mature-transcript or gene span.""" + gene_span = self._gene_spans[transcript_index] + if gene_span is not None: + return _alignment_is_gene_compatible(read, *gene_span) return _alignment_is_compatible( read, self._exons[transcript_index], self._introns[transcript_index] ) @@ -163,15 +211,22 @@ def create_signal_store( sample_names: list[str], sample_roles: list[str], ) -> h5py.File: - """Create the canonical concatenated transcript-space HDF5 store.""" + """Create the canonical concatenated locus-coordinate HDF5 store.""" total_length = sum(tx.length for tx in transcripts) if total_length == 0: - raise ValueError("annotation has zero mature-transcript bases") + raise ValueError("annotation has zero coordinate-space bases") + coordinate_spaces = {tx.coordinate_space for tx in transcripts} + if len(coordinate_spaces) != 1: + raise ValueError("all loci in one signal store must use the same coordinate space") + coordinate_space = coordinate_spaces.pop() store = h5py.File(path, "w") store.attrs["format"] = "transcriptml-rbpnet-signals" store.attrs["format_version"] = "1" - store.attrs["coordinate_system"] = "0-based half-open transcript coordinates" + store.attrs["coordinate_space"] = coordinate_space + store.attrs["coordinate_system"] = ( + "0-based half-open coordinates in annotated 5-prime to 3-prime orientation" + ) strings = h5py.string_dtype("utf-8") store.create_dataset( "transcript_ids", diff --git a/src/transcriptml/rbpnet/windows.py b/src/transcriptml/rbpnet/windows.py index 7127170..7b9fd94 100644 --- a/src/transcriptml/rbpnet/windows.py +++ b/src/transcriptml/rbpnet/windows.py @@ -18,7 +18,7 @@ from transcriptml.progress import ProgressReporter, log_progress from transcriptml.rbpnet.experiment import ProcessedECLIPDataset, RegionRecord -REGION_TYPES = ("5putr", "cds", "3putr", "noncoding_exon") +REGION_TYPES = ("5putr", "cds", "3putr", "noncoding_exon", "intron") @dataclass(frozen=True) @@ -179,6 +179,7 @@ def scan_windows(config: WindowScanConfig) -> dict: "format": "transcriptml-rbpnet-window-scan", "format_version": "1", "source_processed_dir": str(config.processed_dir.resolve()), + "coordinate_space": ds.coordinate_space, "window_size": config.window_size, "stride": config.stride, "min_sminput_tpm": config.min_sminput_tpm, diff --git a/tests/test_rbpnet.py b/tests/test_rbpnet.py index 92a2b4a..dbce0cd 100644 --- a/tests/test_rbpnet.py +++ b/tests/test_rbpnet.py @@ -8,12 +8,19 @@ import pyarrow.parquet as pq import pysam import pytest +from scipy.stats import poisson from transcriptml.data.encoding import encode_rna_sequence -from transcriptml.rbpnet.bundle import RBPNetBundleConfig, load_rbpnet_bundle, make_rbpnet_bundle +from transcriptml.rbpnet.bundle import ( + RBPNetBundleConfig, + _shifted_materialized_interval, + jitter_crop_offset, + load_rbpnet_bundle, + make_rbpnet_bundle, +) from transcriptml.rbpnet.coordinates import Exon, Region, Transcript, annotate_regions from transcriptml.rbpnet.experiment import ProcessedECLIPDataset -from transcriptml.rbpnet.fasta import transcript_sequence +from transcriptml.rbpnet.fasta import reverse_complement, transcript_sequence from transcriptml.rbpnet.preprocessing import PipelineConfig, Sample, preprocess_eclip from transcriptml.rbpnet.selection import SelectionConfig, load_selection_manifest, select_regions from transcriptml.rbpnet.serialization import calculate_tpm @@ -59,6 +66,52 @@ def test_transcript_coordinates_regions_and_minus_sequence(tmp_path): assert transcript_sequence(fasta, minus) == "CCGGTACGT" +def _gene_tx(strand="+", *, coding=False): + tx = Transcript( + f"gene_tx_{strand}", f"gene_{strand}", "", "", + "protein_coding" if coding else "lncRNA", "chr1", strand, + [Exon("chr1", 100, 110, strand), Exon("chr1", 200, 210, strand)], + coordinate_space="gene", + ) + if coding: + tx.feature_intervals = {"CDS": [(105, 110), (200, 205)]} + tx.finalize() + tx.regions = annotate_regions(tx) + return tx + + +def test_gene_coordinates_regions_sequence_and_round_trips(tmp_path): + plus = _gene_tx("+", coding=True) + minus = _gene_tx("-", coding=True) + expected_regions = [ + Region(0, 5, "5putr"), + Region(5, 10, "cds"), + Region(10, 100, "intron"), + Region(100, 105, "cds"), + Region(105, 110, "3putr"), + ] + assert plus.length == minus.length == 110 + assert plus.regions == minus.regions == expected_regions + assert plus.genome_to_transcript("chr1", 150) == 50 + assert minus.genome_to_transcript("chr1", 150) == 59 + assert plus.transcript_to_genome(50) == ("chr1", 150, "+") + assert minus.transcript_to_genome(59) == ("chr1", 150, "-") + for tx in (plus, minus): + for coordinate in (0, 9, 10, 50, 99, 100, 109): + chrom, genomic, strand = tx.transcript_to_genome(coordinate) + assert strand == tx.strand + assert tx.genome_to_transcript(chrom, genomic) == coordinate + + fasta_path = tmp_path / "genome.fa" + genomic = ("ACGT" * 80)[:300] + fasta_path.write_text(">chr1\n" + genomic + "\n") + pysam.faidx(str(fasta_path)) + with pysam.FastaFile(str(fasta_path)) as fasta: + expected = genomic[100:210] + assert transcript_sequence(fasta, plus) == expected + assert transcript_sequence(fasta, minus) == reverse_complement(expected) + + def _alignment(start, cigar, reverse=False): read = pysam.AlignedSegment() read.query_name = "compatibility" @@ -96,6 +149,14 @@ def test_alignment_compatibility_and_library_orientation(): assert not alignment_is_transcript_compatible(intronic, _junction_tx()) assert not alignment_is_transcript_compatible(wrong_junction, _junction_tx()) + gene = _gene_tx("+") + assert alignment_is_transcript_compatible(_alignment(102, ((pysam.CMATCH, 6),)), gene) + assert alignment_is_transcript_compatible(_alignment(140, ((pysam.CMATCH, 20),)), gene) + assert alignment_is_transcript_compatible(junction, gene) + assert not alignment_is_transcript_compatible( + _alignment(205, ((pysam.CMATCH, 10),)), gene + ) + def test_bam_assignment_reports_transcript_incompatibility(tmp_path): tx = _junction_tx() @@ -127,6 +188,47 @@ def test_bam_assignment_reports_transcript_incompatibility(tmp_path): assert counts.tolist() == [1] +def test_gene_space_bam_assignment_retains_intronic_and_spliced_reads(tmp_path): + gene = _gene_tx("+") + mature = _junction_tx("+") + bam_path = tmp_path / "gene_reads.bam" + header = {"HD": {"VN": "1.6", "SO": "coordinate"}, "SQ": [{"SN": "chr1", "LN": 1000}]} + records = [ + ("exonic", 101, ((pysam.CMATCH, 4),), 4), + ("junction", 105, ((pysam.CMATCH, 5), (pysam.CREF_SKIP, 90), (pysam.CMATCH, 5)), 10), + ("intronic", 150, ((pysam.CMATCH, 4),), 4), + ] + with pysam.AlignmentFile(bam_path, "wb", header=header) as bam: + for name, start, cigar, query_length in records: + read = pysam.AlignedSegment() + read.query_name = name + read.query_sequence = "A" * query_length + read.flag = 81 # read1, reverse; opposite-strand RNA is plus + read.reference_id = 0 + read.reference_start = start + read.mapping_quality = 60 + read.cigar = cigar + read.query_qualities = pysam.qualitystring_to_array("I" * query_length) + bam.write(read) + pysam.index(str(bam_path)) + + with create_signal_store(tmp_path / "gene.h5", [gene], ["ip"], ["ip"]) as store: + gene_qc, gene_counts = extract_bam_to_store( + bam_path, 0, store["counts"], [gene], ExonBinIndex([gene]), + "opposite", 1, True, tmp_path, progress=False, + ) + with create_signal_store(tmp_path / "mature.h5", [mature], ["ip"], ["ip"]) as store: + mature_qc, mature_counts = extract_bam_to_store( + bam_path, 0, store["counts"], [mature], ExonBinIndex([mature]), + "opposite", 1, True, tmp_path, progress=False, + ) + assert gene_qc["retained"] == 3 + assert mature_qc["retained"] == 2 + assert mature_qc["no_compatible_transcript"] == 1 + assert gene_counts.tolist() == [3] + assert mature_counts.tolist() == [2] + + def test_preprocess_pipeline_manifest_tpm_effective_sizes_and_missing_contig(tmp_path): fasta = tmp_path / "genome.fa" fasta.write_text(">chr1\n" + "ACGT" * 20 + "\n") @@ -174,6 +276,114 @@ def test_preprocess_pipeline_manifest_tpm_effective_sizes_and_missing_contig(tmp assert int(ds.get_pooled_ip_profile("t1", 0, 20).sum()) == 2 +def test_full_gene_preprocessing_reader_orientation_introns_and_mature_regression(tmp_path): + fasta = tmp_path / "genome.fa" + genomic = ("ACGT" * 300)[:1000] + fasta.write_text(">chr1\n" + genomic + "\n") + gtf = tmp_path / "annotation.gtf" + gtf.write_text( + 'chr1\ttest\ttranscript\t101\t210\t.\t+\t.\tgene_id "gp"; transcript_id "tp"; transcript_type "lncRNA";\n' + 'chr1\ttest\texon\t101\t110\t.\t+\t.\tgene_id "gp"; transcript_id "tp"; exon_number 1;\n' + 'chr1\ttest\texon\t201\t210\t.\t+\t.\tgene_id "gp"; transcript_id "tp"; exon_number 2;\n' + 'chr1\ttest\ttranscript\t401\t510\t.\t-\t.\tgene_id "gm"; transcript_id "tm"; transcript_type "lncRNA";\n' + 'chr1\ttest\texon\t401\t410\t.\t-\t.\tgene_id "gm"; transcript_id "tm"; exon_number 2;\n' + 'chr1\ttest\texon\t501\t510\t.\t-\t.\tgene_id "gm"; transcript_id "tm"; exon_number 1;\n' + ) + bam_path = tmp_path / "reads.bam" + header = {"HD": {"VN": "1.6", "SO": "coordinate"}, "SQ": [{"SN": "chr1", "LN": 1000}]} + records = [ + ("plus_junction", 105, ((pysam.CMATCH, 5), (pysam.CREF_SKIP, 90), (pysam.CMATCH, 5)), 10, 81), + ("plus_intron", 150, ((pysam.CMATCH, 4),), 4, 81), + ("minus_junction", 405, ((pysam.CMATCH, 5), (pysam.CREF_SKIP, 90), (pysam.CMATCH, 5)), 10, 65), + ("minus_intron", 450, ((pysam.CMATCH, 4),), 4, 65), + ] + with pysam.AlignmentFile(bam_path, "wb", header=header) as bam: + for name, start, cigar, query_length, flag in records: + read = pysam.AlignedSegment() + read.query_name = name + read.query_sequence = "A" * query_length + read.flag = flag + read.reference_id = 0 + read.reference_start = start + read.mapping_quality = 60 + read.cigar = cigar + read.query_qualities = pysam.qualitystring_to_array("I" * query_length) + bam.write(read) + pysam.index(str(bam_path)) + + def run(space, output): + return preprocess_eclip(PipelineConfig( + genome_fasta=fasta, + gtf=gtf, + sminput=Sample("sminput", bam_path, "sminput"), + ips=(Sample("ip1", bam_path, "ip"),), + output_dir=output, + coordinate_space=space, + progress=False, + )) + + mature_dir = tmp_path / "mature" + gene_dir = tmp_path / "gene" + mature_qc = run("mature_transcript", mature_dir) + gene_qc = run("gene", gene_dir) + assert mature_qc["annotation"]["transcriptome_bases"] == 40 + assert gene_qc["annotation"]["coordinate_space_bases"] == 220 + assert gene_qc["annotation"]["intronic_bases"] == 180 + assert mature_qc["samples"]["sminput"]["retained"] == 2 + assert gene_qc["samples"]["sminput"]["retained"] == 4 + + with ProcessedECLIPDataset(gene_dir) as ds: + assert ds.coordinate_space == "gene" + assert ds.get_sequence("tp", 0, 110) == genomic[100:210] + assert ds.get_sequence("tm", 0, 110) == reverse_complement(genomic[400:510]) + assert ds.genome_to_coordinate("tp", "chr1", 150) == 50 + assert ds.genome_to_coordinate("tm", "chr1", 450) == 59 + assert ds.coordinate_to_genome("tp", 50) == ("chr1", 150, "+") + assert ds.coordinate_to_genome("tm", 59) == ("chr1", 450, "-") + assert [(b.start, b.end) for b in ds.get_genomic_blocks("tm", 10, 100)] == [(410, 500)] + assert int(ds.get_profile("tp", 50, 60, "sminput").sum()) == 1 + assert int(ds.get_profile("tm", 50, 60, "sminput").sum()) == 1 + + windows = tmp_path / "gene_windows" + summary = scan_windows(WindowScanConfig( + processed_dir=gene_dir, + output_prefix=windows, + window_size=10, + stride=10, + progress=False, + )) + assert summary["coordinate_space"] == "gene" + assert summary["region_type_windows"]["intron"] == 18 + selected = tmp_path / "gene_selected" + select_regions(SelectionConfig( + processed_dir=gene_dir, + windows=windows, + output_prefix=selected, + strategy="broad_coverage", + replicate_mode="combined", + min_total_count=0, + progress=False, + )) + bundle = make_rbpnet_bundle(RBPNetBundleConfig( + processed_dir=gene_dir, + selection_manifest=selected, + output_dir=tmp_path / "gene_bundle", + input_length=20, + profile_length=20, + max_jitter=2, + transcript_end_policy="shift_to_fit", + progress=False, + )) + assert bundle.config["coordinate_space"] == "gene" + with ProcessedECLIPDataset(gene_dir) as ds: + for index, item in enumerate(bundle.metadata): + start, end = item["sequence_materialized_start"], item["sequence_materialized_end"] + np.testing.assert_array_equal( + bundle.X[index], + encode_rna_sequence(ds.get_sequence(item["transcript_id"], start, end)), + ) + + def _write_processed_fixture(root: Path, *, length=12, profiles=None): root.mkdir(parents=True) sequence = ("ACGTGCGTAAAA" * ((length + 11) // 12))[:length] @@ -294,25 +504,26 @@ def test_selection_strategies_ids_serialization_and_stitching(tmp_path): scan_windows(WindowScanConfig( processed_dir=root, output_prefix=windows, window_size=4, stride=2, progress=False, )) - yeo = tmp_path / "yeo" + broad = tmp_path / "broad" summary = select_regions(SelectionConfig( - processed_dir=root, windows=windows, output_prefix=yeo, strategy="yeo_2026", + processed_dir=root, windows=windows, output_prefix=broad, strategy="broad_coverage", min_total_count=1, min_sminput_count=0, min_ip_count=0, progress=False, )) - loaded = load_selection_manifest(yeo) + loaded = load_selection_manifest(broad) assert summary["n_examples"] == len(loaded.rows) > 0 + assert summary["configuration"]["replicate_mode"] == "per_ip" assert len({row["example_id"] for row in loaded.rows}) == len(loaded.rows) - yeo2 = tmp_path / "yeo2" + broad2 = tmp_path / "broad2" select_regions(SelectionConfig( - processed_dir=root, windows=windows, output_prefix=yeo2, strategy="yeo_2026", + processed_dir=root, windows=windows, output_prefix=broad2, strategy="broad_coverage", min_total_count=1, min_sminput_count=0, min_ip_count=0, progress=False, )) assert [r["example_id"] for r in loaded.rows] == [ - r["example_id"] for r in load_selection_manifest(yeo2).rows + r["example_id"] for r in load_selection_manifest(broad2).rows ] - per_ip = tmp_path / "yeo_per_ip" + per_ip = tmp_path / "broad_per_ip" select_regions(SelectionConfig( - processed_dir=root, windows=windows, output_prefix=per_ip, strategy="yeo_2026", + processed_dir=root, windows=windows, output_prefix=per_ip, strategy="broad_coverage", min_total_count=0, min_sminput_count=0, min_ip_count=0, replicate_mode="per_ip", progress=False, )) @@ -332,6 +543,53 @@ def test_selection_strategies_ids_serialization_and_stitching(tmp_path): assert any(row["source_window_count"] > 1 for row in load_selection_manifest(classified).rows) +def test_zero_count_peak_negative_edges_and_broad_coverage_defaults(tmp_path): + profiles = np.zeros((3, 12), dtype=np.uint32) + profiles[0, 0] = 30 # informative input-only window + profiles[1, 4] = 30 # informative IP-only window + root = tmp_path / "processed" + _write_processed_fixture(root, profiles=profiles) + windows = tmp_path / "windows" + scan_windows(WindowScanConfig( + processed_dir=root, output_prefix=windows, window_size=4, stride=4, progress=False, + )) + + classified = tmp_path / "classified" + summary = select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=classified, + strategy="peak_gray_negative", + progress=False, + )) + rows = load_selection_manifest(classified).rows + assert summary["configuration"]["min_total_count"] == 8 + assert summary["configuration"]["min_ip_count"] == 0 + assert summary["configuration"]["min_sminput_count"] == 0 + by_start = {row["selection_start"]: row for row in rows} + assert by_start[0]["selection_state"] == "confident_negative" + assert by_start[0]["ip_pooled_count"] == 0 + assert by_start[0]["sminput_count"] == 30 + assert by_start[4]["selection_state"] == "peak" + assert by_start[4]["ip_pooled_count"] == 30 + assert by_start[4]["sminput_count"] == 0 + assert all(np.isfinite(row["log2_ip_pooled_vs_sminput"]) for row in rows) + + combined = tmp_path / "broad_combined" + combined_summary = select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=combined, + strategy="broad_coverage", replicate_mode="combined", progress=False, + )) + assert combined_summary["n_examples"] == 2 + per_ip = tmp_path / "broad_per_ip" + per_ip_summary = select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=per_ip, + strategy="broad_coverage", replicate_mode="per_ip", progress=False, + )) + assert per_ip_summary["n_examples"] == 3 + assert {row["replicate_id"] for row in load_selection_manifest(per_ip).rows} == {"ipA", "ipB"} + + def test_original_rbpnet_poisson_selection_and_50nt_advance(tmp_path): length = 500 profiles = np.zeros((3, length), dtype=np.uint32) @@ -356,6 +614,33 @@ def test_original_rbpnet_poisson_selection_and_50nt_advance(tmp_path): assert all(right - left >= 50 for left, right in zip(starts, starts[1:])) assert all(row["selection_pvalue"] < 0.01 for row in rows) + explicit = tmp_path / "original_explicit" + select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=explicit, + strategy="original_rbpnet", poisson_null="ip_locus_density", progress=False, + )) + explicit_rows = load_selection_manifest(explicit).rows + assert [row["example_id"] for row in explicit_rows] == [row["example_id"] for row in rows] + np.testing.assert_allclose( + [row["selection_pvalue"] for row in explicit_rows], + [row["selection_pvalue"] for row in rows], + ) + + sminput_null = tmp_path / "original_sminput" + sminput_summary = select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=sminput_null, + strategy="original_rbpnet", poisson_null="sminput", progress=False, + )) + sminput_rows = load_selection_manifest(sminput_null).rows + assert sminput_summary["configuration"]["poisson_null"] == "sminput" + assert sminput_rows + first = sminput_rows[0] + expected_mu = (first["sminput_count"] + 1.0) * (12 / 7) + assert first["selection_null_mean"] == pytest.approx(expected_mu) + assert first["selection_pvalue"] == pytest.approx( + poisson.sf(first["ip_pooled_count"] - 1, expected_mu) + ) + def test_materialized_bundle_exact_arrays_jitter_padding_and_mmap(tmp_path): root = tmp_path / "processed" @@ -366,7 +651,7 @@ def test_materialized_bundle_exact_arrays_jitter_padding_and_mmap(tmp_path): )) selection = tmp_path / "selection" select_regions(SelectionConfig( - processed_dir=root, windows=windows, output_prefix=selection, strategy="yeo_2026", + processed_dir=root, windows=windows, output_prefix=selection, strategy="broad_coverage", min_total_count=0, min_sminput_count=0, min_ip_count=0, progress=False, )) out = tmp_path / "bundle" @@ -413,3 +698,90 @@ def test_materialized_bundle_exact_arrays_jitter_padding_and_mmap(tmp_path): )) assert len(dropped.ids) < len(built.ids) assert all(m["sequence_left_pad"] == m["sequence_right_pad"] == 0 for m in dropped.metadata) + + +def test_shift_to_fit_boundaries_jitter_coordinates_and_short_loci(tmp_path): + assert _shifted_materialized_interval(0, 4, 2, 12) == (0, 8) + assert _shifted_materialized_interval(6, 4, 2, 12) == (2, 10) + assert _shifted_materialized_interval(11, 4, 2, 12) == (4, 12) + assert _shifted_materialized_interval(6, 10, 2, 12) is None + root = tmp_path / "processed" + sequence, profiles = _write_processed_fixture(root) + windows = tmp_path / "windows" + scan_windows(WindowScanConfig( + processed_dir=root, output_prefix=windows, window_size=4, stride=4, progress=False, + )) + selection = tmp_path / "selection" + select_regions(SelectionConfig( + processed_dir=root, windows=windows, output_prefix=selection, + strategy="broad_coverage", min_total_count=0, progress=False, + )) + out = tmp_path / "shifted" + bundle = make_rbpnet_bundle(RBPNetBundleConfig( + processed_dir=root, + selection_manifest=selection, + output_dir=out, + input_length=4, + profile_length=4, + max_jitter=2, + transcript_end_policy="shift_to_fit", + progress=False, + )) + by_selection_start = {item["selection_start"]: (index, item) for index, item in enumerate(bundle.metadata)} + assert by_selection_start[0][1]["sequence_materialized_start"] == 0 + assert by_selection_start[4][1]["sequence_materialized_start"] == 2 + assert by_selection_start[8][1]["sequence_materialized_start"] == 4 + assert all( + item["sequence_left_pad"] == item["sequence_right_pad"] == 0 + for item in bundle.metadata + ) + for index, item in enumerate(bundle.metadata): + start, end = item["sequence_materialized_start"], item["sequence_materialized_end"] + np.testing.assert_array_equal(bundle.X[index], encode_rna_sequence(sequence[start:end])) + pstart, pend = item["profile_materialized_start"], item["profile_materialized_end"] + np.testing.assert_array_equal(bundle.arrays["sminput_profiles"][index], profiles[0, pstart:pend]) + np.testing.assert_array_equal(bundle.arrays["ip_profiles"][index], profiles[1:, pstart:pend]) + assert item["sequence_anchor_offset"] == item["transcript_anchor"] - start + + left_index, left = by_selection_start[0] + offsets = [ + jitter_crop_offset( + anchor=left["transcript_anchor"], + materialized_start=left["sequence_materialized_start"], + locus_length=left["locus_length"], + crop_length=4, + jitter_shift=shift, + ) + for shift in range(-2, 3) + ] + assert offsets == [0, 0, 0, 1, 2] + for shift, offset in zip(range(-2, 3), offsets): + desired = left["transcript_anchor"] - 2 + shift + actual = min(max(desired, 0), len(sequence) - 4) + np.testing.assert_array_equal( + bundle.X[left_index, :, offset:offset + 4], + encode_rna_sequence(sequence[actual:actual + 4]), + ) + + right = by_selection_start[8][1] + assert jitter_crop_offset( + anchor=right["transcript_anchor"], + materialized_start=right["sequence_materialized_start"], + locus_length=right["locus_length"], + crop_length=4, + jitter_shift=2, + ) == 4 + assert bundle.config["transcript_end_policy"] == "shift_to_fit" + assert bundle.config["n_dropped_short_loci"] == 0 + + with pytest.raises(ValueError, match="no examples remain"): + make_rbpnet_bundle(RBPNetBundleConfig( + processed_dir=root, + selection_manifest=selection, + output_dir=tmp_path / "too_short", + input_length=10, + profile_length=10, + max_jitter=2, + transcript_end_policy="shift_to_fit", + progress=False, + )) From eaf5a873dd99db1cf69ffb54fd7c72a05ff73812 Mon Sep 17 00:00:00 2001 From: isvock Date: Tue, 11 Aug 2026 19:52:01 -0700 Subject: [PATCH 04/12] Implemented RBPNet --- README.md | 9 +- docs/api.rst | 18 +- docs/index.rst | 6 +- docs/rbpnet.md | 194 ++++++- docs/training_configuration.md | 77 ++- docs/usage.md | 6 +- src/transcriptml/cli/main.py | 8 +- src/transcriptml/models/__init__.py | 4 + src/transcriptml/models/rbpnet.py | 482 ++++++++++++++++ src/transcriptml/models/registry.py | 2 + src/transcriptml/rbpnet/__init__.py | 33 +- src/transcriptml/rbpnet/cli.py | 4 +- src/transcriptml/rbpnet/dataset.py | 391 +++++++++++++ src/transcriptml/rbpnet/losses.py | 267 +++++++++ src/transcriptml/rbpnet/training.py | 717 ++++++++++++++++++++++++ src/transcriptml/training/__init__.py | 9 +- src/transcriptml/training/evaluation.py | 32 +- src/transcriptml/training/splits.py | 85 +++ src/transcriptml/training/trainer.py | 15 + src/transcriptml/workflows/init_run.py | 49 +- tests/test_cli_analysis.py | 9 + tests/test_rbpnet_model.py | 431 ++++++++++++++ 22 files changed, 2812 insertions(+), 36 deletions(-) create mode 100644 src/transcriptml/models/rbpnet.py create mode 100644 src/transcriptml/rbpnet/dataset.py create mode 100644 src/transcriptml/rbpnet/losses.py create mode 100644 src/transcriptml/rbpnet/training.py create mode 100644 tests/test_rbpnet_model.py diff --git a/README.md b/README.md index b38de91..c44a6c7 100644 --- a/README.md +++ b/README.md @@ -13,13 +13,14 @@ TranscriptML currently supports three main workflows: - **MPRA-LegNet** models MPRA measurements from variable sequence inserts and supports targets such as RNA stability, translation, protein output, etc. -- **RBPNet/eCLIP data** converts FASTA/GTF/BAM inputs into a canonical +- **RBPNet/eCLIP** converts FASTA/GTF/BAM inputs into a canonical mature-transcript or full-gene coordinate experiment, descriptive windows, explicit selection - manifests, and memory-mappable model-ready NumPy bundles. The RBPNet model - and trainer are not implemented yet. + manifests, and memory-mappable model-ready NumPy bundles, then trains a + sequence-only target/control profile model with an optional replicate-aware + enrichment likelihood. In the future, I plan to also support [RiboNN](https://www.nature.com/articles/s41587-025-02712-x) modeling of translation efficiency measurements -and complete [RBPNet](https://link.springer.com/article/10.1186/s13059-023-03015-7) model training and interpretation. +and extend [RBPNet](https://link.springer.com/article/10.1186/s13059-023-03015-7) interpretation and model variants. ## Installation diff --git a/docs/api.rst b/docs/api.rst index 27c4518..81810f4 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -66,6 +66,18 @@ RBPNet/eCLIP data :members: RBPNetBundleConfig, jitter_crop_offset, make_rbpnet_bundle, load_rbpnet_bundle :member-order: bysource +.. automodule:: transcriptml.rbpnet.dataset + :members: RBPNetBatch, RBPNetDataset, collate_rbpnet, deduplicate_locus_indices + :member-order: bysource + +.. automodule:: transcriptml.rbpnet.losses + :members: RBPNetLossConfig, RBPNetLossOutput, multinomial_nll, replicate_binomial_nll, RBPNetObjective + :member-order: bysource + +.. automodule:: transcriptml.rbpnet.training + :members: train_rbpnet_model, evaluate_rbpnet_model, write_rbpnet_predictions + :member-order: bysource + Models ------ @@ -89,6 +101,10 @@ Models :members: SmallCNNConfig, SmallCNN :member-order: bysource +.. automodule:: transcriptml.models.rbpnet + :members: RBPNetConfig, RBPNetOutput, RBPNet, SamePadConv1d, SameLengthConvTranspose1d, theoretical_receptive_field + :member-order: bysource + Training and evaluation ----------------------- @@ -105,7 +121,7 @@ Training and evaluation :member-order: bysource .. automodule:: transcriptml.training.splits - :members: random_split_indices, predefined_split_indices, normalize_splits + :members: random_split_indices, predefined_split_indices, group_split_indices, validate_group_disjoint, normalize_splits :member-order: bysource .. automodule:: transcriptml.training.metrics diff --git a/docs/index.rst b/docs/index.rst index c5342bb..4d5c9f2 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -15,9 +15,9 @@ motif ablations, motif context scans, motif epistasis analyses, and Saluki-specific codon ISM. These analyses can expose learned regulatory sequence features as well as technical artifacts in the model or assay. -RBPNet/eCLIP data preprocessing, descriptive scanning, region selection, and -materialized dataset construction are supported. RBPNet model training remains -planned. +RBPNet/eCLIP preprocessing, descriptive scanning, region selection, +materialized dataset construction, and structured profile/enrichment training +are supported. Start here ---------- diff --git a/docs/rbpnet.md b/docs/rbpnet.md index 64b28c3..10042ee 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -1,10 +1,10 @@ # RBPNet/eCLIP data workflow -TranscriptML includes the complete data path needed before implementing an -RBPNet model. It starts from ordinary eCLIP alignments and ends with fixed-shape, -memory-mappable NumPy arrays. It does **not** implement an RBPNet architecture, -loss, trainer, peak caller intended for general use, GC matching, or training -example sampling. +TranscriptML includes an eCLIP path from ordinary alignments through fixed-shape, +memory-mappable arrays and structured RBPNet training. The scanner is +descriptive and the selectors prepare model examples; none is intended as a +general-purpose peak caller. GC matching and post-selection sampling are not +implemented. ```text FASTA + one-transcript-per-gene GTF + IP BAM(s) + SMInput BAM @@ -29,6 +29,9 @@ FASTA + one-transcript-per-gene GTF + IP BAM(s) + SMInput BAM | v fixed-shape, memory-mappable .npy arrays + | + v + structured RBPNet profile/enrichment training ``` Install the optional assay dependencies with: @@ -404,3 +407,184 @@ The canonical experiment uses HDF5 because it provides compressed lazy slicing over an entire transcriptome. The model bundle uses separate `.npy` files because its selected fixed-shape arrays are simple to inspect and memory-map. Changing selection, context, or jitter does not require reprocessing BAMs. + +## 5. RBPNet model and training + +Create a native starter config and train it with the same TranscriptML command +used by other registered models: + +```bash +transcriptml init-run --workflow rbpnet --out-dir configs/rbpnet +# Edit dataset/output paths and profile_length, then: +transcriptml train configs/rbpnet/train_config.json + +transcriptml evaluate \ + --checkpoint runs/rbpnet/model/best.pt \ + --dataset data/rbpnet_chr21 \ + --out-csv runs/rbpnet/predictions.csv +``` + +`transcriptml models show rbpnet --json` prints every architectural default. +The first model family intentionally requires equal sequence and profile crop +lengths. It uses no observed-control input, pooling, reverse-complement +augmentation, valid convolution, absolute-count head, or larger sequence +context than prediction context. + +### Default architecture + +The RNA4 sequence alone enters a same-padded 1D convolution with 128 filters +and kernel 12 followed by ReLU, then five residual blocks. Each block is a dilated +kernel-6 convolution, BatchNorm, ReLU, dropout 0.25, and residual addition. The +dilations are `[2, 4, 8, 16, 32]`; no positional pooling occurs. Two independent +kernel-25, stride-one transposed-convolution heads produce target and control +positional logits. A global-average-pooled linear head produces the scalar +mixture logit. Important dimensions, normalization, biases, kernels, dilation +schedule, dropout, head kinds, and profile length are configurable. + +The trunk receptive field is reported in checkpoints and `summary.json` and is + +```text +RF = 1 + (initial_kernel - 1) + + sum((residual_kernel - 1) * dilation) +``` + +which is 322 bases for the defaults (160 indexed positions to the left and 161 +to the right under the documented asymmetric even-kernel padding). Explicit +left/right padding preserves output index `i` as index `i`; the extra base of +an even effective kernel is placed on the right. + +| Model parameter | Default | Meaning | +| --- | --- | --- | +| `in_ch` | `4` | RNA4 input channels. | +| `n_filters` | `128` | Shared positional hidden width. | +| `initial_kernel_size` | `12` | Initial same-padded convolution kernel. | +| `n_residual_blocks` | `5` | Number of residual convolutions. | +| `residual_kernel_size` | `6` | Residual convolution kernel. | +| `dilations` | `null` | Explicit schedule; `null` resolves powers of two starting at 2. | +| `normalization` | `batch` | `batch`, `layer`, or `none`. | +| `dropout` | `0.25` | Dropout in each residual branch. | +| `profile_head_type` | `transpose_conv` | `transpose_conv` or ordinary same-padded `conv`. | +| `profile_head_kernel_size` | `25` | Kernel shared by the two separately parameterized profile heads. | +| `profile_head_bias` | `true` | Whether profile heads include a bias. | +| `enrichment_head_type` | `none` | `none`, `linear`, or `mlp`. | +| `enrichment_hidden` | `64` | Hidden width for the optional MLP only. | +| `enrichment_dropout` | `0` | Optional MLP dropout. | +| `profile_length` | `300` | Required input/output crop length, or `null` to accept any length. | + +### Profile model: target, control, and pi + +The two heads define independently normalized distributions +`p_target=softmax(target_logits)` and +`p_control=softmax(control_logits)`. With global mixing logit `a`, +`pi=sigmoid(a)` and the predicted IP distribution is + +```text +p_IP = pi * p_target + (1 - pi) * p_control +``` + +The mixture is evaluated with `logsigmoid` and `logaddexp` for stability. `pi` +is the latent fraction of the **positional IP profile** assigned to the target +component. It is not IP/SMInput enrichment and is never given an IP-vs-SMInput +binomial loss. SMInput is a profile training target, not a neural-network +input. + +By default, individual IP profiles are summed once across their replicate axis +and the complete multinomial NLL is calculated for the pooled IP counts under +`p_IP`. A second complete multinomial NLL compares SMInput counts with +`p_control`. The `lgamma` combinatorial constant is included by default and can +be disabled. A zero-total profile has no positional information, so that locus +is excluded from that profile component's mean instead of producing a NaN. +Components are reduced over informative loci and weighted by +`lambda_ip_profile` and `lambda_sm_profile` (both 1 by default). + +### Optional enrichment model + +Set `"enrichment_head_type": "linear"` to enable the default enrichment head, +or `"mlp"` for a configurable two-layer head. The head average-pools the shared +hidden representation only over the biological selection interval, using the +coordinate-derived mask for the current jittered crop. It supports variable +measurement widths. The result `eta_i` is a sequence-predicted log enrichment, +independent of `pi`. + +For replicate `j`, effective retained-event library sizes supply the known +offset and the observed selection-interval counts supply the binomial data: + +```text +depth_offset_j = log(L_IP_j / L_SM) +logit(p_ij) = eta_i + depth_offset_j +N_ij = IP_ij + SM_i +IP_ij ~ Binomial(N_ij, p_ij) +``` + +The logits-based complete binomial NLL is evaluated independently for each +valid locus-replicate pair and averaged over those pairs. One sequence row and +one `eta_i` therefore use every IP replicate without duplicating the locus. +`IP=0` and `SMInput=0` edge cases are exact and require no pseudocount; only a +pair with both counts zero is excluded because it has no information. Enabling +the head adds `lambda_enrichment * L_enrichment`, with weight 1 by default. + +### Jitter-ready structured batches + +`RBPNetDataset` memory-maps the bundle arrays and returns sequence, pooled and +individual IP profiles, SMInput profile, exact selection counts, effective +library sizes/depth offsets, selection mask, valid-position masks, coordinates, +and identifiers. With `max_train_jitter=J`, a deterministic RNG keyed by +seed/epoch/example samples a shift from `[-J,+J]` for training. Sequence and all +profiles use the same biological crop. The crop start is derived from anchor, +actual materialized start, and locus bounds, so boundary-shifted contexts do +not incorrectly assume offset `J+s`. Evaluation always uses shift zero. + +When enrichment is enabled, every allowed jittered crop must fully contain its +selection interval; invalid bundle/context combinations fail before training. +`max_train_jitter` cannot exceed the materialized bundle margin. + +### Splits, optimization, and outputs + +RBPNet starter configs use a `group` split on `group_gene_id`, keeping all +overlapping loci from one gene together. Transcript, chromosome, metadata, and +explicit predefined groups are also usable through their metadata columns. +Every non-random split is checked for group overlap. Row-random splitting is +rejected unless `allow_random_window_split=true` explicitly acknowledges the +leakage risk. Replicate-specific selection rows describing an identical locus +are deduplicated by default while retaining the complete replicate axis. + +AdamW, Adam, and SGD; plateau, cosine, and step schedulers; clipping; early +stopping; device selection; DataLoader workers; seeds; and mixed precision are +configurable. `history.json` logs total, pooled-IP profile, SMInput profile, and +enrichment losses independently. `best.pt` and `last.pt` retain model, loss, +optimizer, samples, coordinate space, split, receptive-field, and training +provenance. Evaluation CSVs contain `pi`, optional `eta`, and each replicate's +depth-adjusted predicted IP fraction. The raw structured tensors remain +available through the Python model output for future attribution work. + +Profile-only model block: + +```json +{ + "model": { + "name": "rbpnet", + "params": {"profile_length": 300, "enrichment_head_type": "none"} + }, + "loss": {"name": "rbpnet"}, + "max_train_jitter": 0 +} +``` + +To train profiles plus enrichment, change only the head and, if desired, the +independent component weights: + +```json +{ + "model": { + "name": "rbpnet", + "params": {"profile_length": 300, "enrichment_head_type": "linear"} + }, + "loss": { + "name": "rbpnet", + "lambda_ip_profile": 1.0, + "lambda_sm_profile": 1.0, + "lambda_enrichment": 1.0 + }, + "max_train_jitter": 32 +} +``` diff --git a/docs/training_configuration.md b/docs/training_configuration.md index 4ba0772..171b9e8 100644 --- a/docs/training_configuration.md +++ b/docs/training_configuration.md @@ -1,14 +1,15 @@ # Training Configuration -TranscriptML model training is controlled by a JSON or TOML file. The same top-level -training settings are used for Saluki and MPRA-LegNet runs; the main difference -between the workflows is the model selected under `model`. +TranscriptML model training is controlled by a JSON or TOML file. The same +top-level training settings are used across Saluki, MPRA-LegNet, and structured +RBPNet runs; the model and loss determine the batch contract. Create a starter JSON config with: ```bash transcriptml init-run --workflow saluki --out-dir configs/saluki transcriptml init-run --workflow legnet --out-dir configs/legnet +transcriptml init-run --workflow rbpnet --out-dir configs/rbpnet ``` Then train directly: @@ -105,6 +106,60 @@ Fields omitted from this starter, such as `gradient_clip_norm`, section. The Sherlock MPRA workflow has its own editable base config at `scripts/mpra/example_legnet_train_config.json`. +## RBPNet Starter Configuration + +`transcriptml init-run --workflow rbpnet` selects the structured RBPNet trainer. +Edit the bundle path, output path, and `profile_length` to match the bundle: + +```json +{ + "dataset": "__EDIT_ME_RBPNET_BUNDLE_DIR__", + "output_dir": "__EDIT_ME_RUN_DIR__/model", + "model": { + "name": "rbpnet", + "params": { + "profile_length": 300, + "enrichment_head_type": "none" + } + }, + "batch_size": 64, + "epochs": 100, + "learning_rate": 0.001, + "weight_decay": 0.0, + "optimizer": {"name": "adamw"}, + "lr_scheduler": {"name": "reduce_on_plateau", "patience": 3}, + "mixed_precision": false, + "gradient_clip_norm": 0.5, + "patience": 10, + "monitor": "val_loss", + "loss": { + "name": "rbpnet", + "lambda_ip_profile": 1.0, + "lambda_sm_profile": 1.0, + "lambda_enrichment": 1.0 + }, + "device": "auto", + "num_workers": 0, + "mmap_mode": "r", + "seed": 123, + "max_train_jitter": 0, + "deduplicate_loci": true, + "split_source": "config", + "split": { + "method": "group", + "group_col": "group_gene_id", + "val_frac": 0.1, + "test_frac": 0.1 + } +} +``` + +Set `enrichment_head_type` to `linear` (or `mlp`) to add the independent +replicate-aware enrichment likelihood. The three RBPNet component weights are +independent; `lambda_enrichment` has no effect when the head is disabled. See +the [RBPNet guide](rbpnet.md#rbpnet-model-and-training) for the equations, +bundle fields, and jitter semantics. + ## Top-Level Training Settings The following fields are accepted by `transcriptml train`. The default column @@ -116,11 +171,14 @@ above. | --- | --- | --- | --- | | `dataset` | path | required | Dataset bundle containing `X.npy` and its sidecar files. | | `output_dir` | path | required | Directory for checkpoints, history, split information, predictions, and the run summary. | -| `model` | mapping or string | `small_cnn` | Registered model name and optional constructor parameters. Saluki and MPRA starters explicitly select their workflow model. | +| `model` | mapping or string | `small_cnn` | Registered model name and optional constructor parameters. Workflow starters explicitly select their model. | | `batch_size` | integer | `64` | Number of examples per optimizer or evaluation batch. A final singleton training batch is dropped because batch-normalized models cannot train on it reliably. | | `epochs` | integer | `20` | Maximum number of training epochs before early stopping. | -| `learning_rate` | float | `0.001` | Learning rate passed to the AdamW optimizer. | -| `weight_decay` | float | `0.0` | AdamW weight-decay coefficient. | +| `learning_rate` | float | `0.001` | Learning rate. Structured RBPNet passes this to the selected optimizer unless overridden there. | +| `weight_decay` | float | `0.0` | Weight-decay coefficient. | +| `optimizer` | string or mapping | `"adamw"` | Structured RBPNet supports AdamW, Adam, and SGD, with optional optimizer parameters. Scalar workflows retain AdamW. | +| `lr_scheduler` | string, mapping, or `null` | `null` | Structured RBPNet supports plateau, cosine, and step schedulers. | +| `mixed_precision` | boolean | `false` | Enable autocast; CUDA also uses gradient scaling. | | `gradient_clip_norm` | float or `null` | `0.5` | Maximum global gradient norm. Set to `null`, `0`, or a negative value to disable clipping. | | `patience` | integer | `5` | Number of consecutive non-improving epochs tolerated by early stopping. A negative value disables early stopping. | | `monitor` | string or list | `"val_loss"` | Validation metric or metrics used to select `best.pt` and reset early-stopping patience. | @@ -135,6 +193,9 @@ above. | `sequence_controls` | mapping, list, or `null` | `null` | Optional sequence ablations applied before split selection. | | `split_source` | string | `"auto"` | Whether splits come from the bundle or from the `split` block. | | `split` | mapping | random 80/10/10 | Config-defined split settings, used according to `split_source`. | +| `max_train_jitter` | integer | `0` | Structured RBPNet shift range; cannot exceed the bundle's materialized margin. Evaluation remains shift zero. | +| `deduplicate_loci` | boolean | `true` | Collapse repeated eligibility rows for an identical RBPNet locus while retaining all replicate arrays. | +| `allow_random_window_split` | boolean | `false` | Explicitly permit unsafe RBPNet row-random splitting. Grouped splitting is the safe default. | The canonical model mapping contains a registered `name` and a `params` mapping: @@ -161,6 +222,10 @@ The metrics available to `monitor` are: - `val_loss` - `val_pearson` +Structured RBPNet runs instead expose `train_loss`, `val_loss`, and the +corresponding `train_`/`val_` forms of `ip_profile_loss`, `sm_profile_loss`, and +`enrichment_loss`. Profile-only runs report enrichment loss as zero. + Loss metrics improve when they decrease; Pearson metrics improve when they increase. A string can name one metric: diff --git a/docs/usage.md b/docs/usage.md index ca1ce4f..cef1315 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -5,9 +5,8 @@ This page walks through the two main TranscriptML workflows: - **Saluki**, for transcriptome-derived RNA stability measurements. - **MPRA-LegNet**, for single-insert MPRA-style measurements. -The assay-aware RBPNet/eCLIP preprocessing, window scanning, selection, and -bundle workflow has its own [RBPNet data guide](rbpnet.md). RBPNet model -training is intentionally not implemented yet. +The assay-aware RBPNet/eCLIP preprocessing, window scanning, selection, bundle, +and structured training workflow has its own [RBPNet guide](rbpnet.md). For each workflow, the basic pattern is the same: @@ -796,6 +795,7 @@ Create a starter config directory: ```bash transcriptml init-run --workflow saluki --out-dir configs/saluki transcriptml init-run --workflow legnet --out-dir configs/legnet +transcriptml init-run --workflow rbpnet --out-dir configs/rbpnet ``` These commands are intentionally small. They are meant to make the first run diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index 5264e20..6f3aa2f 100644 --- a/src/transcriptml/cli/main.py +++ b/src/transcriptml/cli/main.py @@ -95,7 +95,7 @@ def build_parser() -> argparse.ArgumentParser: add_rbpnet_parser(sub) p = sub.add_parser("init-run", help="Write starter run configuration files") - p.add_argument("--workflow", required=True, choices=["saluki", "legnet"]) + p.add_argument("--workflow", required=True, choices=["saluki", "legnet", "rbpnet"]) p.add_argument("--out-dir", required=True) p.add_argument("--force", action="store_true") @@ -552,7 +552,11 @@ def main(argv: list[str] | None = None) -> None: batch_size=args.batch_size, device=args.device, ) - metrics = {k: v for k, v in result.items() if k not in {"predictions", "targets", "indices"}} + non_summary_fields = { + "predictions", "targets", "indices", "example_ids", "pi", + "enrichment_logit", "depth_offsets", "replicate_names", + } + metrics = {k: v for k, v in result.items() if k not in non_summary_fields} summary_path = Path(evaluate_paths["out_csv"]).with_suffix(".summary.json") log_progress(f"evaluate: writing summary to {summary_path}") summary_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8") diff --git a/src/transcriptml/models/__init__.py b/src/transcriptml/models/__init__.py index 7af044f..b2e104c 100644 --- a/src/transcriptml/models/__init__.py +++ b/src/transcriptml/models/__init__.py @@ -3,6 +3,7 @@ from transcriptml.models.cnn import SmallCNN, SmallCNNConfig from transcriptml.models.legnet import LegNet, LegNetConfig from transcriptml.models.registry import ModelConfig, build_model, load_checkpoint, save_checkpoint +from transcriptml.models.rbpnet import RBPNet, RBPNetConfig, RBPNetOutput from transcriptml.models.reproduce import SalukiExact, SalukiExactConfig from transcriptml.models.saluki import SalukiLike, SalukiLikeConfig @@ -10,6 +11,9 @@ "LegNet", "LegNetConfig", "ModelConfig", + "RBPNet", + "RBPNetConfig", + "RBPNetOutput", "SalukiExact", "SalukiExactConfig", "SalukiLike", diff --git a/src/transcriptml/models/rbpnet.py b/src/transcriptml/models/rbpnet.py new file mode 100644 index 0000000..3661121 --- /dev/null +++ b/src/transcriptml/models/rbpnet.py @@ -0,0 +1,482 @@ +"""Configurable sequence-only RBPNet profile and enrichment model.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass + +import torch +import torch.nn.functional as F +from torch import nn + +from transcriptml.models.common import ChannelLayerNorm, dropout_or_identity + + +@dataclass +class RBPNetConfig: + """Configuration for the first TranscriptML RBPNet model family.""" + + in_ch: int = 4 + n_filters: int = 128 + initial_kernel_size: int = 12 + n_residual_blocks: int = 5 + residual_kernel_size: int = 6 + dilations: list[int] | None = None + normalization: str = "batch" + dropout: float = 0.25 + initial_bias: bool = False + residual_bias: bool = True + profile_head_type: str = "transpose_conv" + profile_head_kernel_size: int = 25 + profile_head_bias: bool = True + enrichment_head_type: str = "none" + enrichment_hidden: int = 64 + enrichment_dropout: float = 0.0 + profile_length: int | None = 300 + batch_norm_eps: float = 1e-5 + batch_norm_momentum: float = 0.1 + + def to_kwargs(self) -> dict[str, object]: + """Return constructor arguments accepted by :class:`RBPNet`.""" + + return asdict(self) + + +@dataclass(frozen=True) +class RBPNetOutput: + """Structured differentiable outputs from :class:`RBPNet`.""" + + target_logits: torch.Tensor + control_logits: torch.Tensor + target_log_probs: torch.Tensor + control_log_probs: torch.Tensor + target_probs: torch.Tensor + control_probs: torch.Tensor + mixing_logit: torch.Tensor + pi: torch.Tensor + ip_log_probs: torch.Tensor + ip_probs: torch.Tensor + enrichment_logit: torch.Tensor | None = None + + +def _validate_positive_int(name: str, value: int) -> int: + value = int(value) + if value <= 0: + raise ValueError(f"{name} must be positive") + return value + + +def resolve_dilations(n_residual_blocks: int, dilations: list[int] | tuple[int, ...] | None) -> tuple[int, ...]: + """Resolve an explicit schedule or powers-of-two defaults starting at two.""" + + n_blocks = int(n_residual_blocks) + if n_blocks < 0: + raise ValueError("n_residual_blocks must be non-negative") + values = tuple(2 ** (index + 1) for index in range(n_blocks)) if dilations is None else tuple( + int(value) for value in dilations + ) + if len(values) != n_blocks: + raise ValueError("dilations length must equal n_residual_blocks") + if any(value <= 0 for value in values): + raise ValueError("all dilations must be positive") + return values + + +def theoretical_receptive_field( + initial_kernel_size: int, + residual_kernel_size: int, + dilations: list[int] | tuple[int, ...], +) -> int: + """Return the position-preserving trunk's theoretical receptive-field width.""" + + initial = _validate_positive_int("initial_kernel_size", initial_kernel_size) + residual = _validate_positive_int("residual_kernel_size", residual_kernel_size) + return 1 + (initial - 1) + sum((residual - 1) * int(dilation) for dilation in dilations) + + +def same_padding(kernel_size: int, dilation: int = 1) -> tuple[int, int]: + """Return explicit left/right padding used to preserve indexed length. + + Even effective kernels necessarily have a half-base geometric center. The + extra base is placed on the right, matching the usual ``same`` convention; + output index ``i`` nevertheless remains output index ``i`` at every layer. + """ + + kernel = _validate_positive_int("kernel_size", kernel_size) + dilation = _validate_positive_int("dilation", dilation) + total = dilation * (kernel - 1) + left = total // 2 + return left, total - left + + +class SamePadConv1d(nn.Module): + """Conv1d with explicit, version-stable asymmetric same padding.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + *, + dilation: int = 1, + bias: bool = True, + ) -> None: + super().__init__() + self.padding = same_padding(kernel_size, dilation) + self.conv = nn.Conv1d( + int(in_channels), + int(out_channels), + kernel_size=int(kernel_size), + dilation=int(dilation), + padding=0, + bias=bool(bias), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + left, right = self.padding + return self.conv(F.pad(x, (left, right))) + + +class SameLengthConvTranspose1d(nn.Module): + """Stride-one transposed convolution cropped to the indexed input length.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + *, + bias: bool = True, + ) -> None: + super().__init__() + self.crop = same_padding(kernel_size) + self.conv = nn.ConvTranspose1d( + int(in_channels), + int(out_channels), + kernel_size=int(kernel_size), + stride=1, + padding=0, + bias=bool(bias), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = self.conv(x) + left, _ = self.crop + return output[..., left : left + x.shape[-1]] + + +def _normalization( + kind: str, + channels: int, + *, + eps: float, + momentum: float, +) -> nn.Module: + kind = str(kind).strip().lower() + if kind in {"batch", "batchnorm", "batch_norm"}: + return nn.BatchNorm1d(int(channels), eps=float(eps), momentum=float(momentum)) + if kind in {"layer", "layernorm", "layer_norm"}: + return ChannelLayerNorm(int(channels), eps=float(eps)) + if kind in {"none", "identity", "off"}: + return nn.Identity() + raise ValueError("normalization must be one of: batch, layer, none") + + +class DilatedResidualBlock(nn.Module): + """RBPNet-style dilated convolution, normalization, ReLU, dropout, add.""" + + def __init__( + self, + channels: int, + kernel_size: int, + dilation: int, + *, + normalization: str, + dropout: float, + bias: bool, + batch_norm_eps: float, + batch_norm_momentum: float, + ) -> None: + super().__init__() + self.dilation = int(dilation) + self.conv = SamePadConv1d( + channels, + channels, + kernel_size, + dilation=dilation, + bias=bias, + ) + self.norm = _normalization( + normalization, + channels, + eps=batch_norm_eps, + momentum=batch_norm_momentum, + ) + self.activation = nn.ReLU() + self.dropout = dropout_or_identity(float(dropout)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = self.dropout(self.activation(self.norm(self.conv(x)))) + if residual.shape != x.shape: + raise RuntimeError("residual convolution changed positional shape") + return x + residual + + +class PositionalProfileHead(nn.Module): + """Modular one-channel positional-logit head.""" + + def __init__( + self, + channels: int, + *, + head_type: str, + kernel_size: int, + bias: bool, + ) -> None: + super().__init__() + head_type = str(head_type).strip().lower() + if head_type in {"transpose_conv", "transposed_conv", "conv_transpose"}: + self.layer = SameLengthConvTranspose1d( + channels, 1, kernel_size, bias=bias + ) + self.head_type = "transpose_conv" + elif head_type in {"conv", "convolution"}: + self.layer = SamePadConv1d(channels, 1, kernel_size, bias=bias) + self.head_type = "conv" + else: + raise ValueError("profile_head_type must be transpose_conv or conv") + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return self.layer(hidden).squeeze(1) + + +def _masked_log_softmax( + logits: torch.Tensor, + mask: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + if mask is None: + log_probs = F.log_softmax(logits, dim=-1) + return log_probs, log_probs.exp() + mask = torch.as_tensor(mask, device=logits.device).bool() + if mask.shape != logits.shape: + raise ValueError( + f"profile_mask shape {tuple(mask.shape)} does not match logits {tuple(logits.shape)}" + ) + if torch.any(mask.sum(dim=-1) == 0): + raise ValueError("every profile_mask row must contain at least one valid position") + normalized = F.log_softmax(logits.masked_fill(~mask, -torch.inf), dim=-1) + log_probs = torch.where(mask, normalized, torch.zeros_like(normalized)) + probs = torch.where(mask, normalized.exp(), torch.zeros_like(normalized)) + return log_probs, probs + + +class RBPNet(nn.Module): + """Sequence-only RBPNet with latent target/control mixture and optional eta.""" + + def __init__( + self, + in_ch: int = 4, + n_filters: int = 128, + initial_kernel_size: int = 12, + n_residual_blocks: int = 5, + residual_kernel_size: int = 6, + dilations: list[int] | tuple[int, ...] | None = None, + normalization: str = "batch", + dropout: float = 0.25, + initial_bias: bool = False, + residual_bias: bool = True, + profile_head_type: str = "transpose_conv", + profile_head_kernel_size: int = 25, + profile_head_bias: bool = True, + enrichment_head_type: str = "none", + enrichment_hidden: int = 64, + enrichment_dropout: float = 0.0, + profile_length: int | None = 300, + batch_norm_eps: float = 1e-5, + batch_norm_momentum: float = 0.1, + ) -> None: + super().__init__() + self.in_ch = _validate_positive_int("in_ch", in_ch) + self.n_filters = _validate_positive_int("n_filters", n_filters) + self.initial_kernel_size = _validate_positive_int( + "initial_kernel_size", initial_kernel_size + ) + self.residual_kernel_size = _validate_positive_int( + "residual_kernel_size", residual_kernel_size + ) + self.dilations = resolve_dilations(n_residual_blocks, dilations) + self.profile_length = None if profile_length is None else _validate_positive_int( + "profile_length", profile_length + ) + if not 0 <= float(dropout) < 1 or not 0 <= float(enrichment_dropout) < 1: + raise ValueError("dropout probabilities must be in [0, 1)") + + self.initial_conv = SamePadConv1d( + self.in_ch, + self.n_filters, + self.initial_kernel_size, + bias=initial_bias, + ) + self.initial_activation = nn.ReLU() + self.residual_blocks = nn.ModuleList( + [ + DilatedResidualBlock( + self.n_filters, + self.residual_kernel_size, + dilation, + normalization=normalization, + dropout=dropout, + bias=residual_bias, + batch_norm_eps=batch_norm_eps, + batch_norm_momentum=batch_norm_momentum, + ) + for dilation in self.dilations + ] + ) + profile_kwargs = { + "head_type": profile_head_type, + "kernel_size": profile_head_kernel_size, + "bias": profile_head_bias, + } + self.target_profile_head = PositionalProfileHead(self.n_filters, **profile_kwargs) + self.control_profile_head = PositionalProfileHead(self.n_filters, **profile_kwargs) + self.mixing_head = nn.Linear(self.n_filters, 1) + + enrichment_kind = str(enrichment_head_type).strip().lower() + if enrichment_kind in {"none", "off", "disabled"}: + self.enrichment_head: nn.Module | None = None + self.enrichment_head_type = "none" + elif enrichment_kind == "linear": + self.enrichment_head = nn.Linear(self.n_filters, 1) + self.enrichment_head_type = "linear" + elif enrichment_kind == "mlp": + hidden = _validate_positive_int("enrichment_hidden", enrichment_hidden) + self.enrichment_head = nn.Sequential( + nn.Linear(self.n_filters, hidden), + nn.ReLU(), + dropout_or_identity(float(enrichment_dropout)), + nn.Linear(hidden, 1), + ) + self.enrichment_head_type = "mlp" + else: + raise ValueError("enrichment_head_type must be one of: none, linear, mlp") + + @property + def receptive_field(self) -> int: + """Theoretical receptive-field width of the shared trunk.""" + + return theoretical_receptive_field( + self.initial_kernel_size, + self.residual_kernel_size, + self.dilations, + ) + + @property + def receptive_field_extents(self) -> tuple[int, int]: + """Return left/right trunk context extents under explicit same padding.""" + + left, right = same_padding(self.initial_kernel_size) + for dilation in self.dilations: + block_left, block_right = same_padding( + self.residual_kernel_size, dilation + ) + left += block_left + right += block_right + return left, right + + @property + def enrichment_enabled(self) -> bool: + return self.enrichment_head is not None + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Return the shared position-preserving hidden representation.""" + + if x.ndim != 3 or x.shape[1] != self.in_ch: + raise ValueError( + f"RBPNet expects (B,{self.in_ch},L), received {tuple(x.shape)}" + ) + if self.profile_length is not None and x.shape[-1] != self.profile_length: + raise ValueError( + f"configured profile_length is {self.profile_length}, received {x.shape[-1]}" + ) + hidden = self.initial_activation(self.initial_conv(x.float())) + for block in self.residual_blocks: + hidden = block(hidden) + if hidden.shape[-1] != x.shape[-1]: + raise RuntimeError("RBPNet trunk changed positional length") + return hidden + + def _pool_measurement( + self, + hidden: torch.Tensor, + measurement_mask: torch.Tensor | None, + ) -> torch.Tensor: + if measurement_mask is None: + return hidden.mean(dim=-1) + weights = torch.as_tensor( + measurement_mask, device=hidden.device, dtype=hidden.dtype + ) + expected = (hidden.shape[0], hidden.shape[-1]) + if weights.shape != expected: + raise ValueError( + f"measurement_mask shape {tuple(weights.shape)} does not match {expected}" + ) + if torch.any(weights < 0): + raise ValueError("measurement_mask weights must be non-negative") + denominator = weights.sum(dim=-1, keepdim=True) + if torch.any(denominator <= 0): + raise ValueError("every measurement_mask row must have positive weight") + return (hidden * weights.unsqueeze(1)).sum(dim=-1) / denominator + + def forward( + self, + sequence: torch.Tensor, + *, + measurement_mask: torch.Tensor | None = None, + profile_mask: torch.Tensor | None = None, + ) -> RBPNetOutput: + """Predict latent profiles, their IP mixture, and optional enrichment.""" + + hidden = self.encode(sequence) + target_logits = self.target_profile_head(hidden) + control_logits = self.control_profile_head(hidden) + if target_logits.shape != control_logits.shape or target_logits.shape[-1] != sequence.shape[-1]: + raise RuntimeError("profile heads did not preserve positional shape") + target_log_probs, target_probs = _masked_log_softmax(target_logits, profile_mask) + control_log_probs, control_probs = _masked_log_softmax(control_logits, profile_mask) + + # Exclude legacy bundle padding from the global mixture summary when a + # validity mask is available. Shift-to-fit bundles are all-valid and + # therefore reduce to ordinary global average pooling. + global_hidden = self._pool_measurement(hidden, profile_mask) + mixing_logit = self.mixing_head(global_hidden).squeeze(-1) + pi = torch.sigmoid(mixing_logit) + log_pi = F.logsigmoid(mixing_logit).unsqueeze(-1) + log_one_minus_pi = F.logsigmoid(-mixing_logit).unsqueeze(-1) + ip_log_probs = torch.logaddexp( + log_pi + target_log_probs, + log_one_minus_pi + control_log_probs, + ) + ip_probs = ip_log_probs.exp() + if profile_mask is not None: + valid = torch.as_tensor(profile_mask, device=sequence.device).bool() + ip_log_probs = torch.where(valid, ip_log_probs, torch.zeros_like(ip_log_probs)) + ip_probs = torch.where(valid, ip_probs, torch.zeros_like(ip_probs)) + + enrichment_logit = None + if self.enrichment_head is not None: + measurement_hidden = self._pool_measurement(hidden, measurement_mask) + enrichment_logit = self.enrichment_head(measurement_hidden).squeeze(-1) + + return RBPNetOutput( + target_logits=target_logits, + control_logits=control_logits, + target_log_probs=target_log_probs, + control_log_probs=control_log_probs, + target_probs=target_probs, + control_probs=control_probs, + mixing_logit=mixing_logit, + pi=pi, + ip_log_probs=ip_log_probs, + ip_probs=ip_probs, + enrichment_logit=enrichment_logit, + ) diff --git a/src/transcriptml/models/registry.py b/src/transcriptml/models/registry.py index c044700..4e73eac 100644 --- a/src/transcriptml/models/registry.py +++ b/src/transcriptml/models/registry.py @@ -9,6 +9,7 @@ from transcriptml.models.cnn import SmallCNN, SmallCNNConfig from transcriptml.models.legnet import LegNet, LegNetConfig +from transcriptml.models.rbpnet import RBPNet, RBPNetConfig from transcriptml.models.reproduce import SalukiExact, SalukiExactConfig from transcriptml.models.saluki import SalukiLike, SalukiLikeConfig @@ -35,6 +36,7 @@ class ModelSpec: "saluki_like": ModelSpec(SalukiLike, SalukiLikeConfig), "saluki_gru": ModelSpec(SalukiLike, SalukiLikeConfig), "legnet": ModelSpec(LegNet, LegNetConfig), + "rbpnet": ModelSpec(RBPNet, RBPNetConfig), "saluki_exact": ModelSpec(SalukiExact, SalukiExactConfig), } diff --git a/src/transcriptml/rbpnet/__init__.py b/src/transcriptml/rbpnet/__init__.py index e7e2f73..1bccbfa 100644 --- a/src/transcriptml/rbpnet/__init__.py +++ b/src/transcriptml/rbpnet/__init__.py @@ -1,20 +1,23 @@ -"""RBPNet/eCLIP transcript-oriented locus data preparation. - -The public API deliberately stops at model-ready data. Neural-network -architectures, losses, and training are not part of this module. -""" +"""RBPNet/eCLIP preprocessing, structured data, losses, and training APIs.""" __all__ = [ "PipelineConfig", "ProcessedECLIPDataset", + "RBPNetBatch", "RBPNetBundleConfig", + "RBPNetDataset", + "RBPNetLossConfig", + "RBPNetObjective", "Sample", "SelectionConfig", "WindowScanConfig", "make_rbpnet_bundle", + "evaluate_rbpnet_model", "preprocess_eclip", "scan_windows", "select_regions", + "train_rbpnet_model", + "write_rbpnet_predictions", ] @@ -41,4 +44,24 @@ def __getattr__(name: str): from transcriptml.rbpnet.bundle import RBPNetBundleConfig, make_rbpnet_bundle return {"RBPNetBundleConfig": RBPNetBundleConfig, "make_rbpnet_bundle": make_rbpnet_bundle}[name] + if name in {"RBPNetBatch", "RBPNetDataset"}: + from transcriptml.rbpnet.dataset import RBPNetBatch, RBPNetDataset + + return {"RBPNetBatch": RBPNetBatch, "RBPNetDataset": RBPNetDataset}[name] + if name in {"RBPNetLossConfig", "RBPNetObjective"}: + from transcriptml.rbpnet.losses import RBPNetLossConfig, RBPNetObjective + + return {"RBPNetLossConfig": RBPNetLossConfig, "RBPNetObjective": RBPNetObjective}[name] + if name in {"evaluate_rbpnet_model", "train_rbpnet_model", "write_rbpnet_predictions"}: + from transcriptml.rbpnet.training import ( + evaluate_rbpnet_model, + train_rbpnet_model, + write_rbpnet_predictions, + ) + + return { + "evaluate_rbpnet_model": evaluate_rbpnet_model, + "train_rbpnet_model": train_rbpnet_model, + "write_rbpnet_predictions": write_rbpnet_predictions, + }[name] raise AttributeError(f"module 'transcriptml.rbpnet' has no attribute {name!r}") diff --git a/src/transcriptml/rbpnet/cli.py b/src/transcriptml/rbpnet/cli.py index 6784503..d980df1 100644 --- a/src/transcriptml/rbpnet/cli.py +++ b/src/transcriptml/rbpnet/cli.py @@ -13,7 +13,9 @@ def add_rbpnet_parser(subparsers) -> None: """Add the nested ``transcriptml rbpnet`` command family.""" - root = subparsers.add_parser("rbpnet", help="Prepare eCLIP data for future RBPNet models") + root = subparsers.add_parser( + "rbpnet", help="Preprocess eCLIP and construct RBPNet datasets" + ) commands = root.add_subparsers(dest="rbpnet_command", required=True) preprocess = commands.add_parser( diff --git a/src/transcriptml/rbpnet/dataset.py b/src/transcriptml/rbpnet/dataset.py new file mode 100644 index 0000000..99675fe --- /dev/null +++ b/src/transcriptml/rbpnet/dataset.py @@ -0,0 +1,391 @@ +"""Jitter-aware structured batches over materialized RBPNet bundles.""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Mapping, Sequence + +import numpy as np +import torch +from torch.utils.data import Dataset + +from transcriptml.data.bundle import DatasetBundle + + +@dataclass(frozen=True) +class RBPNetBatch: + """One structured RBPNet mini-batch.""" + + sequence: torch.Tensor + pooled_ip_profile: torch.Tensor + sminput_profile: torch.Tensor + individual_ip_profiles: torch.Tensor + ip_measurement_counts: torch.Tensor + sminput_measurement_counts: torch.Tensor + ip_library_sizes: torch.Tensor + sminput_library_size: torch.Tensor + depth_offsets: torch.Tensor + measurement_mask: torch.Tensor + profile_valid_mask: torch.Tensor + sequence_valid_mask: torch.Tensor + jitter_shift: torch.Tensor + crop_start: torch.Tensor + selection_start: torch.Tensor + selection_end: torch.Tensor + indices: torch.Tensor + example_ids: tuple[str, ...] + replicate_names: tuple[str, ...] + + def to(self, device: torch.device | str) -> "RBPNetBatch": + """Move tensor fields to a device while retaining identifiers.""" + + values = {} + for item in fields(self): + value = getattr(self, item.name) + values[item.name] = value.to(device) if isinstance(value, torch.Tensor) else value + return RBPNetBatch(**values) + + +def _required_array(bundle: DatasetBundle, name: str) -> np.ndarray: + try: + return bundle.arrays[name] + except KeyError as exc: + raise ValueError(f"RBPNet bundle is missing named array {name!r}") from exc + + +def _sample_metadata(bundle: DatasetBundle) -> tuple[str, tuple[str, ...], int, np.ndarray]: + metadata = bundle.config.get("sample_metadata") + if not isinstance(metadata, Mapping): + raise ValueError("RBPNet bundle lacks sample_metadata") + sminput = metadata.get("sminput") + ips = metadata.get("ip") + if not isinstance(sminput, Mapping) or not isinstance(ips, Sequence) or not ips: + raise ValueError("RBPNet bundle sample_metadata must define sminput and IP samples") + input_name = str(sminput.get("name", "")) + input_size = int(sminput.get("effective_library_size", 0)) + ip_names = tuple(str(sample.get("name", "")) for sample in ips) + ip_sizes = np.asarray( + [int(sample.get("effective_library_size", 0)) for sample in ips], + dtype=np.int64, + ) + if not input_name or any(not name for name in ip_names): + raise ValueError("RBPNet sample names must be non-empty") + if input_size <= 0 or np.any(ip_sizes <= 0): + raise ValueError("RBPNet effective library sizes must be positive") + axis_order = tuple(str(name) for name in metadata.get("ip_axis_order", ip_names)) + if axis_order != ip_names: + raise ValueError("RBPNet ip_axis_order disagrees with IP sample metadata") + return input_name, ip_names, input_size, ip_sizes + + +class RBPNetDataset(Dataset): + """Lazy fixed-crop view of one memory-mappable RBPNet bundle. + + Training jitter is deterministic for a given ``(seed, epoch, index)`` and + therefore works consistently with zero or multiple DataLoader workers. + Evaluation uses shift zero unless an explicit shift is requested through + :meth:`item_for_shift`. + """ + + def __init__( + self, + bundle: DatasetBundle, + *, + crop_length: int | None = None, + max_train_jitter: int = 0, + training: bool = False, + seed: int = 123, + require_full_measurement_interval: bool = True, + ) -> None: + if bundle.config.get("bundle_format") != "transcriptml-rbpnet-bundle": + raise ValueError("structured RBPNet training requires a TranscriptML RBPNet bundle") + if bundle.metadata is None: + raise ValueError("RBPNet bundle metadata is required for coordinate-aware crops") + self.bundle = bundle + self.metadata = bundle.metadata + self.training = bool(training) + self.seed = int(seed) + self.epoch = 0 + self.require_full_measurement_interval = bool(require_full_measurement_interval) + base_input_length = int(bundle.config.get("input_length", 0)) + base_profile_length = int(bundle.config.get("profile_length", 0)) + if base_input_length <= 0 or base_profile_length <= 0: + raise ValueError("RBPNet bundle must record positive input_length and profile_length") + if base_input_length != base_profile_length: + raise ValueError( + "this RBPNet family currently requires equal sequence and profile crop lengths" + ) + self.crop_length = base_input_length if crop_length is None else int(crop_length) + if self.crop_length != base_input_length: + raise ValueError( + f"crop_length {self.crop_length} must match bundle input_length {base_input_length}" + ) + self.bundle_max_jitter = int(bundle.config.get("max_jitter", 0)) + self.max_train_jitter = int(max_train_jitter) + if self.max_train_jitter < 0: + raise ValueError("max_train_jitter must be non-negative") + if self.max_train_jitter > self.bundle_max_jitter: + raise ValueError( + f"max_train_jitter {self.max_train_jitter} exceeds materialized margin " + f"{self.bundle_max_jitter}" + ) + self.boundary_policy = str(bundle.config.get("transcript_end_policy", "drop")) + if self.boundary_policy not in {"shift_to_fit", "drop", "pad"}: + raise ValueError(f"unsupported RBPNet boundary policy {self.boundary_policy!r}") + + self.X = bundle.X + self.sminput_profiles = _required_array(bundle, "sminput_profiles") + self.ip_profiles = _required_array(bundle, "ip_profiles") + self.selection_sminput_counts = _required_array(bundle, "selection_sminput_counts") + self.selection_ip_counts = _required_array(bundle, "selection_ip_counts") + self.sequence_valid_masks = bundle.arrays.get("sequence_valid_mask") + self.profile_valid_masks = bundle.arrays.get("profile_valid_mask") + self.sminput_name, self.replicate_names, self.sminput_library_size, self.ip_library_sizes = ( + _sample_metadata(bundle) + ) + self.depth_offsets = np.log( + self.ip_library_sizes.astype(np.float64) / float(self.sminput_library_size) + ).astype(np.float32) + self._validate_shapes() + if self.require_full_measurement_interval: + self._validate_measurement_intervals() + + def _validate_shapes(self) -> None: + n = int(self.X.shape[0]) + if self.X.ndim != 3 or self.X.shape[1] != 4: + raise ValueError("RBPNet X must have shape (N, 4, materialized_length)") + if self.sminput_profiles.ndim != 2 or self.sminput_profiles.shape[0] != n: + raise ValueError("sminput_profiles must have shape (N, materialized_length)") + if self.ip_profiles.ndim != 3 or self.ip_profiles.shape[:2] != ( + n, len(self.replicate_names) + ): + raise ValueError("ip_profiles must have shape (N, R, materialized_length)") + expected_width = self.crop_length + 2 * self.bundle_max_jitter + if self.X.shape[-1] != expected_width or self.sminput_profiles.shape[-1] != expected_width: + raise ValueError("materialized sequence/profile width disagrees with bundle jitter contract") + if self.ip_profiles.shape[-1] != expected_width: + raise ValueError("IP profile width disagrees with sequence/profile width") + if self.selection_sminput_counts.shape != (n,): + raise ValueError("selection_sminput_counts must have shape (N,)") + if self.selection_ip_counts.shape != (n, len(self.replicate_names)): + raise ValueError("selection_ip_counts must have shape (N, R)") + for name, masks in ( + ("sequence_valid_mask", self.sequence_valid_masks), + ("profile_valid_mask", self.profile_valid_masks), + ): + if masks is not None and masks.shape != (n, expected_width): + raise ValueError(f"{name} must have shape (N, materialized_length)") + + def _metadata_coordinates(self, index: int) -> tuple[int, int, int, int, int, int, int]: + row = self.metadata[int(index)] + required = ("transcript_anchor", "selection_start", "selection_end", "locus_length") + missing = [name for name in required if name not in row] + if missing: + raise ValueError( + f"RBPNet metadata row {index} lacks coordinate fields: {', '.join(missing)}" + ) + sequence_start = int( + row.get("sequence_materialized_start", row.get("sequence_context_start")) + ) + profile_start = int( + row.get("profile_materialized_start", row.get("profile_context_start")) + ) + return ( + int(row["transcript_anchor"]), + int(row["selection_start"]), + int(row["selection_end"]), + int(row["locus_length"]), + sequence_start, + profile_start, + self.crop_length, + ) + + def _crop_offsets(self, index: int, jitter_shift: int) -> tuple[int, int, int]: + anchor, _, _, locus_length, sequence_start, profile_start, length = ( + self._metadata_coordinates(index) + ) + shift = int(jitter_shift) + if abs(shift) > self.max_train_jitter: + raise ValueError( + f"requested jitter shift {shift} exceeds configured range " + f"[-{self.max_train_jitter}, {self.max_train_jitter}]" + ) + if self.boundary_policy == "pad": + sequence_offset = self.bundle_max_jitter + shift + profile_offset = self.bundle_max_jitter + shift + crop_start = sequence_start + sequence_offset + else: + if locus_length < length: + raise ValueError("locus is shorter than the requested RBPNet crop") + desired_start = anchor - length // 2 + shift + crop_start = min(max(desired_start, 0), locus_length - length) + sequence_offset = crop_start - sequence_start + profile_offset = crop_start - profile_start + if sequence_offset < 0 or sequence_offset + length > self.X.shape[-1]: + raise ValueError("coordinate-derived sequence crop lies outside materialized context") + if profile_offset < 0 or profile_offset + length > self.sminput_profiles.shape[-1]: + raise ValueError("coordinate-derived profile crop lies outside materialized context") + profile_crop_start = profile_start + profile_offset + if profile_crop_start != crop_start: + raise ValueError("sequence and profile crops do not describe the same biological interval") + return sequence_offset, profile_offset, crop_start + + def _measurement_mask(self, index: int, crop_start: int) -> np.ndarray: + _, selection_start, selection_end, _, _, _, length = self._metadata_coordinates(index) + if selection_end <= selection_start: + raise ValueError(f"RBPNet example {index} has an empty selection interval") + crop_end = crop_start + length + overlap_start = max(selection_start, crop_start) + overlap_end = min(selection_end, crop_end) + fully_contained = selection_start >= crop_start and selection_end <= crop_end + if self.require_full_measurement_interval and not fully_contained: + raise ValueError( + f"selection interval {selection_start}-{selection_end} for example " + f"{self.bundle.ids[index]} is not fully contained in jittered crop " + f"{crop_start}-{crop_end}; reduce jitter or use a larger model context" + ) + mask = np.zeros(length, dtype=np.float32) + if overlap_end > overlap_start: + mask[overlap_start - crop_start : overlap_end - crop_start] = 1.0 + if not np.any(mask): + raise ValueError(f"selection interval for example {self.bundle.ids[index]} misses model crop") + return mask + + def _validate_measurement_intervals(self) -> None: + shifts = {-self.max_train_jitter, self.max_train_jitter, 0} + for index in range(len(self)): + for shift in shifts: + _, _, crop_start = self._crop_offsets(index, shift) + self._measurement_mask(index, crop_start) + + def set_epoch(self, epoch: int) -> None: + """Set the deterministic training-jitter epoch.""" + + self.epoch = int(epoch) + + def __len__(self) -> int: + return int(self.X.shape[0]) + + def _sample_shift(self, index: int) -> int: + if not self.training or self.max_train_jitter == 0: + return 0 + rng = np.random.default_rng( + np.random.SeedSequence([self.seed, self.epoch, int(index)]) + ) + return int(rng.integers(-self.max_train_jitter, self.max_train_jitter + 1)) + + def item_for_shift(self, index: int, jitter_shift: int) -> dict[str, object]: + """Return one example using an explicit jitter shift, useful for testing.""" + + i = int(index) + sequence_offset, profile_offset, crop_start = self._crop_offsets(i, jitter_shift) + end_sequence = sequence_offset + self.crop_length + end_profile = profile_offset + self.crop_length + sequence = np.asarray(self.X[i, :, sequence_offset:end_sequence], dtype=np.float32) + individual_ip = np.asarray( + self.ip_profiles[i, :, profile_offset:end_profile], dtype=np.float32 + ) + sminput = np.asarray( + self.sminput_profiles[i, profile_offset:end_profile], dtype=np.float32 + ) + sequence_valid = ( + np.ones(self.crop_length, dtype=bool) + if self.sequence_valid_masks is None + else np.asarray( + self.sequence_valid_masks[i, sequence_offset:end_sequence], dtype=bool + ) + ) + profile_valid = ( + np.ones(self.crop_length, dtype=bool) + if self.profile_valid_masks is None + else np.asarray( + self.profile_valid_masks[i, profile_offset:end_profile], dtype=bool + ) + ) + return { + "sequence": sequence, + "pooled_ip_profile": individual_ip.sum(axis=0, dtype=np.float32), + "sminput_profile": sminput, + "individual_ip_profiles": individual_ip, + "ip_measurement_counts": np.asarray( + self.selection_ip_counts[i], dtype=np.float32 + ), + "sminput_measurement_counts": np.float32(self.selection_sminput_counts[i]), + "ip_library_sizes": self.ip_library_sizes.astype(np.float32, copy=False), + "sminput_library_size": np.float32(self.sminput_library_size), + "depth_offsets": self.depth_offsets, + "measurement_mask": self._measurement_mask(i, crop_start), + "profile_valid_mask": profile_valid, + "sequence_valid_mask": sequence_valid, + "jitter_shift": np.int64(jitter_shift), + "crop_start": np.int64(crop_start), + "selection_start": np.int64(self.metadata[i]["selection_start"]), + "selection_end": np.int64(self.metadata[i]["selection_end"]), + "index": np.int64(i), + "example_id": str(self.bundle.ids[i]), + "replicate_names": self.replicate_names, + } + + def __getitem__(self, index: int) -> dict[str, object]: + return self.item_for_shift(int(index), self._sample_shift(int(index))) + + +def collate_rbpnet(batch: list[Mapping[str, object]]) -> RBPNetBatch: + """Stack structured examples without obscuring replicate/sample axes.""" + + if not batch: + raise ValueError("cannot collate an empty RBPNet batch") + replicate_names = tuple(batch[0]["replicate_names"]) + if any(tuple(item["replicate_names"]) != replicate_names for item in batch): + raise ValueError("RBPNet batch mixes incompatible replicate axes") + + def stack(name: str, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: + return torch.as_tensor(np.stack([np.asarray(item[name]) for item in batch]), dtype=dtype) + + return RBPNetBatch( + sequence=stack("sequence"), + pooled_ip_profile=stack("pooled_ip_profile"), + sminput_profile=stack("sminput_profile"), + individual_ip_profiles=stack("individual_ip_profiles"), + ip_measurement_counts=stack("ip_measurement_counts"), + sminput_measurement_counts=stack("sminput_measurement_counts"), + ip_library_sizes=stack("ip_library_sizes"), + sminput_library_size=stack("sminput_library_size"), + depth_offsets=stack("depth_offsets"), + measurement_mask=stack("measurement_mask"), + profile_valid_mask=stack("profile_valid_mask", dtype=torch.bool), + sequence_valid_mask=stack("sequence_valid_mask", dtype=torch.bool), + jitter_shift=stack("jitter_shift", dtype=torch.long).reshape(-1), + crop_start=stack("crop_start", dtype=torch.long).reshape(-1), + selection_start=stack("selection_start", dtype=torch.long).reshape(-1), + selection_end=stack("selection_end", dtype=torch.long).reshape(-1), + indices=stack("index", dtype=torch.long).reshape(-1), + example_ids=tuple(str(item["example_id"]) for item in batch), + replicate_names=replicate_names, + ) + + +def deduplicate_locus_indices( + bundle: DatasetBundle, + indices: Sequence[int], +) -> tuple[list[int], int]: + """Remove replicate-eligibility duplicate rows while retaining all R tracks.""" + + if bundle.metadata is None: + raise ValueError("cannot deduplicate RBPNet loci without bundle metadata") + kept: list[int] = [] + first_by_key: dict[tuple[object, ...], int] = {} + for raw_index in indices: + index = int(raw_index) + row = bundle.metadata[index] + key = ( + row.get("coordinate_space", bundle.config.get("coordinate_space")), + row.get("transcript_id"), + int(row.get("selection_start", -1)), + int(row.get("selection_end", -1)), + int(row.get("transcript_anchor", -1)), + ) + if key not in first_by_key: + first_by_key[key] = index + kept.append(index) + return kept, len(indices) - len(kept) diff --git a/src/transcriptml/rbpnet/losses.py b/src/transcriptml/rbpnet/losses.py new file mode 100644 index 0000000..98f0797 --- /dev/null +++ b/src/transcriptml/rbpnet/losses.py @@ -0,0 +1,267 @@ +"""Stable structured likelihoods for TranscriptML RBPNet models.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Mapping + +import torch +import torch.nn.functional as F +from torch import nn + +from transcriptml.models.rbpnet import RBPNetOutput +from transcriptml.rbpnet.dataset import RBPNetBatch + + +@dataclass(frozen=True) +class ReducedLikelihood: + """A reduced likelihood and exact aggregation terms.""" + + loss: torch.Tensor + numerator: torch.Tensor + denominator: torch.Tensor + per_observation: torch.Tensor + valid: torch.Tensor + + +@dataclass +class RBPNetLossConfig: + """Weights and reporting choices for the structured RBPNet objective.""" + + name: str = "rbpnet" + lambda_ip_profile: float = 1.0 + lambda_sm_profile: float = 1.0 + lambda_enrichment: float = 1.0 + include_multinomial_constant: bool = True + include_binomial_constant: bool = True + + @classmethod + def from_config(cls, config: str | Mapping[str, object] | None) -> "RBPNetLossConfig": + if config is None: + return cls() + if isinstance(config, str): + values: dict[str, object] = {"name": config} + else: + values = dict(config) + name = str(values.pop("name", "rbpnet")).strip().lower() + values.pop("enrichment_enabled", None) + values.pop("effective_lambda_enrichment", None) + if name not in {"rbpnet", "rbpnet_profile", "rbpnet_profile_enrichment"}: + raise ValueError( + "RBPNet training requires loss.name='rbpnet', not " + repr(name) + ) + result = cls(name="rbpnet", **values) + for field_name in ( + "lambda_ip_profile", + "lambda_sm_profile", + "lambda_enrichment", + ): + if float(getattr(result, field_name)) < 0: + raise ValueError(f"{field_name} must be non-negative") + if result.lambda_ip_profile == result.lambda_sm_profile == result.lambda_enrichment == 0: + raise ValueError("at least one RBPNet loss weight must be positive") + return result + + def to_dict(self, *, enrichment_enabled: bool | None = None) -> dict[str, object]: + values = asdict(self) + if enrichment_enabled is not None: + values["enrichment_enabled"] = bool(enrichment_enabled) + values["effective_lambda_enrichment"] = ( + float(self.lambda_enrichment) if enrichment_enabled else 0.0 + ) + return values + + +@dataclass(frozen=True) +class RBPNetLossOutput: + """Total differentiable loss plus independently aggregatable components.""" + + loss: torch.Tensor + components: Mapping[str, torch.Tensor] + numerators: Mapping[str, torch.Tensor] + denominators: Mapping[str, torch.Tensor] + + +def _reduce_valid(nll: torch.Tensor, valid: torch.Tensor) -> ReducedLikelihood: + valid = valid.bool() + numerator = torch.where(valid, nll, torch.zeros_like(nll)).sum() + denominator = valid.sum().to(dtype=nll.dtype) + loss = numerator / denominator.clamp_min(1.0) + return ReducedLikelihood(loss, numerator, denominator, nll, valid) + + +def multinomial_nll( + log_probs: torch.Tensor, + counts: torch.Tensor, + *, + valid_positions: torch.Tensor | None = None, + include_constant: bool = True, +) -> ReducedLikelihood: + """Mean multinomial NLL over loci with nonzero profile totals. + + Zero-total profiles contain no positional information and are excluded from + the mean. With ``include_constant=True`` (the default), this is the complete + multinomial NLL, including the ``lgamma`` combinatorial term. + """ + + log_probs = log_probs.float() + counts = counts.to(device=log_probs.device, dtype=log_probs.dtype) + if log_probs.ndim != 2 or counts.shape != log_probs.shape: + raise ValueError("multinomial log_probs and counts must have matching (B, L) shapes") + if torch.any(counts < 0) or not torch.all(torch.isfinite(counts)): + raise ValueError("multinomial counts must be finite and non-negative") + if valid_positions is not None: + mask = valid_positions.to(device=log_probs.device).bool() + if mask.shape != counts.shape: + raise ValueError("valid_positions must match multinomial count shape") + if torch.any((~mask) & (counts != 0)): + raise ValueError("multinomial counts occur outside the valid profile mask") + total = counts.sum(dim=-1) + safe_terms = torch.where(counts > 0, counts * log_probs, torch.zeros_like(counts)) + log_likelihood = safe_terms.sum(dim=-1) + if include_constant: + log_likelihood = log_likelihood + torch.lgamma(total + 1) - torch.lgamma( + counts + 1 + ).sum(dim=-1) + nll = -log_likelihood + return _reduce_valid(nll, total > 0) + + +def replicate_binomial_nll( + eta: torch.Tensor, + ip_counts: torch.Tensor, + sminput_counts: torch.Tensor, + depth_offsets: torch.Tensor, + *, + include_constant: bool = True, +) -> ReducedLikelihood: + """Binomial NLL over valid locus-replicate observations. + + ``eta`` is one sequence-derived log enrichment per locus. Known effective + library sizes enter only through ``depth_offsets = log(L_IP/L_SM)``. + """ + + eta = eta.float().reshape(-1) + ip_counts = ip_counts.to(device=eta.device, dtype=eta.dtype) + sminput_counts = sminput_counts.to(device=eta.device, dtype=eta.dtype).reshape(-1) + depth_offsets = depth_offsets.to(device=eta.device, dtype=eta.dtype) + if ip_counts.ndim != 2 or ip_counts.shape[0] != eta.shape[0]: + raise ValueError("ip_counts must have shape (B, R) aligned to eta") + if sminput_counts.shape != eta.shape: + raise ValueError("sminput_counts must have shape (B,)") + if depth_offsets.ndim == 1: + if depth_offsets.shape[0] != ip_counts.shape[1]: + raise ValueError("one-dimensional depth_offsets must have shape (R,)") + offsets = depth_offsets.unsqueeze(0).expand_as(ip_counts) + elif depth_offsets.shape == ip_counts.shape: + offsets = depth_offsets + else: + raise ValueError("depth_offsets must have shape (R,) or (B, R)") + if torch.any(ip_counts < 0) or torch.any(sminput_counts < 0): + raise ValueError("IP and SMInput measurement counts must be non-negative") + failures = sminput_counts.unsqueeze(-1).expand_as(ip_counts) + total = ip_counts + failures + logits = eta.unsqueeze(-1) + offsets + # N * softplus(logit) - k * logit is the logits-based binomial + # cross-entropy. Subtract log(N choose k) for the complete NLL. + nll = total * F.softplus(logits) - ip_counts * logits + if include_constant: + log_choose = ( + torch.lgamma(total + 1) + - torch.lgamma(ip_counts + 1) + - torch.lgamma(failures + 1) + ) + nll = nll - log_choose + return _reduce_valid(nll, total > 0) + + +class RBPNetObjective(nn.Module): + """Target/control profile objective with an optional independent eta head.""" + + def __init__( + self, + config: RBPNetLossConfig | Mapping[str, object] | str | None = None, + *, + enrichment_enabled: bool, + ) -> None: + super().__init__() + self.config = ( + config + if isinstance(config, RBPNetLossConfig) + else RBPNetLossConfig.from_config(config) + ) + self.enrichment_enabled = bool(enrichment_enabled) + effective_weight = ( + float(self.config.lambda_ip_profile) + + float(self.config.lambda_sm_profile) + + ( + float(self.config.lambda_enrichment) + if self.enrichment_enabled + else 0.0 + ) + ) + if effective_weight == 0: + raise ValueError("enabled RBPNet loss components cannot all have zero weight") + + def forward(self, output: RBPNetOutput, batch: RBPNetBatch) -> RBPNetLossOutput: + pooled_from_replicates = batch.individual_ip_profiles.sum(dim=1) + if not torch.equal(pooled_from_replicates, batch.pooled_ip_profile): + raise ValueError("pooled IP profile does not equal the sum over replicate profiles") + ip = multinomial_nll( + output.ip_log_probs, + batch.pooled_ip_profile, + valid_positions=batch.profile_valid_mask, + include_constant=self.config.include_multinomial_constant, + ) + sm = multinomial_nll( + output.control_log_probs, + batch.sminput_profile, + valid_positions=batch.profile_valid_mask, + include_constant=self.config.include_multinomial_constant, + ) + zero = output.ip_log_probs.sum() * 0.0 + if self.enrichment_enabled: + if output.enrichment_logit is None: + raise ValueError("enrichment-enabled objective requires model enrichment_logit") + enrichment = replicate_binomial_nll( + output.enrichment_logit, + batch.ip_measurement_counts, + batch.sminput_measurement_counts, + batch.depth_offsets, + include_constant=self.config.include_binomial_constant, + ) + else: + enrichment = ReducedLikelihood( + loss=zero, + numerator=zero, + denominator=zero.detach(), + per_observation=zero.reshape(1), + valid=torch.zeros(1, dtype=torch.bool, device=zero.device), + ) + total = ( + float(self.config.lambda_ip_profile) * ip.loss + + float(self.config.lambda_sm_profile) * sm.loss + + ( + float(self.config.lambda_enrichment) * enrichment.loss + if self.enrichment_enabled + else zero + ) + ) + return RBPNetLossOutput( + loss=total, + components={ + "ip_profile_loss": ip.loss, + "sm_profile_loss": sm.loss, + "enrichment_loss": enrichment.loss, + }, + numerators={ + "ip_profile_loss": ip.numerator, + "sm_profile_loss": sm.numerator, + "enrichment_loss": enrichment.numerator, + }, + denominators={ + "ip_profile_loss": ip.denominator, + "sm_profile_loss": sm.denominator, + "enrichment_loss": enrichment.denominator, + }, + ) diff --git a/src/transcriptml/rbpnet/training.py b/src/transcriptml/rbpnet/training.py new file mode 100644 index 0000000..ad3da80 --- /dev/null +++ b/src/transcriptml/rbpnet/training.py @@ -0,0 +1,717 @@ +"""Structured RBPNet training integrated with TranscriptML checkpoints/configs.""" + +from __future__ import annotations + +import csv +import json +from dataclasses import asdict +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np +import torch +from torch.utils.data import DataLoader, Subset + +from transcriptml.data.bundle import DatasetBundle +from transcriptml.devices import resolve_device +from transcriptml.models.rbpnet import RBPNet +from transcriptml.models.registry import build_model, normalize_model_config, save_checkpoint +from transcriptml.progress import ProgressReporter, log_progress +from transcriptml.rbpnet.dataset import ( + RBPNetBatch, + RBPNetDataset, + collate_rbpnet, + deduplicate_locus_indices, +) +from transcriptml.rbpnet.losses import RBPNetLossConfig, RBPNetObjective +from transcriptml.training.splits import ( + group_split_indices, + normalize_splits, + predefined_split_indices, + random_split_indices, + validate_group_disjoint, +) + + +def _config_dict(value: str | Mapping[str, Any] | None, default_name: str) -> dict[str, Any]: + if value is None: + return {"name": default_name, "params": {}} + if isinstance(value, str): + return {"name": value, "params": {}} + result = dict(value) + params = dict(result.pop("params", {}) or {}) + params.update(result) + return {"name": str(params.pop("name", default_name)).lower(), "params": params} + + +def _build_optimizer(model: torch.nn.Module, cfg) -> tuple[torch.optim.Optimizer, dict[str, Any]]: + config = _config_dict(getattr(cfg, "optimizer", None), "adamw") + name = config["name"] + params = dict(config["params"]) + params.setdefault("lr", float(cfg.learning_rate)) + params.setdefault("weight_decay", float(cfg.weight_decay)) + if name == "adamw": + optimizer = torch.optim.AdamW(model.parameters(), **params) + elif name == "adam": + optimizer = torch.optim.Adam(model.parameters(), **params) + elif name == "sgd": + optimizer = torch.optim.SGD(model.parameters(), **params) + else: + raise ValueError("RBPNet optimizer must be one of: adamw, adam, sgd") + return optimizer, {"name": name, "params": params} + + +def _build_scheduler( + optimizer: torch.optim.Optimizer, + value: str | Mapping[str, Any] | None, + *, + epochs: int, +) -> tuple[object | None, dict[str, Any] | None, bool]: + if value is None: + return None, None, False + config = _config_dict(value, "none") + name = config["name"] + params = dict(config["params"]) + if name in {"none", "off", "disabled"}: + return None, {"name": "none", "params": {}}, False + if name in {"reduce_on_plateau", "plateau"}: + params.setdefault("mode", "min") + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **params) + return scheduler, {"name": "reduce_on_plateau", "params": params}, True + if name in {"cosine", "cosine_annealing"}: + params.setdefault("T_max", int(epochs)) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **params) + return scheduler, {"name": "cosine", "params": params}, False + if name in {"step", "step_lr"}: + params.setdefault("step_size", 10) + params.setdefault("gamma", 0.1) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, **params) + return scheduler, {"name": "step", "params": params}, False + raise ValueError("RBPNet lr_scheduler must be one of: none, reduce_on_plateau, cosine, step") + + +def _build_grad_scaler(*, enabled: bool): + """Build a CUDA scaler across the supported PyTorch API variants.""" + + if hasattr(torch, "amp") and hasattr(torch.amp, "GradScaler"): + try: + return torch.amp.GradScaler("cuda", enabled=enabled) + except TypeError: # PyTorch versions without the device argument. + return torch.amp.GradScaler(enabled=enabled) + return torch.cuda.amp.GradScaler(enabled=enabled) + + +def _make_config_splits(bundle: DatasetBundle, cfg) -> tuple[dict[str, list[int]], str, bool]: + split_cfg = dict(cfg.split or {}) + method = str(split_cfg.get("method", "group")).lower() + if method == "group": + if bundle.metadata is None: + raise ValueError("group split requested but RBPNet bundle has no metadata") + group_col = str(split_cfg.get("group_col", "group_gene_id")) + return ( + group_split_indices( + bundle.metadata, + group_col=group_col, + val_frac=float(split_cfg.get("val_frac", 0.1)), + test_frac=float(split_cfg.get("test_frac", 0.1)), + seed=int(split_cfg.get("seed", cfg.seed)), + ), + group_col, + False, + ) + if method == "metadata": + if bundle.metadata is None: + raise ValueError("metadata split requested but RBPNet bundle has no metadata") + return ( + predefined_split_indices( + bundle.metadata, + split_col=str(split_cfg.get("split_col", "split")), + ), + str(split_cfg.get("group_col", "group_gene_id")), + False, + ) + if method == "predefined": + return ( + normalize_splits(split_cfg["splits"]), + str(split_cfg.get("group_col", "group_gene_id")), + False, + ) + if method == "random": + if not bool(getattr(cfg, "allow_random_window_split", False)): + raise ValueError( + "row-level random splitting is unsafe for overlapping RBPNet windows; " + "use split.method='group' (recommended) or explicitly set " + "allow_random_window_split=true" + ) + return ( + random_split_indices( + int(bundle.X.shape[0]), + val_frac=float(split_cfg.get("val_frac", 0.1)), + test_frac=float(split_cfg.get("test_frac", 0.1)), + seed=int(split_cfg.get("seed", cfg.seed)), + ), + str(split_cfg.get("group_col", "group_gene_id")), + True, + ) + raise ValueError(f"Unknown RBPNet split method {method!r}") + + +def _select_rbpnet_splits( + bundle: DatasetBundle, + cfg, +) -> tuple[dict[str, list[int]], str, str, bool]: + source = str(cfg.split_source or "auto").strip().lower() + if source not in {"auto", "bundle", "config"}: + raise ValueError("split_source must be one of: auto, bundle, config") + if source == "bundle" or (source == "auto" and bundle.splits is not None): + if bundle.splits is None: + raise ValueError("split_source='bundle' requested but dataset bundle has no splits") + splits = normalize_splits(bundle.splits) + group_col = str(dict(cfg.split or {}).get("group_col", "group_gene_id")) + allow_random = False + source_used = "bundle" + else: + splits, group_col, allow_random = _make_config_splits(bundle, cfg) + source_used = "config" + if not allow_random: + if bundle.metadata is None: + raise ValueError("leakage validation requires RBPNet bundle metadata") + validate_group_disjoint(splits, bundle.metadata, group_col=group_col) + return splits, source_used, group_col, allow_random + + +def _deduplicate_splits( + bundle: DatasetBundle, + splits: Mapping[str, Sequence[int]], + *, + enabled: bool, +) -> tuple[dict[str, list[int]], int]: + normalized = normalize_splits(splits) + if not enabled: + return normalized, 0 + result: dict[str, list[int]] = {} + dropped = 0 + for name, indices in normalized.items(): + result[name], count = deduplicate_locus_indices(bundle, indices) + dropped += count + return result, dropped + + +def _loader( + dataset: RBPNetDataset, + indices: Sequence[int], + batch_size: int, + *, + shuffle: bool, + num_workers: int, + pin_memory: bool, +) -> DataLoader | None: + if not indices: + return None + return DataLoader( + Subset(dataset, [int(index) for index in indices]), + batch_size=int(batch_size), + shuffle=bool(shuffle), + num_workers=int(num_workers), + pin_memory=bool(pin_memory), + persistent_workers=False, + collate_fn=collate_rbpnet, + ) + + +def _aggregate_loss( + numerators: Mapping[str, float], + denominators: Mapping[str, float], + loss_config: RBPNetLossConfig, + *, + enrichment_enabled: bool, +) -> dict[str, float]: + components = { + name: float(numerators[name] / denominators[name]) + if denominators[name] > 0 + else 0.0 + for name in numerators + } + total = ( + float(loss_config.lambda_ip_profile) * components["ip_profile_loss"] + + float(loss_config.lambda_sm_profile) * components["sm_profile_loss"] + + ( + float(loss_config.lambda_enrichment) * components["enrichment_loss"] + if enrichment_enabled + else 0.0 + ) + ) + return {"loss": total, **components} + + +def _run_loader( + model: RBPNet, + loader: DataLoader | None, + *, + device: torch.device, + objective: RBPNetObjective, + optimizer: torch.optim.Optimizer | None = None, + gradient_clip_norm: float | None = None, + mixed_precision: bool = False, + scaler: Any | None = None, + progress: bool = True, + progress_label: str = "RBPNet batches", + return_predictions: bool = False, +) -> dict[str, Any]: + names = ("ip_profile_loss", "sm_profile_loss", "enrichment_loss") + if loader is None: + return { + "loss": float("nan"), + **{name: float("nan") for name in names}, + "n_examples": 0, + "indices": [], + "example_ids": [], + "pi": np.empty(0, dtype=np.float32), + "enrichment_logit": None, + } + training = optimizer is not None + model.train(training) + numerators = {name: 0.0 for name in names} + denominators = {name: 0.0 for name in names} + pis: list[np.ndarray] = [] + etas: list[np.ndarray] = [] + indices: list[int] = [] + example_ids: list[str] = [] + n_examples = 0 + reporter = ProgressReporter( + progress_label, + total=len(loader), + unit="batches", + enabled=progress, + percent_step=25.0, + ) + amp_enabled = bool(mixed_precision) + amp_dtype = torch.float16 if device.type == "cuda" else torch.bfloat16 + for batch in loader: + assert isinstance(batch, RBPNetBatch) + batch = batch.to(device) + n_examples += int(batch.sequence.shape[0]) + if training: + optimizer.zero_grad(set_to_none=True) + with torch.set_grad_enabled(training): + with torch.autocast( + device_type=device.type, + dtype=amp_dtype, + enabled=amp_enabled, + ): + output = model( + batch.sequence, + measurement_mask=batch.measurement_mask, + profile_mask=batch.profile_valid_mask, + ) + loss_output = objective(output, batch) + if training: + if scaler is not None and scaler.is_enabled(): + scaler.scale(loss_output.loss).backward() + if gradient_clip_norm is not None and float(gradient_clip_norm) > 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_( + model.parameters(), float(gradient_clip_norm) + ) + scaler.step(optimizer) + scaler.update() + else: + loss_output.loss.backward() + if gradient_clip_norm is not None and float(gradient_clip_norm) > 0: + torch.nn.utils.clip_grad_norm_( + model.parameters(), float(gradient_clip_norm) + ) + optimizer.step() + for name in names: + numerators[name] += float(loss_output.numerators[name].detach().cpu()) + denominators[name] += float(loss_output.denominators[name].detach().cpu()) + if return_predictions: + pis.append(output.pi.detach().float().cpu().numpy()) + if output.enrichment_logit is not None: + etas.append(output.enrichment_logit.detach().float().cpu().numpy()) + indices.extend(int(value) for value in batch.indices.detach().cpu().tolist()) + example_ids.extend(batch.example_ids) + reporter.update() + reporter.close() + metrics: dict[str, Any] = _aggregate_loss( + numerators, + denominators, + objective.config, + enrichment_enabled=objective.enrichment_enabled, + ) + metrics["n_examples"] = n_examples + if return_predictions: + metrics.update( + { + "indices": indices, + "example_ids": example_ids, + "pi": np.concatenate(pis) if pis else np.empty(0, dtype=np.float32), + "enrichment_logit": ( + np.concatenate(etas) if etas else None + ), + } + ) + return metrics + + +def _is_better(value: float, best: float | None, name: str) -> bool: + if not np.isfinite(value): + return False + if best is None: + return True + return value < best if name.endswith("loss") else value > best + + +def _monitor_names(value: str | Sequence[str]) -> tuple[str, ...]: + names = ( + [part.strip() for part in value.split(",") if part.strip()] + if isinstance(value, str) + else [str(part).strip() for part in value if str(part).strip()] + ) + if not names: + raise ValueError("monitor must name at least one metric") + return tuple(names) + + +def _write_predictions( + path: str | Path, + metrics: Mapping[str, Any], + *, + depth_offsets: np.ndarray, + replicate_names: Sequence[str], +) -> None: + eta = metrics.get("enrichment_logit") + fieldnames = ["index", "id", "pi"] + if eta is not None: + fieldnames.append("enrichment_logit") + fieldnames.extend(f"predicted_ip_fraction_{name}" for name in replicate_names) + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + for position, (index, identifier, pi) in enumerate( + zip(metrics["indices"], metrics["example_ids"], metrics["pi"]) + ): + row: dict[str, object] = { + "index": int(index), + "id": str(identifier), + "pi": float(pi), + } + if eta is not None: + value = float(eta[position]) + row["enrichment_logit"] = value + probabilities = 1.0 / (1.0 + np.exp(-(value + depth_offsets))) + for name, probability in zip(replicate_names, probabilities): + row[f"predicted_ip_fraction_{name}"] = float(probability) + writer.writerow(row) + + +def train_rbpnet_model(bundle: DatasetBundle, cfg) -> dict[str, Any]: + """Train a registered RBPNet model without routing through scalar targets.""" + + torch.manual_seed(int(cfg.seed)) + np.random.seed(int(cfg.seed)) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(int(cfg.seed)) + device = resolve_device(cfg.device) + model_config = normalize_model_config(cfg.model) + if model_config.name != "rbpnet": + raise ValueError("train_rbpnet_model requires model.name='rbpnet'") + model = build_model(model_config).to(device) + if not isinstance(model, RBPNet): + raise TypeError("registered rbpnet model did not build an RBPNet instance") + crop_length = int(bundle.config.get("input_length", 0)) + if model.profile_length is not None and model.profile_length != crop_length: + raise ValueError( + f"model profile_length {model.profile_length} does not match bundle crop " + f"length {crop_length}" + ) + loss_config = RBPNetLossConfig.from_config(cfg.loss) + objective = RBPNetObjective( + loss_config, + enrichment_enabled=model.enrichment_enabled, + ).to(device) + splits, split_source, group_col, random_split = _select_rbpnet_splits(bundle, cfg) + splits, n_deduplicated = _deduplicate_splits( + bundle, + splits, + enabled=bool(getattr(cfg, "deduplicate_loci", True)), + ) + if not splits.get("train"): + raise ValueError("RBPNet training split is empty after locus deduplication") + + train_dataset = RBPNetDataset( + bundle, + crop_length=crop_length, + max_train_jitter=int(getattr(cfg, "max_train_jitter", 0)), + training=True, + seed=int(cfg.seed), + require_full_measurement_interval=model.enrichment_enabled, + ) + eval_dataset = RBPNetDataset( + bundle, + crop_length=crop_length, + max_train_jitter=0, + training=False, + seed=int(cfg.seed), + require_full_measurement_interval=model.enrichment_enabled, + ) + pin_memory = device.type == "cuda" + train_loader = _loader( + train_dataset, + splits["train"], + cfg.batch_size, + shuffle=True, + num_workers=cfg.num_workers, + pin_memory=pin_memory, + ) + val_loader = _loader( + eval_dataset, + splits.get("val", []), + cfg.batch_size, + shuffle=False, + num_workers=cfg.num_workers, + pin_memory=pin_memory, + ) + optimizer, optimizer_config = _build_optimizer(model, cfg) + scheduler, scheduler_config, scheduler_uses_metric = _build_scheduler( + optimizer, + getattr(cfg, "lr_scheduler", None), + epochs=int(cfg.epochs), + ) + mixed_precision = bool(getattr(cfg, "mixed_precision", False)) + if mixed_precision and device.type not in {"cpu", "cuda"}: + raise ValueError("RBPNet mixed precision currently supports CPU and CUDA devices") + scaler = _build_grad_scaler(enabled=mixed_precision and device.type == "cuda") + out = Path(cfg.output_dir) + out.mkdir(parents=True, exist_ok=True) + parameter_count = sum(parameter.numel() for parameter in model.parameters()) + split_counts = {name: len(splits.get(name, [])) for name in ("train", "val", "test")} + log_progress( + ( + f"RBPNet training: device={device}, parameters={parameter_count:,}, " + f"receptive_field={model.receptive_field}, jitter={train_dataset.max_train_jitter}, " + f"train={split_counts['train']}, val={split_counts['val']}, " + f"test={split_counts['test']}" + ), + enabled=cfg.progress, + ) + + history: list[dict[str, float | int]] = [] + monitors = _monitor_names(cfg.monitor) + best_values: dict[str, float | None] = {name: None for name in monitors} + best_epoch = -1 + stale = 0 + for epoch in range(1, int(cfg.epochs) + 1): + train_dataset.set_epoch(epoch) + train_metrics = _run_loader( + model, + train_loader, + device=device, + objective=objective, + optimizer=optimizer, + gradient_clip_norm=cfg.gradient_clip_norm, + mixed_precision=mixed_precision, + scaler=scaler, + progress=cfg.progress, + progress_label=f"epoch {epoch} RBPNet train", + ) + val_metrics = _run_loader( + model, + val_loader, + device=device, + objective=objective, + mixed_precision=mixed_precision, + progress=cfg.progress, + progress_label=f"epoch {epoch} RBPNet val", + ) + row: dict[str, float | int] = { + "epoch": epoch, + "learning_rate": float(optimizer.param_groups[0]["lr"]), + "train_loss": float(train_metrics["loss"]), + "val_loss": float(val_metrics["loss"]), + } + for component in ("ip_profile_loss", "sm_profile_loss", "enrichment_loss"): + row[f"train_{component}"] = float(train_metrics[component]) + row[f"val_{component}"] = float(val_metrics[component]) + missing = [name for name in monitors if name not in row] + if missing: + raise ValueError(f"Unknown RBPNet monitor metric(s): {', '.join(missing)}") + values = {name: float(row[name]) for name in monitors} + improved = any(_is_better(values[name], best_values[name], name) for name in monitors) + history.append(row) + checkpoint_extra = { + "splits": splits, + "split_source_used": split_source, + "train_config": asdict(cfg), + "loss_config": loss_config.to_dict( + enrichment_enabled=model.enrichment_enabled + ), + "optimizer_config": optimizer_config, + "lr_scheduler_config": scheduler_config, + "coordinate_space": bundle.config.get("coordinate_space"), + "sample_metadata": bundle.config.get("sample_metadata"), + "parameter_count": parameter_count, + "receptive_field": model.receptive_field, + } + if improved: + best_values = values + best_epoch = epoch + stale = 0 + save_checkpoint( + out / "best.pt", + model, + model_config, + epoch=epoch, + metrics=row, + optimizer_state=optimizer.state_dict(), + extra=checkpoint_extra, + ) + else: + stale += 1 + save_checkpoint( + out / "last.pt", + model, + model_config, + epoch=epoch, + metrics=row, + optimizer_state=optimizer.state_dict(), + extra=checkpoint_extra, + ) + if scheduler is not None: + metric = float(val_metrics["loss"]) + if not np.isfinite(metric): + metric = float(train_metrics["loss"]) + scheduler.step(metric) if scheduler_uses_metric else scheduler.step() + log_progress( + ( + f"epoch {epoch}/{cfg.epochs}: train={row['train_loss']:.6g}, " + f"val={row['val_loss']:.6g}, " + f"IP={row['val_ip_profile_loss']:.6g}, " + f"SM={row['val_sm_profile_loss']:.6g}, " + f"enrichment={row['val_enrichment_loss']:.6g}" + ), + enabled=cfg.progress, + ) + if int(cfg.patience) >= 0 and stale > int(cfg.patience): + break + + test_loader = _loader( + eval_dataset, + splits.get("test", []), + cfg.batch_size, + shuffle=False, + num_workers=cfg.num_workers, + pin_memory=pin_memory, + ) + test_metrics = _run_loader( + model, + test_loader, + device=device, + objective=objective, + mixed_precision=mixed_precision, + progress=cfg.progress, + progress_label="RBPNet test", + return_predictions=True, + ) + if test_metrics["indices"]: + _write_predictions( + out / "test_predictions.csv", + test_metrics, + depth_offsets=eval_dataset.depth_offsets, + replicate_names=eval_dataset.replicate_names, + ) + (out / "history.json").write_text(json.dumps(history, indent=2), encoding="utf-8") + (out / "splits.json").write_text(json.dumps(splits, indent=2), encoding="utf-8") + summary = { + "trainer": "rbpnet", + "best_epoch": best_epoch, + "monitor": list(monitors), + "best_monitor_values": best_values, + "epochs_run": len(history), + "loss": loss_config.to_dict(enrichment_enabled=model.enrichment_enabled), + "optimizer": optimizer_config, + "lr_scheduler": scheduler_config, + "mixed_precision": mixed_precision, + "max_train_jitter": train_dataset.max_train_jitter, + "split_source_used": split_source, + "split_group_col": group_col, + "unsafe_random_window_split": random_split, + "split_counts": split_counts, + "deduplicate_loci": bool(getattr(cfg, "deduplicate_loci", True)), + "n_deduplicated_rows": n_deduplicated, + "parameter_count": parameter_count, + "receptive_field": model.receptive_field, + "receptive_field_extents": list(model.receptive_field_extents), + "replicate_names": list(eval_dataset.replicate_names), + "depth_offsets": eval_dataset.depth_offsets.tolist(), + "test_loss": float(test_metrics["loss"]), + "test_ip_profile_loss": float(test_metrics["ip_profile_loss"]), + "test_sm_profile_loss": float(test_metrics["sm_profile_loss"]), + "test_enrichment_loss": float(test_metrics["enrichment_loss"]), + } + (out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + return {"model": model, "history": history, "splits": splits, "summary": summary} + + +@torch.no_grad() +def evaluate_rbpnet_model( + model: RBPNet, + bundle: DatasetBundle, + *, + indices: Sequence[int] | None = None, + batch_size: int = 128, + device: str | torch.device = "cpu", + loss_config: str | Mapping[str, object] | None = None, + progress: bool = True, +) -> dict[str, Any]: + """Deterministically evaluate structured profile/enrichment likelihoods.""" + + resolved_device = resolve_device(device) + model = model.to(resolved_device) + dataset = RBPNetDataset( + bundle, + max_train_jitter=0, + training=False, + require_full_measurement_interval=model.enrichment_enabled, + ) + selected = list(range(len(dataset))) if indices is None else [int(i) for i in indices] + loader = _loader( + dataset, + selected, + batch_size, + shuffle=False, + num_workers=0, + pin_memory=resolved_device.type == "cuda", + ) + objective = RBPNetObjective( + RBPNetLossConfig.from_config(loss_config), + enrichment_enabled=model.enrichment_enabled, + ).to(resolved_device) + metrics = _run_loader( + model, + loader, + device=resolved_device, + objective=objective, + progress=progress, + progress_label="evaluate RBPNet", + return_predictions=True, + ) + metrics["depth_offsets"] = dataset.depth_offsets + metrics["replicate_names"] = dataset.replicate_names + return metrics + + +def write_rbpnet_predictions( + path: str | Path, + metrics: Mapping[str, Any], +) -> None: + """Write pi, eta, and replicate-specific predicted IP fractions.""" + + _write_predictions( + path, + metrics, + depth_offsets=np.asarray(metrics["depth_offsets"], dtype=np.float64), + replicate_names=metrics["replicate_names"], + ) diff --git a/src/transcriptml/training/__init__.py b/src/transcriptml/training/__init__.py index aa57922..2a26d47 100644 --- a/src/transcriptml/training/__init__.py +++ b/src/transcriptml/training/__init__.py @@ -2,13 +2,19 @@ from transcriptml.training.evaluation import evaluate_checkpoint, evaluate_fold_checkpoints, predict_to_csv from transcriptml.training.metrics import mse, pearson_corr -from transcriptml.training.splits import predefined_split_indices, random_split_indices +from transcriptml.training.splits import ( + group_split_indices, + predefined_split_indices, + random_split_indices, + validate_group_disjoint, +) from transcriptml.training.trainer import TrainConfig, train_from_config, train_model __all__ = [ "TrainConfig", "evaluate_checkpoint", "evaluate_fold_checkpoints", + "group_split_indices", "mse", "pearson_corr", "predict_to_csv", @@ -16,4 +22,5 @@ "random_split_indices", "train_from_config", "train_model", + "validate_group_disjoint", ] diff --git a/src/transcriptml/training/evaluation.py b/src/transcriptml/training/evaluation.py index 39531a7..4fb70f9 100644 --- a/src/transcriptml/training/evaluation.py +++ b/src/transcriptml/training/evaluation.py @@ -498,7 +498,7 @@ def evaluate_checkpoint( device = resolve_device(device) log_progress(f"evaluate: loading checkpoint {checkpoint_path}", enabled=progress) - model, _ = load_checkpoint(checkpoint_path, map_location=device) + model, checkpoint = load_checkpoint(checkpoint_path, map_location=device) log_progress(f"evaluate: loading dataset {dataset_path}", enabled=progress) bundle = load_bundle(dataset_path, mmap_mode="r") indices = None @@ -506,6 +506,36 @@ def evaluate_checkpoint( if not bundle.splits or split not in bundle.splits: raise ValueError(f"Dataset has no split '{split}'") indices = [int(i) for i in bundle.splits[split]] + if checkpoint.get("model_config", {}).get("name") == "rbpnet": + from transcriptml.models.rbpnet import RBPNet + from transcriptml.rbpnet.training import ( + evaluate_rbpnet_model, + write_rbpnet_predictions, + ) + + if not isinstance(model, RBPNet): + raise TypeError("rbpnet checkpoint did not reconstruct an RBPNet model") + result = evaluate_rbpnet_model( + model, + bundle, + indices=indices, + batch_size=batch_size, + device=device, + loss_config=checkpoint.get("loss_config"), + progress=progress, + ) + if out_csv is not None: + log_progress(f"evaluate: writing RBPNet predictions to {out_csv}", enabled=progress) + write_rbpnet_predictions(out_csv, result) + # Keep CLI summary serialization compact while preserving the scalar + # prediction convention for callers that expect a ``predictions`` key. + result["predictions"] = ( + result["enrichment_logit"] + if result.get("enrichment_logit") is not None + else result["pi"] + ) + result["targets"] = None + return result log_progress( f"evaluate: running on {len(indices) if indices is not None else bundle.X.shape[0]} examples", enabled=progress, diff --git a/src/transcriptml/training/splits.py b/src/transcriptml/training/splits.py index c3a4f91..6c4e4df 100644 --- a/src/transcriptml/training/splits.py +++ b/src/transcriptml/training/splits.py @@ -99,6 +99,91 @@ def predefined_split_indices( return splits +def group_split_indices( + metadata: Sequence[Mapping[str, object]], + *, + group_col: str = "group_gene_id", + val_frac: float = 0.1, + test_frac: float = 0.1, + seed: int | None = None, +) -> dict[str, list[int]]: + """Split complete biological groups while approximately balancing rows. + + This is the safe default for overlapping RBPNet windows. Groups are + shuffled reproducibly and then assigned to test, validation, and training + without ever dividing a group between splits. + """ + + if not metadata: + raise ValueError("metadata must contain at least one example") + if not (0 <= val_frac < 1) or not (0 <= test_frac < 1): + raise ValueError("val_frac and test_frac must be in [0, 1)") + if val_frac + test_frac >= 1: + raise ValueError("val_frac + test_frac must be < 1") + groups: dict[str, list[int]] = {} + for index, row in enumerate(metadata): + value = row.get(group_col) + if value is None or str(value).strip() == "": + raise ValueError( + f"Missing group column '{group_col}' for example index {index}" + ) + groups.setdefault(str(value), []).append(index) + if len(groups) < 1 + int(val_frac > 0) + int(test_frac > 0): + raise ValueError("not enough biological groups for requested train/val/test splits") + + keys = np.asarray(sorted(groups), dtype=object) + np.random.default_rng(seed).shuffle(keys) + target_test = int(round(test_frac * len(metadata))) + target_val = int(round(val_frac * len(metadata))) + if test_frac > 0: + target_test = max(1, target_test) + if val_frac > 0: + target_val = max(1, target_val) + splits = {"train": [], "val": [], "test": []} + for key in keys.tolist(): + rows = groups[str(key)] + if test_frac > 0 and len(splits["test"]) < target_test: + destination = "test" + elif val_frac > 0 and len(splits["val"]) < target_val: + destination = "val" + else: + destination = "train" + splits[destination].extend(rows) + if not splits["train"]: + raise ValueError("group split leaves no training examples") + _check_no_overlap(splits) + validate_group_disjoint(splits, metadata, group_col=group_col) + return splits + + +def validate_group_disjoint( + splits: Mapping[str, Sequence[int]], + metadata: Sequence[Mapping[str, object]], + *, + group_col: str = "group_gene_id", +) -> None: + """Raise when one biological group occurs in multiple dataset splits.""" + + owner: dict[str, str] = {} + for split_name, indices in splits.items(): + for raw_index in indices: + index = int(raw_index) + if index < 0 or index >= len(metadata): + raise ValueError(f"split index {index} is outside metadata bounds") + value = metadata[index].get(group_col) + if value is None or str(value).strip() == "": + raise ValueError( + f"Missing group column '{group_col}' for example index {index}" + ) + group = str(value) + previous = owner.setdefault(group, split_name) + if previous != split_name: + raise ValueError( + f"Biological group {group!r} appears in both '{previous}' and " + f"'{split_name}' using metadata column '{group_col}'" + ) + + def normalize_splits(splits: Mapping[str, Sequence[int]]) -> dict[str, list[int]]: """Normalize split indices to mutable integer lists with standard keys. diff --git a/src/transcriptml/training/trainer.py b/src/transcriptml/training/trainer.py index 819fc54..c9506b0 100644 --- a/src/transcriptml/training/trainer.py +++ b/src/transcriptml/training/trainer.py @@ -33,6 +33,9 @@ class TrainConfig: epochs: int = 20 learning_rate: float = 1e-3 weight_decay: float = 0.0 + optimizer: str | Mapping[str, Any] = "adamw" + lr_scheduler: str | Mapping[str, Any] | None = None + mixed_precision: bool = False gradient_clip_norm: float | None = 0.5 patience: int = 5 monitor: str | Sequence[str] = "val_loss" @@ -46,6 +49,9 @@ class TrainConfig: head_layernorm: bool = False sequence_controls: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None split_source: str = "auto" + max_train_jitter: int = 0 + allow_random_window_split: bool = False + deduplicate_loci: bool = True split: Mapping[str, Any] = field( default_factory=lambda: {"method": "random", "val_frac": 0.1, "test_frac": 0.1} ) @@ -506,6 +512,15 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) """ cfg = _as_train_config(config) + requested_model = normalize_model_config(cfg.model) + if requested_model.name == "rbpnet": + if cfg.sequence_controls: + raise ValueError( + "generic sequence_controls are not supported for structured RBPNet training" + ) + from transcriptml.rbpnet.training import train_rbpnet_model + + return train_rbpnet_model(bundle, cfg) _seed_everything(cfg.seed) device = resolve_device(cfg.device) out = Path(cfg.output_dir) diff --git a/src/transcriptml/workflows/init_run.py b/src/transcriptml/workflows/init_run.py index f278457..a1ea4de 100644 --- a/src/transcriptml/workflows/init_run.py +++ b/src/transcriptml/workflows/init_run.py @@ -50,7 +50,48 @@ def _train_config(workflow: str) -> dict[str, Any]: "split_source": "auto", "split": {"method": "random", "val_frac": 0.1, "test_frac": 0.1}, } - raise ValueError("workflow must be one of: saluki, legnet") + if workflow == "rbpnet": + return { + "dataset": "__EDIT_ME_RBPNET_BUNDLE_DIR__", + "output_dir": "__EDIT_ME_RUN_DIR__/model", + "model": { + "name": "rbpnet", + "params": { + "profile_length": 300, + "enrichment_head_type": "none", + }, + }, + "batch_size": 64, + "epochs": 100, + "learning_rate": 0.001, + "weight_decay": 0.0, + "optimizer": {"name": "adamw"}, + "lr_scheduler": {"name": "reduce_on_plateau", "patience": 3}, + "mixed_precision": False, + "gradient_clip_norm": 0.5, + "patience": 10, + "monitor": "val_loss", + "loss": { + "name": "rbpnet", + "lambda_ip_profile": 1.0, + "lambda_sm_profile": 1.0, + "lambda_enrichment": 1.0, + }, + "device": "auto", + "num_workers": 0, + "mmap_mode": "r", + "seed": 123, + "max_train_jitter": 0, + "deduplicate_loci": True, + "split_source": "config", + "split": { + "method": "group", + "group_col": "group_gene_id", + "val_frac": 0.1, + "test_frac": 0.1, + }, + } + raise ValueError("workflow must be one of: saluki, legnet, rbpnet") def _readme(workflow: str) -> str: @@ -73,14 +114,14 @@ def init_run(workflow: str, out_dir: str | Path, *, force: bool = False) -> Path """Write starter configs for a TranscriptML run. Args: - workflow: Workflow template name, ``saluki`` or ``legnet``. + workflow: Workflow template name: ``saluki``, ``legnet``, or ``rbpnet``. out_dir: Directory to create or populate. force: Allow writing into a non-empty output directory. """ workflow = str(workflow).strip().lower() - if workflow not in {"saluki", "legnet"}: - raise ValueError("workflow must be one of: saluki, legnet") + if workflow not in {"saluki", "legnet", "rbpnet"}: + raise ValueError("workflow must be one of: saluki, legnet, rbpnet") out = Path(out_dir) if out.exists() and any(out.iterdir()) and not force: raise FileExistsError(f"Output directory is not empty: {out}. Use --force to overwrite template files.") diff --git a/tests/test_cli_analysis.py b/tests/test_cli_analysis.py index fdcac0b..e116536 100644 --- a/tests/test_cli_analysis.py +++ b/tests/test_cli_analysis.py @@ -45,6 +45,15 @@ def test_init_run_cli_writes_templates(tmp_path): assert not (out_dir / "run_config.json").exists() assert (out_dir / "README.md").exists() + rbpnet_dir = tmp_path / "rbpnet_run" + main(["init-run", "--workflow", "rbpnet", "--out-dir", str(rbpnet_dir)]) + rbpnet = json.loads((rbpnet_dir / "train_config.json").read_text(encoding="utf-8")) + assert rbpnet["model"]["name"] == "rbpnet" + assert rbpnet["model"]["params"]["profile_length"] == 300 + assert rbpnet["split"]["method"] == "group" + assert rbpnet["split"]["group_col"] == "group_gene_id" + assert rbpnet["max_train_jitter"] == 0 + def test_plot_ism_cli_demo_writes_png(tmp_path): out_path = tmp_path / "ism.png" diff --git a/tests/test_rbpnet_model.py b/tests/test_rbpnet_model.py new file mode 100644 index 0000000..39b4220 --- /dev/null +++ b/tests/test_rbpnet_model.py @@ -0,0 +1,431 @@ +import math +from dataclasses import replace + +import numpy as np +import pytest +import torch + +from transcriptml.data.bundle import DatasetBundle, save_bundle +from transcriptml.models.rbpnet import ( + RBPNet, + SameLengthConvTranspose1d, + SamePadConv1d, +) +from transcriptml.models.registry import build_model +from transcriptml.rbpnet.dataset import RBPNetDataset, collate_rbpnet +from transcriptml.rbpnet.losses import ( + RBPNetObjective, + multinomial_nll, + replicate_binomial_nll, +) +from transcriptml.training.splits import group_split_indices, validate_group_disjoint +from transcriptml.training.evaluation import evaluate_checkpoint +from transcriptml.training.trainer import train_model + + +def _synthetic_bundle(n=9, length=16, jitter=2): + width = length + 2 * jitter + rng = np.random.default_rng(8) + bases = rng.integers(0, 4, size=(n, width)) + X = np.zeros((n, 4, width), dtype=np.uint8) + for index in range(n): + X[index, bases[index], np.arange(width)] = 1 + ip = rng.poisson(0.5, size=(n, 2, width)).astype(np.uint32) + sm = rng.poisson(0.7, size=(n, width)).astype(np.uint32) + metadata = [] + selection_ip = np.zeros((n, 2), dtype=np.uint64) + selection_sm = np.zeros(n, dtype=np.uint64) + for index in range(n): + if index == 0: + anchor, materialized_start, selection = 1, 0, (0, 4) + elif index == 1: + anchor, materialized_start, selection = 98, 80, (96, 100) + else: + anchor = 30 + index * 3 + materialized_start = anchor - width // 2 + selection = (anchor - 2, anchor + 3 + index % 2) + local_start = selection[0] - materialized_start + local_end = selection[1] - materialized_start + selection_ip[index] = ip[index, :, local_start:local_end].sum(axis=1) + selection_sm[index] = sm[index, local_start:local_end].sum() + metadata.append( + { + "example_id": f"ex{index}", + "gene_id": f"g{index // 3}", + "transcript_id": f"tx{index // 3}", + "chromosome": "chr1", + "strand": "-" if index % 2 else "+", + "coordinate_space": "mature_transcript", + "locus_length": 100, + "transcript_anchor": anchor, + "selection_start": selection[0], + "selection_end": selection[1], + "sequence_materialized_start": materialized_start, + "sequence_materialized_end": materialized_start + width, + "profile_materialized_start": materialized_start, + "profile_materialized_end": materialized_start + width, + "group_gene_id": f"g{index // 3}", + "group_transcript_id": f"tx{index // 3}", + "group_chromosome": "chr1", + } + ) + arrays = { + "sminput_profiles": sm, + "ip_profiles": ip, + "sequence_valid_mask": np.ones((n, width), dtype=np.uint8), + "profile_valid_mask": np.ones((n, width), dtype=np.uint8), + "profile_sminput_totals": sm.sum(axis=1, dtype=np.uint64), + "profile_ip_totals": ip.sum(axis=2, dtype=np.uint64), + "selection_sminput_counts": selection_sm, + "selection_ip_counts": selection_ip, + } + return DatasetBundle( + X=X, + ids=[f"ex{index}" for index in range(n)], + metadata=metadata, + arrays=arrays, + config={ + "bundle_format": "transcriptml-rbpnet-bundle", + "bundle_format_version": "1", + "coordinate_space": "mature_transcript", + "input_length": length, + "profile_length": length, + "max_jitter": jitter, + "materialized_sequence_length": width, + "materialized_profile_length": width, + "transcript_end_policy": "shift_to_fit", + "sample_metadata": { + "sminput": {"name": "sminput", "effective_library_size": 100}, + "ip": [ + {"name": "ip1", "effective_library_size": 50}, + {"name": "ip2", "effective_library_size": 200}, + ], + "ip_axis_order": ["ip1", "ip2"], + }, + }, + ) + + +def _small_model(*, enrichment="linear"): + return RBPNet( + n_filters=8, + initial_kernel_size=3, + n_residual_blocks=1, + residual_kernel_size=3, + dilations=[1], + normalization="none", + dropout=0.0, + profile_head_kernel_size=3, + enrichment_head_type=enrichment, + profile_length=16, + ) + + +def test_default_rbpnet_architecture_probabilities_and_receptive_field(): + model = RBPNet() + assert len(model.residual_blocks) == 5 + assert model.dilations == (2, 4, 8, 16, 32) + assert model.receptive_field == 322 + assert model.receptive_field_extents == (160, 161) + assert model.target_profile_head.head_type == "transpose_conv" + output = model(torch.randn(2, 4, 300)) + assert output.target_logits.shape == output.control_logits.shape == (2, 300) + assert output.mixing_logit.shape == output.pi.shape == (2,) + assert output.enrichment_logit is None + torch.testing.assert_close(output.target_probs.sum(dim=-1), torch.ones(2)) + torch.testing.assert_close(output.control_probs.sum(dim=-1), torch.ones(2)) + torch.testing.assert_close(output.ip_probs.sum(dim=-1), torch.ones(2)) + assert torch.all((output.pi >= 0) & (output.pi <= 1)) + expected = ( + output.pi[:, None] * output.target_probs + + (1 - output.pi[:, None]) * output.control_probs + ) + torch.testing.assert_close(output.ip_probs, expected) + + +def test_architecture_override_same_length_and_indexed_alignment(): + model = build_model( + { + "name": "rbpnet", + "params": { + "n_filters": 7, + "initial_kernel_size": 4, + "n_residual_blocks": 2, + "residual_kernel_size": 4, + "dilations": [1, 3], + "normalization": "layer", + "dropout": 0.0, + "profile_head_type": "conv", + "profile_head_kernel_size": 6, + "profile_length": 31, + "enrichment_head_type": "mlp", + "enrichment_hidden": 5, + }, + } + ) + output = model(torch.randn(3, 4, 31), measurement_mask=torch.ones(3, 31)) + assert model.receptive_field == 16 + assert output.ip_probs.shape == (3, 31) + assert output.enrichment_logit.shape == (3,) + + conv = SamePadConv1d(1, 1, 12, bias=False) + conv.conv.weight.data.zero_() + conv.conv.weight.data[0, 0, conv.padding[0]] = 1 + impulse = torch.zeros(1, 1, 25) + impulse[0, 0, 11] = 1 + torch.testing.assert_close(conv(impulse), impulse) + + transpose = SameLengthConvTranspose1d(1, 1, 6, bias=False) + transpose.conv.weight.data.zero_() + transpose.conv.weight.data[0, 0, transpose.crop[0]] = 1 + torch.testing.assert_close(transpose(impulse), impulse) + + residual_model = _small_model(enrichment="none") + block = residual_model.residual_blocks[0] + block.conv.conv.weight.data.zero_() + if block.conv.conv.bias is not None: + block.conv.conv.bias.data.zero_() + hidden = torch.randn(2, 8, 16) + torch.testing.assert_close(block(hidden), hidden) + + +def test_multinomial_and_binomial_likelihoods_zero_edges_and_offsets(): + probabilities = torch.tensor([[0.25, 0.75], [0.5, 0.5]]) + counts = torch.tensor([[1.0, 2.0], [0.0, 0.0]]) + result = multinomial_nll(probabilities.log(), counts) + expected = -math.log(3 * 0.25 * 0.75**2) + assert result.denominator.item() == 1 + assert result.loss.item() == pytest.approx(expected) + empty = multinomial_nll(torch.log(torch.tensor([[0.5, 0.5]])), torch.zeros(1, 2)) + assert empty.loss.item() == 0 + assert empty.denominator.item() == 0 + + eta = torch.tensor([0.0, math.log(2.0)]) + ip = torch.tensor([[0.0, 3.0], [4.0, 0.0]]) + sm = torch.tensor([3.0, 0.0]) + offsets = torch.tensor([0.0, math.log(2.0)]) + observed = replicate_binomial_nll(eta, ip, sm, offsets) + logits = eta[:, None] + offsets[None, :] + total = ip + sm[:, None] + expected_values = -torch.distributions.Binomial( + total_count=total, logits=logits + ).log_prob(ip) + valid = total > 0 + assert observed.loss.item() == pytest.approx(expected_values[valid].mean().item()) + assert logits[0, 0].item() == 0.0 + assert logits[0, 1].item() == pytest.approx(math.log(2.0)) + assert torch.isfinite(observed.per_observation).all() + + +def test_dataset_measurement_masks_jitter_boundaries_and_deterministic_eval(): + bundle = _synthetic_bundle() + training = RBPNetDataset(bundle, max_train_jitter=2, training=True, seed=11) + left_minus = training.item_for_shift(0, -2) + left_plus = training.item_for_shift(0, 2) + assert left_minus["crop_start"] == left_plus["crop_start"] == 0 + assert left_minus["measurement_mask"].sum() == 4 + right = training.item_for_shift(1, 2) + assert right["crop_start"] == 84 + assert right["measurement_mask"].sum() == 4 + variable = training.item_for_shift(3, 0) + assert variable["measurement_mask"].sum() == 6 + np.testing.assert_array_equal( + left_minus["pooled_ip_profile"], + left_minus["individual_ip_profiles"].sum(axis=0), + ) + assert training.depth_offsets.tolist() == pytest.approx( + [math.log(0.5), math.log(2.0)] + ) + + evaluation = RBPNetDataset(bundle, max_train_jitter=0, training=False) + first = evaluation[4] + second = evaluation[4] + assert first["jitter_shift"] == second["jitter_shift"] == 0 + np.testing.assert_array_equal(first["sequence"], second["sequence"]) + training.set_epoch(7) + a = training[4] + training.set_epoch(7) + b = training[4] + assert a["jitter_shift"] == b["jitter_shift"] + np.testing.assert_array_equal(a["sequence"], b["sequence"]) + + interior_left = training.item_for_shift(4, -2) + interior_right = training.item_for_shift(4, 2) + materialized_start = bundle.metadata[4]["sequence_materialized_start"] + left_offset = int(interior_left["crop_start"]) - materialized_start + right_offset = int(interior_right["crop_start"]) - materialized_start + assert (left_offset, right_offset) == (0, 4) + np.testing.assert_array_equal( + interior_left["sequence"], bundle.X[4, :, left_offset : left_offset + 16] + ) + np.testing.assert_array_equal( + interior_right["individual_ip_profiles"], + bundle.arrays["ip_profiles"][4, :, right_offset : right_offset + 16], + ) + assert np.flatnonzero(interior_left["measurement_mask"])[0] == ( + np.flatnonzero(interior_right["measurement_mask"])[0] + 4 + ) + + +def test_enrichment_pooling_uses_weighted_variable_measurement_mask(): + model = _small_model(enrichment="linear") + hidden = torch.arange(2 * 8 * 16, dtype=torch.float32).reshape(2, 8, 16) + masks = torch.zeros(2, 16) + masks[0, 2:5] = 1 + masks[1, 4:10] = torch.tensor([1, 1, 2, 2, 1, 1], dtype=torch.float32) + pooled = model._pool_measurement(hidden, masks) + torch.testing.assert_close(pooled[0], hidden[0, :, 2:5].mean(dim=-1)) + expected_second = (hidden[1, :, 4:10] * masks[1, 4:10]).sum(dim=-1) / 8 + torch.testing.assert_close(pooled[1], expected_second) + + +def test_loss_pools_ip_replicates_and_gradients_reach_every_enabled_head(): + dataset = RBPNetDataset(_synthetic_bundle(), max_train_jitter=0, training=False) + batch = collate_rbpnet([dataset[index] for index in range(3)]) + model = _small_model(enrichment="linear") + output = model( + batch.sequence, + measurement_mask=batch.measurement_mask, + profile_mask=batch.profile_valid_mask, + ) + objective = RBPNetObjective(enrichment_enabled=True) + loss = objective(output, batch) + assert torch.isfinite(loss.loss) + loss.loss.backward() + for module in ( + model.initial_conv, + model.target_profile_head, + model.control_profile_head, + model.mixing_head, + model.enrichment_head, + ): + assert module is not None + assert any(parameter.grad is not None for parameter in module.parameters()) + + # Enrichment consumes eta and known offsets only; changing pi leaves its + # likelihood unchanged, i.e. there is deliberately no pi-BNLL. + enrichment_a = objective(output, batch).components["enrichment_loss"] + changed_pi = replace( + output, + mixing_logit=torch.full_like(output.mixing_logit, 100.0), + pi=torch.ones_like(output.pi), + ) + enrichment_b = objective(changed_pi, batch).components["enrichment_loss"] + torch.testing.assert_close(enrichment_a, enrichment_b) + + +def test_group_split_and_structured_training_profile_only_and_enrichment(tmp_path): + bundle = _synthetic_bundle() + splits = group_split_indices( + bundle.metadata, group_col="group_gene_id", val_frac=0.2, test_frac=0.2, seed=3 + ) + validate_group_disjoint(splits, bundle.metadata, group_col="group_gene_id") + owners = {} + for split, indices in splits.items(): + for index in indices: + group = bundle.metadata[index]["group_gene_id"] + assert group not in owners or owners[group] == split + owners[group] = split + + base_config = { + "dataset": "unused", + "batch_size": 3, + "epochs": 1, + "patience": 0, + "progress": False, + "learning_rate": 0.005, + "max_train_jitter": 2, + "model": { + "name": "rbpnet", + "params": { + "n_filters": 8, + "initial_kernel_size": 3, + "n_residual_blocks": 1, + "residual_kernel_size": 3, + "dilations": [1], + "normalization": "none", + "dropout": 0.0, + "profile_head_kernel_size": 3, + "profile_length": 16, + }, + }, + "loss": {"name": "rbpnet"}, + "split_source": "config", + "split": { + "method": "group", + "group_col": "group_gene_id", + "val_frac": 0.2, + "test_frac": 0.2, + "seed": 3, + }, + } + profile_config = dict(base_config) + profile_config["output_dir"] = str(tmp_path / "profile") + profile_config["model"] = { + "name": "rbpnet", + "params": {**base_config["model"]["params"], "enrichment_head_type": "none"}, + } + profile = train_model(bundle, profile_config) + assert profile["summary"]["trainer"] == "rbpnet" + assert profile["summary"]["loss"]["effective_lambda_enrichment"] == 0 + assert np.isfinite(profile["history"][0]["train_loss"]) + + enrichment_config = dict(base_config) + enrichment_config["output_dir"] = str(tmp_path / "enrichment") + enrichment_config["model"] = { + "name": "rbpnet", + "params": {**base_config["model"]["params"], "enrichment_head_type": "linear"}, + } + enriched = train_model(bundle, enrichment_config) + assert enriched["summary"]["loss"]["effective_lambda_enrichment"] == 1 + assert np.isfinite(enriched["history"][0]["train_enrichment_loss"]) + assert (tmp_path / "enrichment" / "best.pt").is_file() + assert (tmp_path / "enrichment" / "test_predictions.csv").is_file() + + bundle_dir = tmp_path / "bundle" + save_bundle(bundle, bundle_dir) + predictions_path = tmp_path / "checkpoint_predictions.csv" + evaluated = evaluate_checkpoint( + tmp_path / "enrichment" / "best.pt", + bundle_dir, + predictions_path, + batch_size=3, + progress=False, + ) + assert evaluated["pi"].shape == (len(bundle.ids),) + assert evaluated["enrichment_logit"].shape == (len(bundle.ids),) + assert predictions_path.is_file() + + unsafe = dict(base_config) + unsafe["output_dir"] = str(tmp_path / "unsafe") + unsafe["split"] = {"method": "random", "val_frac": 0.2, "test_frac": 0.2} + with pytest.raises(ValueError, match="row-level random splitting is unsafe"): + train_model(bundle, unsafe) + + +def test_tiny_batch_can_overfit(): + torch.manual_seed(4) + dataset = RBPNetDataset(_synthetic_bundle(n=3), max_train_jitter=0, training=False) + batch = collate_rbpnet([dataset[index] for index in range(3)]) + model = _small_model(enrichment="linear") + objective = RBPNetObjective(enrichment_enabled=True) + optimizer = torch.optim.Adam(model.parameters(), lr=0.02) + + def step(update): + output = model( + batch.sequence, + measurement_mask=batch.measurement_mask, + profile_mask=batch.profile_valid_mask, + ) + value = objective(output, batch).loss + if update: + optimizer.zero_grad() + value.backward() + optimizer.step() + return float(value.detach()) + + initial = step(False) + for _ in range(25): + step(True) + final = step(False) + assert final < initial From 3db81335704e3a89fbe7b8b2a97e7cb1ffb2e008 Mon Sep 17 00:00:00 2001 From: isvock Date: Wed, 12 Aug 2026 10:15:02 -0700 Subject: [PATCH 05/12] Optimized signal writing and CV splitting --- docs/api.rst | 4 + docs/rbpnet.md | 33 +++ docs/training_configuration.md | 46 +++ scripts/README.md | 1 + scripts/rbpnet/README.md | 48 ++++ scripts/rbpnet/create_chromosome_cv_plan.sh | 32 +++ scripts/rbpnet/example_train_config.json | 40 +++ scripts/rbpnet/preprocess.sh | 48 ++++ scripts/rbpnet/rbpnet_config.sh | 96 +++++++ scripts/rbpnet/scan_select_bundle.sh | 97 +++++++ scripts/rbpnet/submit_train_cv.sh | 21 ++ scripts/rbpnet/train_cv_fold.sh | 48 ++++ src/transcriptml/cli/main.py | 70 ++++- src/transcriptml/rbpnet/cli.py | 17 ++ src/transcriptml/rbpnet/preprocessing.py | 6 + src/transcriptml/rbpnet/signals.py | 177 ++++++++++-- src/transcriptml/rbpnet/training.py | 35 +++ src/transcriptml/training/trainer.py | 55 +++- src/transcriptml/workflows/__init__.py | 21 +- src/transcriptml/workflows/chromosome_cv.py | 297 ++++++++++++++++++++ tests/test_chromosome_cv.py | 184 ++++++++++++ tests/test_rbpnet.py | 51 ++++ tests/test_rbpnet_model.py | 54 +++- 23 files changed, 1460 insertions(+), 21 deletions(-) create mode 100644 scripts/rbpnet/README.md create mode 100644 scripts/rbpnet/create_chromosome_cv_plan.sh create mode 100644 scripts/rbpnet/example_train_config.json create mode 100644 scripts/rbpnet/preprocess.sh create mode 100644 scripts/rbpnet/rbpnet_config.sh create mode 100644 scripts/rbpnet/scan_select_bundle.sh create mode 100644 scripts/rbpnet/submit_train_cv.sh create mode 100644 scripts/rbpnet/train_cv_fold.sh create mode 100644 src/transcriptml/workflows/chromosome_cv.py create mode 100644 tests/test_chromosome_cv.py diff --git a/docs/api.rst b/docs/api.rst index 81810f4..6439632 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -173,6 +173,10 @@ Plotting Run setup --------- +.. automodule:: transcriptml.workflows.chromosome_cv + :members: ChromosomeCVPlan, ChromosomeCVResolution, create_chromosome_cv_plan, save_chromosome_cv_plan, load_chromosome_cv_plan, resolve_chromosome_cv_plan + :member-order: bysource + .. automodule:: transcriptml.workflows.init_run :members: init_run :member-order: bysource diff --git a/docs/rbpnet.md b/docs/rbpnet.md index 10042ee..2385c4c 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -130,6 +130,15 @@ Each sample's `effective_library_size` is exactly the number of retained read1 Pooled-IP CPM uses the sum of IP counts divided by the sum of IP effective library sizes. +Signal rows are sparse-aggregated and then written one complete touched HDF5 +chunk at a time. Zero-only chunks retain the HDF5 fill value and are not +allocated. The default is shuffled gzip level 1, chosen to reduce write and +slice-decompression time while retaining compact sparse storage. Use +`--signal-compression gzip --signal-compression-level 4` for smaller but slower +files, `--signal-compression lzf` for faster/larger files, or +`--signal-compression none` when storage is unimportant. Compression and chunk +length are recorded as `signals.h5` attributes and in preprocessing provenance. + ### HDF5 signal layout HDF5 is a hierarchical binary container: datasets behave like typed, @@ -548,6 +557,30 @@ rejected unless `allow_random_window_split=true` explicitly acknowledges the leakage risk. Replicate-specific selection rows describing an identical locus are deduplicated by default while retaining the complete replicate axis. +For chromosome cross-validation, create one plan from the final bundle and +reuse it for every training job: + +```bash +transcriptml cv create-chromosome-plan \ + --dataset data/rbpnet \ + --group-col group_chromosome \ + --n-folds 5 \ + --output runs/rbpnet/cv5.json + +transcriptml train configs/rbpnet/train_config.json \ + --dataset data/rbpnet \ + --cv-plan runs/rbpnet/cv5.json \ + --fold 0 \ + --output-dir runs/rbpnet/fold0/model +``` + +Chromosomes are sorted by decreasing example count and greedily assigned to +the currently smallest fold group, with deterministic ties. Run `k` uses group +`k` for test, `(k+1) mod N` for validation, and every remaining group for +training. Resolution verifies that the dataset's chromosome membership and +counts still exactly match the immutable, content-hashed plan. Training records +both the plan path and its validated `plan_id` in summaries and checkpoints. + AdamW, Adam, and SGD; plateau, cosine, and step schedulers; clipping; early stopping; device selection; DataLoader workers; seeds; and mixed precision are configurable. `history.json` logs total, pooled-IP profile, SMInput profile, and diff --git a/docs/training_configuration.md b/docs/training_configuration.md index 171b9e8..812cbd0 100644 --- a/docs/training_configuration.md +++ b/docs/training_configuration.md @@ -196,6 +196,8 @@ above. | `max_train_jitter` | integer | `0` | Structured RBPNet shift range; cannot exceed the bundle's materialized margin. Evaluation remains shift zero. | | `deduplicate_loci` | boolean | `true` | Collapse repeated eligibility rows for an identical RBPNet locus while retaining all replicate arrays. | | `allow_random_window_split` | boolean | `false` | Explicitly permit unsafe RBPNet row-random splitting. Grouped splitting is the safe default. | +| `cv_plan` | path or `null` | `null` | Saved balanced chromosome CV plan. Must be provided together with `fold`; it takes precedence over `split_source`. | +| `fold` | integer or `null` | `null` | Zero-based chromosome CV test-fold index. Validation is the following fold modulo `n_folds`. | The canonical model mapping contains a registered `name` and a `params` mapping: @@ -532,6 +534,50 @@ Cross-validation fold preparation writes `splits.json` inside each fold bundle. The usual CV workflow therefore uses those fold assignments under the default `"auto"` setting. +### Balanced chromosome CV plans + +Create one content-hashed plan from a dataset's chromosome grouping metadata: + +```bash +transcriptml cv create-chromosome-plan \ + --dataset data/rbpnet \ + --group-col group_chromosome \ + --n-folds 5 \ + --output cv/cv5.json +``` + +The JSON records the grouping column, example count on every chromosome, +chromosome membership and total examples for every fold group, algorithm and +tie-breaking rules, format and TranscriptML versions, and a SHA-256 `plan_id`. +Generation sorts chromosomes from largest to smallest and assigns each to the +group with the smallest current example count; ties use chromosome name and +then fold index deterministically. This balances examples rather than numbers +of chromosomes. Plans require at least three folds and at least one chromosome +per fold. The plan path and validated `plan_id` are recorded in training +summaries and checkpoints. + +For run `k`, test is group `k`, validation is `(k+1) mod N`, and training is all +remaining groups. Inspect or materialize one resolution with: + +```bash +transcriptml cv resolve-plan \ + --dataset data/rbpnet --cv-plan cv/cv5.json --fold 0 \ + --output cv/fold0_splits.json +``` + +Training accepts overrides suitable for a Slurm job array: + +```bash +transcriptml train configs/rbpnet/train_config.json \ + --dataset data/rbpnet \ + --cv-plan cv/cv5.json \ + --fold "${SLURM_ARRAY_TASK_ID}" \ + --output-dir "cv/fold${SLURM_ARRAY_TASK_ID}/model" +``` + +Resolution rejects missing chromosomes, new chromosomes, or changed example +counts rather than silently applying an obsolete plan. + ### Random Splits ```json diff --git a/scripts/README.md b/scripts/README.md index 5d3a989..085daeb 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -17,6 +17,7 @@ These scripts are intentionally Sherlock-specific and deliberately small. For a - `motif_ablation_by_fold.sh` and `motif_ablation_all_folds.sh`: motif ablations across the configured motif list. - `motif_epistasis_by_fold.sh` and `motif_epistasis_all_folds.sh`: motif epistasis across the configured motif-pair list. - `mpra/`: MPRA 3-prime UTR insert workflows for building 4-channel LegNet input, training LegNet, and running single-nucleotide ISM. See `mpra/README.md`. +- `rbpnet/`: eCLIP preprocessing, scan/selection/bundle construction, immutable balanced chromosome CV planning, and one RBPNet training job per chromosome fold. See `rbpnet/README.md`. ## Configure A Run diff --git a/scripts/rbpnet/README.md b/scripts/rbpnet/README.md new file mode 100644 index 0000000..6b71985 --- /dev/null +++ b/scripts/rbpnet/README.md @@ -0,0 +1,48 @@ +# RBPNet Sherlock workflow + +This directory implements the staged workflow: + +```text +FASTA/GTF/IP BAMs/SMInput BAM + -> canonical eCLIP preprocessing + -> window scan, selection, and RBPNet bundle + -> immutable balanced chromosome CV plan + -> one independent training job per fold +``` + +Copy this directory to a writable run directory and edit `rbpnet_config.sh`. +At minimum set `TRANSCRIPTML_REPO`, `GENOME_FASTA`, `GTF`, `SMINPUT_BAM`, and +the `IP_BAMS` array. Use `sample_name=path` values, for example: + +```bash +SMINPUT_BAM="sminput=/oak/project/PUM2_sminput.bam" +IP_BAMS=( + "ip1=/oak/project/PUM2_ip1.bam" + "ip2=/oak/project/PUM2_ip2.bam" +) +``` + +Then run or submit each stage in order: + +```bash +mkdir -p slurm_output +sbatch scripts/rbpnet/preprocess.sh +sbatch scripts/rbpnet/scan_select_bundle.sh +sbatch scripts/rbpnet/create_chromosome_cv_plan.sh +bash scripts/rbpnet/submit_train_cv.sh +``` + +The data-construction script chooses stride 1 automatically for +`original_rbpnet` and stride 50 for the other selectors unless +`WINDOW_STRIDE` is explicitly set. Its default selector is +`peak_gray_negative`; all thresholds remain editable in the config. + +The CV stage counts examples per chromosome, greedily balances whole +chromosomes across `N_FOLDS`, and writes `CV_PLAN` once. Every fold job loads +that same file. Fold `k` uses group `k` as test, group `(k+1) mod N` as +validation, and all other groups for training. `train_cv_fold.sh 0` can be run +interactively for one fold without Slurm. + +`example_train_config.json` enables the independent linear enrichment head and +32-nt training jitter. Change `enrichment_head_type` to `none` for profile-only +RBPNet, and keep `profile_length`/`max_train_jitter` consistent with the bundle. diff --git a/scripts/rbpnet/create_chromosome_cv_plan.sh b/scripts/rbpnet/create_chromosome_cv_plan.sh new file mode 100644 index 0000000..161dfad --- /dev/null +++ b/scripts/rbpnet/create_chromosome_cv_plan.sh @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --partition=akundaje +#SBATCH --job-name=tml_rbp_cvplan +#SBATCH --cpus-per-task=1 +#SBATCH --mem=16G +#SBATCH --time=01:00:00 +#SBATCH --output=slurm_output/%x_%j.out +#SBATCH --error=slurm_output/%x_%j.err + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/rbpnet/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts/rbpnet" + elif [[ -f "${SLURM_SUBMIT_DIR}/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/rbpnet_config.sh" +setup_transcriptml_env + +if [[ ! -f "${BUNDLE_DIR}/metadata.json" ]]; then + echo "Missing RBPNet bundle metadata at ${BUNDLE_DIR}; run scan_select_bundle.sh first." >&2 + exit 1 +fi + +transcriptml cv create-chromosome-plan \ + --dataset "${BUNDLE_DIR}" \ + --output "${CV_PLAN}" \ + --n-folds "${N_FOLDS}" \ + --group-col "${CHROMOSOME_GROUP_COL}" diff --git a/scripts/rbpnet/example_train_config.json b/scripts/rbpnet/example_train_config.json new file mode 100644 index 0000000..baf49ac --- /dev/null +++ b/scripts/rbpnet/example_train_config.json @@ -0,0 +1,40 @@ +{ + "dataset": "SCRIPT_OVERRIDES_THIS_FROM_BUNDLE_DIR", + "output_dir": "SCRIPT_OVERRIDES_THIS_FOR_EACH_FOLD", + "model": { + "name": "rbpnet", + "params": { + "profile_length": 300, + "enrichment_head_type": "linear" + } + }, + "batch_size": 64, + "epochs": 100, + "learning_rate": 0.001, + "weight_decay": 0.0, + "optimizer": {"name": "adamw"}, + "lr_scheduler": {"name": "reduce_on_plateau", "patience": 3}, + "mixed_precision": true, + "gradient_clip_norm": 0.5, + "patience": 10, + "monitor": "val_loss", + "loss": { + "name": "rbpnet", + "lambda_ip_profile": 1.0, + "lambda_sm_profile": 1.0, + "lambda_enrichment": 1.0 + }, + "device": "auto", + "num_workers": 4, + "mmap_mode": "r", + "seed": 123, + "max_train_jitter": 32, + "deduplicate_loci": true, + "split_source": "config", + "split": { + "method": "group", + "group_col": "group_gene_id", + "val_frac": 0.1, + "test_frac": 0.1 + } +} diff --git a/scripts/rbpnet/preprocess.sh b/scripts/rbpnet/preprocess.sh new file mode 100644 index 0000000..e202c52 --- /dev/null +++ b/scripts/rbpnet/preprocess.sh @@ -0,0 +1,48 @@ +#!/bin/bash +#SBATCH --partition=akundaje +#SBATCH --job-name=tml_rbp_preprocess +#SBATCH --cpus-per-task=4 +#SBATCH --mem=48G +#SBATCH --time=24:00:00 +#SBATCH --output=slurm_output/%x_%j.out +#SBATCH --error=slurm_output/%x_%j.err + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/rbpnet/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts/rbpnet" + elif [[ -f "${SLURM_SUBMIT_DIR}/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/rbpnet_config.sh" +setup_transcriptml_env + +if [[ -z "${GENOME_FASTA}" || -z "${GTF}" || -z "${SMINPUT_BAM}" || ${#IP_BAMS[@]} -eq 0 ]]; then + echo "Set GENOME_FASTA, GTF, SMINPUT_BAM, and at least one IP_BAMS entry in rbpnet_config.sh." >&2 + exit 1 +fi + +command=( + transcriptml rbpnet preprocess + --genome-fasta "${GENOME_FASTA}" + --gtf "${GTF}" + --sminput-bam "${SMINPUT_BAM}" + --coordinate-space "${COORDINATE_SPACE}" + --read1-rna-strand "${READ1_RNA_STRAND}" + --min-mapq "${MIN_MAPQ}" + --signal-compression "${SIGNAL_COMPRESSION}" + --output-dir "${PROCESSED_DIR}" +) +for bam in "${IP_BAMS[@]}"; do + command+=(--ip-bam "${bam}") +done +if [[ "${SIGNAL_COMPRESSION}" == "gzip" ]]; then + command+=(--signal-compression-level "${SIGNAL_COMPRESSION_LEVEL}") +fi +if [[ "${OVERWRITE}" == "1" ]]; then + command+=(--overwrite) +fi +"${command[@]}" diff --git a/scripts/rbpnet/rbpnet_config.sh b/scripts/rbpnet/rbpnet_config.sh new file mode 100644 index 0000000..0c24a79 --- /dev/null +++ b/scripts/rbpnet/rbpnet_config.sh @@ -0,0 +1,96 @@ +#!/bin/bash + +# Shared Sherlock defaults for the eCLIP -> RBPNet chromosome-CV workflow. +# Copy scripts/rbpnet to a writable run directory and edit this file there. + +_RBPNET_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_CONFIG_DIR="${SCRIPT_CONFIG_DIR:-${_RBPNET_SCRIPT_DIR}}" + +# Advanced option: source another shell config before applying defaults below. +if [[ -n "${TRANSCRIPTML_RUN_CONFIG:-}" ]]; then + if [[ ! -f "${TRANSCRIPTML_RUN_CONFIG}" ]]; then + echo "TRANSCRIPTML_RUN_CONFIG does not exist: ${TRANSCRIPTML_RUN_CONFIG}" >&2 + return 1 2>/dev/null || exit 1 + fi + TRANSCRIPTML_RUN_CONFIG_DIR="$(cd "$(dirname "${TRANSCRIPTML_RUN_CONFIG}")" && pwd)" + TRANSCRIPTML_RUN_CONFIG="${TRANSCRIPTML_RUN_CONFIG_DIR}/$(basename "${TRANSCRIPTML_RUN_CONFIG}")" + source "${TRANSCRIPTML_RUN_CONFIG}" +else + TRANSCRIPTML_RUN_CONFIG_DIR="" +fi + +CONDA_ENV="${CONDA_ENV:-transcript-ml}" +SHERLOCK_CONDA_ROOT="${SHERLOCK_CONDA_ROOT:-${GROUP_HOME:-${HOME}}/miniconda}" +TRANSCRIPTML_REPO="${TRANSCRIPTML_REPO:-}" +_TRANSCRIPTML_REPO_CANDIDATE="$(cd "${_RBPNET_SCRIPT_DIR}/../.." && pwd)" +if [[ -z "${TRANSCRIPTML_REPO}" && -d "${_TRANSCRIPTML_REPO_CANDIDATE}/src/transcriptml" ]]; then + TRANSCRIPTML_REPO="${_TRANSCRIPTML_REPO_CANDIDATE}" +fi + +# Standard assay inputs. SMINPUT_BAM and each IP_BAMS entry may use LABEL=PATH. +GENOME_FASTA="${GENOME_FASTA:-}" +GTF="${GTF:-}" +SMINPUT_BAM="${SMINPUT_BAM:-}" +if ! declare -p IP_BAMS >/dev/null 2>&1; then + IP_BAMS=() +fi + +RUN_NAME="${RUN_NAME:-RBPNet_eCLIP}" +RUN_ROOT="${RUN_ROOT:-/scratch/users/${USER:-user}/TranscriptML/${RUN_NAME}}" +PROCESSED_DIR="${PROCESSED_DIR:-${RUN_ROOT}/processed/eclip}" +WINDOW_PREFIX="${WINDOW_PREFIX:-${RUN_ROOT}/windows/windows_100nt}" +SELECTION_PREFIX="${SELECTION_PREFIX:-${RUN_ROOT}/selection/selected_regions}" +BUNDLE_DIR="${BUNDLE_DIR:-${RUN_ROOT}/data/rbpnet}" +CV_PLAN="${CV_PLAN:-${RUN_ROOT}/cv/cv5_chromosomes.json}" +CV_ROOT="${CV_ROOT:-${RUN_ROOT}/cv}" + +COORDINATE_SPACE="${COORDINATE_SPACE:-mature_transcript}" +READ1_RNA_STRAND="${READ1_RNA_STRAND:-opposite}" +MIN_MAPQ="${MIN_MAPQ:-1}" +SIGNAL_COMPRESSION="${SIGNAL_COMPRESSION:-gzip}" +SIGNAL_COMPRESSION_LEVEL="${SIGNAL_COMPRESSION_LEVEL:-1}" +OVERWRITE="${OVERWRITE:-0}" + +# Leave WINDOW_STRIDE empty to use 1 for original_rbpnet and 50 otherwise. +WINDOW_SIZE="${WINDOW_SIZE:-100}" +WINDOW_STRIDE="${WINDOW_STRIDE:-}" +MIN_SMINPUT_TPM="${MIN_SMINPUT_TPM:-0}" +SELECTION_STRATEGY="${SELECTION_STRATEGY:-peak_gray_negative}" +MIN_TOTAL_COUNT="${MIN_TOTAL_COUNT:-8}" +MIN_IP_COUNT="${MIN_IP_COUNT:-0}" +MIN_SMINPUT_COUNT="${MIN_SMINPUT_COUNT:-0}" +PEAK_FDR="${PEAK_FDR:-0.05}" +PEAK_MIN_LOG2_RATIO="${PEAK_MIN_LOG2_RATIO:-1.0}" +NEGATIVE_FDR="${NEGATIVE_FDR:-0.05}" +NEGATIVE_MAX_LOG2_RATIO="${NEGATIVE_MAX_LOG2_RATIO:--0.5}" +STITCH_GAP="${STITCH_GAP:-0}" +POISSON_NULL="${POISSON_NULL:-ip_locus_density}" +REPLICATE_MODE="${REPLICATE_MODE:-per_ip}" + +INPUT_LENGTH="${INPUT_LENGTH:-300}" +PROFILE_LENGTH="${PROFILE_LENGTH:-300}" +MAX_JITTER="${MAX_JITTER:-32}" +TRANSCRIPT_END_POLICY="${TRANSCRIPT_END_POLICY:-shift_to_fit}" + +N_FOLDS="${N_FOLDS:-5}" +CHROMOSOME_GROUP_COL="${CHROMOSOME_GROUP_COL:-group_chromosome}" +BASE_TRAIN_CONFIG="${BASE_TRAIN_CONFIG:-${SCRIPT_CONFIG_DIR}/example_train_config.json}" +DEVICE="${DEVICE:-cuda}" + +setup_transcriptml_env() { + module load gcc/10.1.0 + module load openblas/0.3.10 + source "${SHERLOCK_CONDA_ROOT}/etc/profile.d/conda.sh" + conda activate "${CONDA_ENV}" + if [[ -n "${TRANSCRIPTML_REPO}" ]]; then + if [[ ! -d "${TRANSCRIPTML_REPO}/src/transcriptml" ]]; then + echo "TRANSCRIPTML_REPO is not a TranscriptML checkout: ${TRANSCRIPTML_REPO}" >&2 + return 1 + fi + cd "${TRANSCRIPTML_REPO}" + export PYTHONPATH="${TRANSCRIPTML_REPO}/src:${PYTHONPATH:-}" + elif ! command -v transcriptml >/dev/null 2>&1; then + echo "Set TRANSCRIPTML_REPO or install TranscriptML in ${CONDA_ENV}." >&2 + return 1 + fi +} diff --git a/scripts/rbpnet/scan_select_bundle.sh b/scripts/rbpnet/scan_select_bundle.sh new file mode 100644 index 0000000..de2e4d0 --- /dev/null +++ b/scripts/rbpnet/scan_select_bundle.sh @@ -0,0 +1,97 @@ +#!/bin/bash +#SBATCH --partition=akundaje +#SBATCH --job-name=tml_rbp_bundle +#SBATCH --cpus-per-task=4 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm_output/%x_%j.out +#SBATCH --error=slurm_output/%x_%j.err + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/rbpnet/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts/rbpnet" + elif [[ -f "${SLURM_SUBMIT_DIR}/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/rbpnet_config.sh" +setup_transcriptml_env + +if [[ ! -f "${PROCESSED_DIR}/signals.h5" ]]; then + echo "Missing ${PROCESSED_DIR}/signals.h5; run preprocess.sh first." >&2 + exit 1 +fi + +stride="${WINDOW_STRIDE}" +if [[ -z "${stride}" ]]; then + if [[ "${SELECTION_STRATEGY}" == "original_rbpnet" ]]; then + stride=1 + else + stride=50 + fi +fi + +overwrite_args=() +if [[ "${OVERWRITE}" == "1" ]]; then + overwrite_args+=(--overwrite) +fi + +transcriptml rbpnet scan-windows \ + --processed-dir "${PROCESSED_DIR}" \ + --window-size "${WINDOW_SIZE}" \ + --stride "${stride}" \ + --min-sminput-tpm "${MIN_SMINPUT_TPM}" \ + --output-prefix "${WINDOW_PREFIX}" \ + "${overwrite_args[@]}" + +selection_args=( + --processed-dir "${PROCESSED_DIR}" + --windows "${WINDOW_PREFIX}.parquet" + --strategy "${SELECTION_STRATEGY}" + --min-sminput-tpm "${MIN_SMINPUT_TPM}" + --output-prefix "${SELECTION_PREFIX}" +) +case "${SELECTION_STRATEGY}" in + original_rbpnet) + selection_args+=(--poisson-null "${POISSON_NULL}") + ;; + broad_coverage) + selection_args+=( + --min-total-count "${MIN_TOTAL_COUNT}" + --min-ip-count "${MIN_IP_COUNT}" + --min-sminput-count "${MIN_SMINPUT_COUNT}" + --replicate-mode "${REPLICATE_MODE}" + ) + ;; + peak_gray_negative) + selection_args+=( + --min-total-count "${MIN_TOTAL_COUNT}" + --min-ip-count "${MIN_IP_COUNT}" + --min-sminput-count "${MIN_SMINPUT_COUNT}" + --peak-fdr "${PEAK_FDR}" + --peak-min-log2-ratio "${PEAK_MIN_LOG2_RATIO}" + --negative-fdr "${NEGATIVE_FDR}" + --negative-max-log2-ratio "${NEGATIVE_MAX_LOG2_RATIO}" + --stitch-gap "${STITCH_GAP}" + ) + ;; + *) + echo "Unknown SELECTION_STRATEGY: ${SELECTION_STRATEGY}" >&2 + exit 1 + ;; +esac +selection_args+=("${overwrite_args[@]}") +transcriptml rbpnet select-regions "${selection_args[@]}" + +transcriptml rbpnet make-bundle \ + --processed-dir "${PROCESSED_DIR}" \ + --selection-manifest "${SELECTION_PREFIX}.parquet" \ + --output-dir "${BUNDLE_DIR}" \ + --input-length "${INPUT_LENGTH}" \ + --profile-length "${PROFILE_LENGTH}" \ + --max-jitter "${MAX_JITTER}" \ + --transcript-end-policy "${TRANSCRIPT_END_POLICY}" \ + "${overwrite_args[@]}" diff --git a/scripts/rbpnet/submit_train_cv.sh b/scripts/rbpnet/submit_train_cv.sh new file mode 100644 index 0000000..8844991 --- /dev/null +++ b/scripts/rbpnet/submit_train_cv.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/rbpnet/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts/rbpnet" + elif [[ -f "${SLURM_SUBMIT_DIR}/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/rbpnet_config.sh" + +if [[ ! -f "${CV_PLAN}" ]]; then + echo "Missing chromosome CV plan ${CV_PLAN}; run create_chromosome_cv_plan.sh first." >&2 + exit 1 +fi + +mkdir -p "${CV_ROOT}" slurm_output +sbatch --array="0-$((N_FOLDS - 1))" "${SCRIPT_DIR}/train_cv_fold.sh" diff --git a/scripts/rbpnet/train_cv_fold.sh b/scripts/rbpnet/train_cv_fold.sh new file mode 100644 index 0000000..74a49d9 --- /dev/null +++ b/scripts/rbpnet/train_cv_fold.sh @@ -0,0 +1,48 @@ +#!/bin/bash +#SBATCH --partition=akundaje +#SBATCH --job-name=tml_rbp_cv +#SBATCH --cpus-per-task=4 +#SBATCH --gpus=1 +#SBATCH --mem=48G +#SBATCH --time=24:00:00 +#SBATCH -C GPU_MEM:48GB +#SBATCH --output=slurm_output/%x_%A_%a.out +#SBATCH --error=slurm_output/%x_%A_%a.err + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/rbpnet/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts/rbpnet" + elif [[ -f "${SLURM_SUBMIT_DIR}/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/rbpnet_config.sh" +setup_transcriptml_env + +FOLD="${SLURM_ARRAY_TASK_ID:-${1:-}}" +if [[ -z "${FOLD}" ]]; then + echo "Provide a fold argument or run as a Slurm array job." >&2 + exit 1 +fi +if [[ ! "${FOLD}" =~ ^[0-9]+$ || "${FOLD}" -ge "${N_FOLDS}" ]]; then + echo "Fold must be an integer in [0, $((N_FOLDS - 1))]; got ${FOLD}." >&2 + exit 1 +fi +if [[ ! -f "${CV_PLAN}" ]]; then + echo "Missing chromosome CV plan ${CV_PLAN}; run create_chromosome_cv_plan.sh first." >&2 + exit 1 +fi +if [[ ! -f "${BUNDLE_DIR}/metadata.json" ]]; then + echo "Missing RBPNet bundle metadata at ${BUNDLE_DIR}." >&2 + exit 1 +fi + +FOLD_DIR="${CV_ROOT}/fold${FOLD}" +transcriptml train "${BASE_TRAIN_CONFIG}" \ + --cv-plan "${CV_PLAN}" \ + --fold "${FOLD}" \ + --dataset "${BUNDLE_DIR}" \ + --output-dir "${FOLD_DIR}/model" diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index 6f3aa2f..cdbde26 100644 --- a/src/transcriptml/cli/main.py +++ b/src/transcriptml/cli/main.py @@ -118,6 +118,21 @@ def build_parser() -> argparse.ArgumentParser: p_fold.add_argument("--n-folds", type=int, default=10) p_fold.add_argument("--seed", type=int, default=42) p_fold.add_argument("--val-offset", type=int, default=1) + p_plan = cv_sub.add_parser( + "create-chromosome-plan", + help="Balance complete chromosomes into an immutable N-fold CV plan", + ) + p_plan.add_argument("--dataset", required=True, help="DatasetBundle directory") + p_plan.add_argument("--output", required=True, help="Output chromosome-plan JSON") + p_plan.add_argument("--n-folds", type=int, required=True) + p_plan.add_argument("--group-col", default="group_chromosome") + p_resolve = cv_sub.add_parser( + "resolve-plan", help="Resolve one saved chromosome CV run to split indices" + ) + p_resolve.add_argument("--dataset", required=True, help="DatasetBundle directory") + p_resolve.add_argument("--cv-plan", required=True, help="Saved chromosome-plan JSON") + p_resolve.add_argument("--fold", type=int, required=True) + p_resolve.add_argument("--output", help="Optional output splits JSON") p_ensemble = cv_sub.add_parser( "ensemble-predict", help="Average predictions from fold checkpoints on one shared dataset", @@ -176,6 +191,10 @@ def build_parser() -> argparse.ArgumentParser: p = sub.add_parser("train", help="Train from a JSON/TOML config") p.add_argument("config") + p.add_argument("--cv-plan", help="Immutable chromosome CV plan JSON") + p.add_argument("--fold", type=int, help="Zero-based CV test fold") + p.add_argument("--dataset", help="Override config dataset path") + p.add_argument("--output-dir", help="Override config output_dir (useful for job arrays)") p = sub.add_parser("evaluate", help="Evaluate a checkpoint on a dataset bundle") p.add_argument("checkpoint", nargs="?", metavar="CHECKPOINT", help="Checkpoint path; prefer --checkpoint") @@ -406,6 +425,49 @@ def main(argv: list[str] | None = None) -> None: ) print(config_path) return + if args.cv_command in {"create-chromosome-plan", "resolve-plan"}: + from transcriptml.data.bundle import load_bundle + from transcriptml.workflows import ( + create_chromosome_cv_plan, + load_chromosome_cv_plan, + resolve_chromosome_cv_plan, + save_chromosome_cv_plan, + ) + + bundle = load_bundle(args.dataset, mmap_mode="r") + if bundle.metadata is None: + raise SystemExit("Dataset bundle has no metadata for chromosome CV") + if args.cv_command == "create-chromosome-plan": + plan = create_chromosome_cv_plan( + bundle.metadata, + n_folds=args.n_folds, + group_col=args.group_col, + ) + output = save_chromosome_cv_plan(plan, args.output) + print(output) + return + plan = load_chromosome_cv_plan(args.cv_plan) + resolution = resolve_chromosome_cv_plan( + plan, bundle.metadata, fold=args.fold + ) + result = { + "plan_id": plan.plan_id, + "fold": resolution.fold, + "validation_fold": resolution.validation_fold, + "chromosomes": { + name: list(values) for name, values in resolution.groups.items() + }, + "indices": resolution.indices, + } + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + Path(args.output).write_text( + json.dumps(result, indent=2) + "\n", encoding="utf-8" + ) + print(args.output) + else: + print(json.dumps(result, indent=2)) + return if args.cv_command == "ensemble-predict": from transcriptml.progress import log_progress from transcriptml.training.evaluation import evaluate_fold_checkpoints @@ -537,7 +599,13 @@ def main(argv: list[str] | None = None) -> None: if args.command == "train": from transcriptml.training.trainer import train_from_config - train_from_config(args.config) + train_from_config( + args.config, + cv_plan=args.cv_plan, + fold=args.fold, + dataset=args.dataset, + output_dir=args.output_dir, + ) return if args.command == "evaluate": from transcriptml.progress import log_progress diff --git a/src/transcriptml/rbpnet/cli.py b/src/transcriptml/rbpnet/cli.py index d980df1..fd6f4eb 100644 --- a/src/transcriptml/rbpnet/cli.py +++ b/src/transcriptml/rbpnet/cli.py @@ -41,6 +41,18 @@ def add_rbpnet_parser(subparsers) -> None: ) preprocess.add_argument("--min-mapq", type=int, default=1) preprocess.add_argument("--include-duplicates", action="store_true") + preprocess.add_argument( + "--signal-compression", + choices=("gzip", "lzf", "none"), + default="gzip", + help="HDF5 signal compression (default: gzip)", + ) + preprocess.add_argument( + "--signal-compression-level", + type=int, + default=None, + help="gzip level 0-9 (default: 1; invalid for lzf/none)", + ) preprocess.add_argument("--overwrite", action="store_true") preprocess.add_argument("--no-progress", action="store_true") @@ -157,6 +169,9 @@ def run_rbpnet_command(args: argparse.Namespace, parser: argparse.ArgumentParser if args.rbpnet_command == "preprocess": from transcriptml.rbpnet.preprocessing import PipelineConfig, preprocess_eclip + compression_level = args.signal_compression_level + if compression_level is None and args.signal_compression == "gzip": + compression_level = 1 qc = preprocess_eclip(PipelineConfig( genome_fasta=args.genome_fasta, gtf=args.gtf, @@ -167,6 +182,8 @@ def run_rbpnet_command(args: argparse.Namespace, parser: argparse.ArgumentParser read1_rna_strand=args.read1_rna_strand, min_mapq=args.min_mapq, exclude_duplicates=not args.include_duplicates, + signal_compression=args.signal_compression, + signal_compression_level=compression_level, overwrite=args.overwrite, progress=not args.no_progress, )) diff --git a/src/transcriptml/rbpnet/preprocessing.py b/src/transcriptml/rbpnet/preprocessing.py index d437f8c..f7af823 100644 --- a/src/transcriptml/rbpnet/preprocessing.py +++ b/src/transcriptml/rbpnet/preprocessing.py @@ -46,6 +46,8 @@ class PipelineConfig: read1_rna_strand: str = "opposite" min_mapq: int = 1 exclude_duplicates: bool = True + signal_compression: str | None = "gzip" + signal_compression_level: int | None = 1 overwrite: bool = False progress: bool = True @@ -136,6 +138,8 @@ def preprocess_eclip(config: PipelineConfig) -> dict: transcripts, sample_names, [sample.role for sample in samples], + compression=config.signal_compression, + compression_level=config.signal_compression_level, ) as store: for row, sample in enumerate(samples): log_progress( @@ -195,6 +199,8 @@ def preprocess_eclip(config: PipelineConfig) -> dict: "read1_rna_strand": config.read1_rna_strand, "min_mapq": config.min_mapq, "exclude_duplicates": config.exclude_duplicates, + "signal_compression": config.signal_compression, + "signal_compression_level": config.signal_compression_level, "assignment": assignment, }, "annotation": { diff --git a/src/transcriptml/rbpnet/signals.py b/src/transcriptml/rbpnet/signals.py index 8564c7a..02f4b14 100644 --- a/src/transcriptml/rbpnet/signals.py +++ b/src/transcriptml/rbpnet/signals.py @@ -6,6 +6,7 @@ import tempfile from collections import Counter from pathlib import Path +from typing import Iterable, Iterator import h5py import numpy as np @@ -210,6 +211,10 @@ def create_signal_store( transcripts: list[Transcript], sample_names: list[str], sample_roles: list[str], + *, + compression: str | None = "gzip", + compression_level: int | None = 1, + chunk_length: int = 1_048_576, ) -> h5py.File: """Create the canonical concatenated locus-coordinate HDF5 store.""" @@ -237,20 +242,155 @@ def create_signal_store( store.create_dataset("transcript_lengths", data=np.asarray([tx.length for tx in transcripts], dtype=np.int64)) store.create_dataset("sample_names", data=np.asarray(sample_names, dtype=object), dtype=strings) store.create_dataset("sample_roles", data=np.asarray(sample_roles, dtype=object), dtype=strings) - chunk = min(total_length, 1_048_576) + compression, compression_opts = _normalize_compression( + compression, compression_level + ) + chunk = min(total_length, int(chunk_length)) + if chunk <= 0: + raise ValueError("chunk_length must be positive") + store.attrs["signal_compression"] = "none" if compression is None else compression + store.attrs["signal_compression_level"] = ( + -1 if compression_opts is None else int(compression_opts) + ) + store.attrs["signal_chunk_length"] = chunk store.create_dataset( "counts", shape=(len(sample_names), total_length), dtype=np.uint32, chunks=(1, chunk), - compression="gzip", - compression_opts=4, + compression=compression, + compression_opts=compression_opts, shuffle=True, fillvalue=0, ) return store +def _normalize_compression( + compression: str | None, + compression_level: int | None, +) -> tuple[str | None, int | None]: + """Validate signal compression and return h5py keyword values.""" + + if compression is None or str(compression).strip().lower() in { + "none", "off", "uncompressed", + }: + if compression_level not in {None, 0}: + raise ValueError("compression_level is only valid with gzip") + return None, None + name = str(compression).strip().lower() + if name == "gzip": + level = 1 if compression_level is None else int(compression_level) + if level < 0 or level > 9: + raise ValueError("gzip compression_level must be between 0 and 9") + return name, level + if name == "lzf": + if compression_level not in {None, 0}: + raise ValueError("lzf does not accept a compression level") + return name, None + raise ValueError("signal compression must be one of: gzip, lzf, none") + + +def _sqlite_sparse_batches( + cursor: sqlite3.Cursor, + *, + batch_size: int = 100_000, +) -> Iterator[tuple[np.ndarray, np.ndarray]]: + """Yield sorted sparse position/count arrays from the aggregation table.""" + + while True: + rows = cursor.fetchmany(int(batch_size)) + if not rows: + return + positions = np.fromiter( + (item[0] for item in rows), dtype=np.int64, count=len(rows) + ) + values = np.fromiter( + (item[1] for item in rows), dtype=np.uint64, count=len(rows) + ) + yield positions, values + + +def write_sparse_signal_chunkwise( + dataset: h5py.Dataset, + row: int, + batches: Iterable[tuple[np.ndarray, np.ndarray]], + *, + reporter: ProgressReporter | None = None, +) -> int: + """Write a sorted sparse signal row using one dense write per HDF5 chunk. + + The input positions must be unique and strictly increasing across batches. + Only chunks containing at least one nonzero count are materialized. A + dense ``uint32`` buffer is retained when a sparse input batch ends partway + through a chunk, ensuring that even such chunks are written exactly once. + """ + + if dataset.ndim != 2 or dataset.chunks is None: + raise ValueError("signal dataset must be a two-dimensional chunked dataset") + if row < 0 or row >= dataset.shape[0]: + raise IndexError("signal dataset row is outside bounds") + row_chunk, chunk_length = (int(value) for value in dataset.chunks) + if row_chunk != 1: + raise ValueError("signal dataset must use one sample row per HDF5 chunk") + total_length = int(dataset.shape[1]) + uint32_max = np.iinfo(np.uint32).max + active_chunk = -1 + active_start = 0 + active_buffer: np.ndarray | None = None + previous_position = -1 + written_positions = 0 + + def flush() -> None: + nonlocal active_buffer + if active_buffer is None: + return + dataset[row, active_start : active_start + active_buffer.size] = active_buffer + active_buffer = None + + for raw_positions, raw_values in batches: + positions = np.asarray(raw_positions, dtype=np.int64) + values = np.asarray(raw_values, dtype=np.uint64) + if positions.ndim != 1 or values.ndim != 1 or positions.shape != values.shape: + raise ValueError("sparse signal positions and values must be aligned vectors") + if positions.size == 0: + continue + if positions[0] <= previous_position or np.any(np.diff(positions) <= 0): + raise ValueError("sparse signal positions must be unique and strictly increasing") + if positions[0] < 0 or positions[-1] >= total_length: + raise IndexError("sparse signal position is outside dataset bounds") + if np.any(values == 0): + raise ValueError("sparse signal values must be positive") + if values.max(initial=0) > uint32_max: + raise OverflowError("a crosslink-position count exceeds uint32") + + offset = 0 + while offset < positions.size: + chunk_index = int(positions[offset] // chunk_length) + chunk_end_position = min((chunk_index + 1) * chunk_length, total_length) + end = int(np.searchsorted(positions, chunk_end_position, side="left")) + if chunk_index != active_chunk: + flush() + active_chunk = chunk_index + active_start = chunk_index * chunk_length + active_buffer = np.zeros( + chunk_end_position - active_start, dtype=np.uint32 + ) + assert active_buffer is not None + local_positions = positions[offset:end] - active_start + active_buffer[local_positions] = values[offset:end].astype( + np.uint32, copy=False + ) + written = end - offset + written_positions += written + if reporter is not None: + reporter.update(written) + offset = end + previous_position = int(positions[-1]) + flush() + return written_positions + + def _flush_counts(connection: sqlite3.Connection, counts: Counter[int]) -> None: if not counts: return @@ -358,17 +498,17 @@ def extract_bam_to_store( unit="positions", enabled=progress, ) - while True: - rows = cursor.fetchmany(100_000) - if not rows: - break - unique_positions += len(rows) - positions = np.fromiter((item[0] for item in rows), dtype=np.int64, count=len(rows)) - values64 = np.fromiter((item[1] for item in rows), dtype=np.uint64, count=len(rows)) - if values64.max(initial=0) > np.iinfo(np.uint32).max: - raise OverflowError(f"a crosslink-position count in {bam_path} exceeds uint32") - dataset[row, positions] = values64.astype(np.uint32) - reporter.update(len(rows)) + try: + unique_positions = write_sparse_signal_chunkwise( + dataset, + row, + _sqlite_sparse_batches(cursor), + reporter=reporter, + ) + except OverflowError as exc: + raise OverflowError( + f"a crosslink-position count in {bam_path} exceeds uint32" + ) from exc reporter.close() connection.close() qc["unique_crosslink_positions"] = unique_positions @@ -389,13 +529,15 @@ def write_ip_pooled(store: h5py.File, ip_rows: list[int], *, progress: bool = Tr counts = store["counts"] total = counts.shape[1] chunk = counts.chunks[1] + compression = counts.compression + compression_opts = counts.compression_opts if compression == "gzip" else None pooled = store.create_dataset( "ip_pooled", shape=(total,), dtype=np.uint32, chunks=(chunk,), - compression="gzip", - compression_opts=4, + compression=compression, + compression_opts=compression_opts, shuffle=True, fillvalue=0, ) @@ -411,5 +553,6 @@ def write_ip_pooled(store: h5py.File, ip_rows: list[int], *, progress: bool = Tr values = counts[ip_rows, start:end].astype(np.uint64).sum(axis=0) if values.max(initial=0) > np.iinfo(np.uint32).max: raise OverflowError("pooled IP count exceeds uint32") - pooled[start:end] = values.astype(np.uint32) + if np.any(values): + pooled[start:end] = values.astype(np.uint32) pooled.attrs["source_rows"] = np.asarray(ip_rows, dtype=np.int64) diff --git a/src/transcriptml/rbpnet/training.py b/src/transcriptml/rbpnet/training.py index ad3da80..7a62e78 100644 --- a/src/transcriptml/rbpnet/training.py +++ b/src/transcriptml/rbpnet/training.py @@ -160,6 +160,26 @@ def _select_rbpnet_splits( bundle: DatasetBundle, cfg, ) -> tuple[dict[str, list[int]], str, str, bool]: + if getattr(cfg, "cv_plan", None) is not None or getattr(cfg, "fold", None) is not None: + if getattr(cfg, "cv_plan", None) is None or getattr(cfg, "fold", None) is None: + raise ValueError("cv_plan and fold must be provided together") + if bundle.metadata is None: + raise ValueError("chromosome CV requires RBPNet bundle metadata") + from transcriptml.workflows.chromosome_cv import ( + load_chromosome_cv_plan, + resolve_chromosome_cv_plan, + ) + + plan = load_chromosome_cv_plan(cfg.cv_plan) + resolution = resolve_chromosome_cv_plan( + plan, bundle.metadata, fold=cfg.fold + ) + return ( + normalize_splits(resolution.indices), + "cv_plan", + plan.group_col, + False, + ) source = str(cfg.split_source or "auto").strip().lower() if source not in {"auto", "bundle", "config"}: raise ValueError("split_source must be one of: auto, bundle, config") @@ -180,6 +200,16 @@ def _select_rbpnet_splits( return splits, source_used, group_col, allow_random +def _cv_plan_id(cfg) -> str | None: + """Return the validated plan identifier recorded by CV training artifacts.""" + + if getattr(cfg, "cv_plan", None) is None: + return None + from transcriptml.workflows.chromosome_cv import load_chromosome_cv_plan + + return load_chromosome_cv_plan(cfg.cv_plan).plan_id + + def _deduplicate_splits( bundle: DatasetBundle, splits: Mapping[str, Sequence[int]], @@ -433,6 +463,7 @@ def train_rbpnet_model(bundle: DatasetBundle, cfg) -> dict[str, Any]: enrichment_enabled=model.enrichment_enabled, ).to(device) splits, split_source, group_col, random_split = _select_rbpnet_splits(bundle, cfg) + cv_plan_id = _cv_plan_id(cfg) splits, n_deduplicated = _deduplicate_splits( bundle, splits, @@ -544,6 +575,7 @@ def train_rbpnet_model(bundle: DatasetBundle, cfg) -> dict[str, Any]: checkpoint_extra = { "splits": splits, "split_source_used": split_source, + "cv_plan_id": cv_plan_id, "train_config": asdict(cfg), "loss_config": loss_config.to_dict( enrichment_enabled=model.enrichment_enabled @@ -636,6 +668,9 @@ def train_rbpnet_model(bundle: DatasetBundle, cfg) -> dict[str, Any]: "mixed_precision": mixed_precision, "max_train_jitter": train_dataset.max_train_jitter, "split_source_used": split_source, + "cv_plan": cfg.cv_plan, + "cv_plan_id": cv_plan_id, + "fold": cfg.fold, "split_group_col": group_col, "unsafe_random_window_split": random_split, "split_counts": split_counts, diff --git a/src/transcriptml/training/trainer.py b/src/transcriptml/training/trainer.py index c9506b0..faf2745 100644 --- a/src/transcriptml/training/trainer.py +++ b/src/transcriptml/training/trainer.py @@ -49,6 +49,8 @@ class TrainConfig: head_layernorm: bool = False sequence_controls: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None split_source: str = "auto" + cv_plan: str | None = None + fold: int | None = None max_train_jitter: int = 0 allow_random_window_split: bool = False deduplicate_loci: bool = True @@ -135,6 +137,21 @@ def _select_splits(bundle: DatasetBundle, cfg: TrainConfig) -> tuple[dict[str, l cfg: Training configuration containing split source and strategy. """ + if cfg.cv_plan is not None or cfg.fold is not None: + if cfg.cv_plan is None or cfg.fold is None: + raise ValueError("cv_plan and fold must be provided together") + if bundle.metadata is None: + raise ValueError("chromosome CV requires dataset bundle metadata") + from transcriptml.workflows.chromosome_cv import ( + load_chromosome_cv_plan, + resolve_chromosome_cv_plan, + ) + + plan = load_chromosome_cv_plan(cfg.cv_plan) + resolution = resolve_chromosome_cv_plan( + plan, bundle.metadata, fold=cfg.fold + ) + return normalize_splits(resolution.indices), "cv_plan" source = str(cfg.split_source or "auto").strip().lower() if source not in {"auto", "bundle", "config"}: raise ValueError("split_source must be one of: auto, bundle, config") @@ -157,6 +174,16 @@ def _make_splits(bundle: DatasetBundle, cfg: TrainConfig) -> dict[str, list[int] return splits +def _cv_plan_id(cfg: TrainConfig) -> str | None: + """Return the validated plan identifier recorded by CV training artifacts.""" + + if cfg.cv_plan is None: + return None + from transcriptml.workflows.chromosome_cv import load_chromosome_cv_plan + + return load_chromosome_cv_plan(cfg.cv_plan).plan_id + + class _ArrayRegressionDataset(Dataset): def __init__(self, X: np.ndarray, y: np.ndarray, aux_arrays: Mapping[str, np.ndarray] | None = None): """Wrap NumPy arrays as a PyTorch regression dataset. @@ -546,6 +573,7 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) else: y_train = bundle.y splits, split_source_used = _select_splits(bundle, cfg) + cv_plan_id = _cv_plan_id(cfg) split_counts = {name: len(splits.get(name, [])) for name in ("train", "val", "test")} model_config = normalize_model_config(cfg.model) if cfg.head_layernorm and model_config.name != "saluki_exact": @@ -699,6 +727,7 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) extra={ "splits": splits, "split_source_used": split_source_used, + "cv_plan_id": cv_plan_id, "train_config": asdict(cfg), "loss_config": normalized_loss_config, }, @@ -715,6 +744,7 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) extra={ "splits": splits, "split_source_used": split_source_used, + "cv_plan_id": cv_plan_id, "train_config": asdict(cfg), "loss_config": normalized_loss_config, }, @@ -781,6 +811,9 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) "head_layernorm": bool(cfg.head_layernorm), "split_source_requested": cfg.split_source, "split_source_used": split_source_used, + "cv_plan": cfg.cv_plan, + "cv_plan_id": cv_plan_id, + "fold": cfg.fold, "split_counts": split_counts, "test_loss": test_loss_metrics.get("loss"), "test_mse": test_result.get("loss"), @@ -800,15 +833,35 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) return {"model": model, "history": history, "splits": splits, "summary": summary} -def train_from_config(config_path: str | Path, *, progress: bool | None = None) -> dict[str, Any]: +def train_from_config( + config_path: str | Path, + *, + progress: bool | None = None, + cv_plan: str | Path | None = None, + fold: int | None = None, + dataset: str | Path | None = None, + output_dir: str | Path | None = None, +) -> dict[str, Any]: """Load a training config and train its requested model. Args: config_path: Path to a JSON or TOML training configuration file. progress: Optional override for whether progress messages are emitted. + cv_plan: Optional chromosome CV plan overriding the config. + fold: Optional zero-based CV test fold overriding the config. + dataset: Optional dataset-directory override. + output_dir: Optional output-directory override. """ cfg = TrainConfig(**_load_config(config_path)) + if cv_plan is not None: + cfg.cv_plan = str(cv_plan) + if fold is not None: + cfg.fold = int(fold) + if dataset is not None: + cfg.dataset = str(dataset) + if output_dir is not None: + cfg.output_dir = str(output_dir) if progress is not None: cfg.progress = bool(progress) log_progress(f"training: loading dataset {cfg.dataset}", enabled=cfg.progress) diff --git a/src/transcriptml/workflows/__init__.py b/src/transcriptml/workflows/__init__.py index f232762..731d5c3 100644 --- a/src/transcriptml/workflows/__init__.py +++ b/src/transcriptml/workflows/__init__.py @@ -1,6 +1,25 @@ """Workflow template helpers for TranscriptML.""" +from transcriptml.workflows.chromosome_cv import ( + ChromosomeCVPlan, + ChromosomeCVResolution, + create_chromosome_cv_plan, + load_chromosome_cv_plan, + resolve_chromosome_cv_plan, + save_chromosome_cv_plan, +) from transcriptml.workflows.cv import find_fold_checkpoints, load_fold_test_indices, prepare_cv_fold from transcriptml.workflows.init_run import init_run -__all__ = ["find_fold_checkpoints", "init_run", "load_fold_test_indices", "prepare_cv_fold"] +__all__ = [ + "ChromosomeCVPlan", + "ChromosomeCVResolution", + "create_chromosome_cv_plan", + "find_fold_checkpoints", + "init_run", + "load_chromosome_cv_plan", + "load_fold_test_indices", + "prepare_cv_fold", + "resolve_chromosome_cv_plan", + "save_chromosome_cv_plan", +] diff --git a/src/transcriptml/workflows/chromosome_cv.py b/src/transcriptml/workflows/chromosome_cv.py new file mode 100644 index 0000000..17745e4 --- /dev/null +++ b/src/transcriptml/workflows/chromosome_cv.py @@ -0,0 +1,297 @@ +"""Immutable, example-balanced chromosome cross-validation plans.""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Mapping, Sequence + +from transcriptml import __version__ + + +PLAN_FORMAT = "transcriptml-chromosome-cv-plan" +PLAN_FORMAT_VERSION = "1" +PLAN_ALGORITHM = "largest_chromosome_first_greedy" + + +def _natural_key(value: str) -> tuple[tuple[int, object], ...]: + return tuple( + (0, int(part)) if part.isdigit() else (1, part.lower()) + for part in re.split(r"(\d+)", str(value)) + if part + ) + + +def _plan_digest(payload: Mapping[str, object]) -> str: + encoded = json.dumps( + payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +@dataclass(frozen=True) +class ChromosomeCVPlan: + """A versioned assignment of complete chromosomes to fold groups.""" + + n_folds: int + group_col: str + n_examples: int + chromosome_counts: Mapping[str, int] + fold_groups: tuple[tuple[str, ...], ...] + fold_example_counts: tuple[int, ...] + plan_id: str + algorithm: str = PLAN_ALGORITHM + transcriptml_version: str = __version__ + + def _content_dict(self) -> dict[str, object]: + return { + "format": PLAN_FORMAT, + "format_version": PLAN_FORMAT_VERSION, + "n_folds": int(self.n_folds), + "group_col": self.group_col, + "n_examples": int(self.n_examples), + "chromosome_counts": { + chrom: int(self.chromosome_counts[chrom]) + for chrom in sorted(self.chromosome_counts, key=_natural_key) + }, + "fold_groups": [ + { + "fold": fold, + "chromosomes": list(chromosomes), + "example_count": int(self.fold_example_counts[fold]), + } + for fold, chromosomes in enumerate(self.fold_groups) + ], + "generation": { + "algorithm": self.algorithm, + "group_order": "descending example count, then natural chromosome name", + "fold_tie_break": "lowest current example count, then lowest fold index", + "transcriptml_version": self.transcriptml_version, + }, + } + + def to_dict(self) -> dict[str, object]: + """Serialize the complete, self-validating plan.""" + + payload = self._content_dict() + payload["plan_id"] = self.plan_id + return payload + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> "ChromosomeCVPlan": + """Validate and reconstruct a plan from JSON-like data.""" + + if value.get("format") != PLAN_FORMAT: + raise ValueError("not a TranscriptML chromosome CV plan") + if str(value.get("format_version")) != PLAN_FORMAT_VERSION: + raise ValueError("unsupported chromosome CV plan format version") + raw_folds = value.get("fold_groups") + if not isinstance(raw_folds, list): + raise ValueError("chromosome CV plan fold_groups must be a list") + folds: list[tuple[str, ...]] = [] + fold_counts: list[int] = [] + for expected_fold, raw in enumerate(raw_folds): + if not isinstance(raw, Mapping) or int(raw.get("fold", -1)) != expected_fold: + raise ValueError("chromosome CV plan folds must be consecutively indexed") + chromosomes = raw.get("chromosomes") + if not isinstance(chromosomes, list): + raise ValueError("each chromosome CV fold must list chromosomes") + folds.append(tuple(str(chrom) for chrom in chromosomes)) + fold_counts.append(int(raw.get("example_count", -1))) + raw_counts = value.get("chromosome_counts") + if not isinstance(raw_counts, Mapping): + raise ValueError("chromosome CV plan lacks chromosome_counts") + counts = {str(chrom): int(count) for chrom, count in raw_counts.items()} + plan = cls( + n_folds=int(value.get("n_folds", 0)), + group_col=str(value.get("group_col", "")), + n_examples=int(value.get("n_examples", -1)), + chromosome_counts=counts, + fold_groups=tuple(folds), + fold_example_counts=tuple(fold_counts), + plan_id=str(value.get("plan_id", "")), + algorithm=str( + value.get("generation", {}).get("algorithm", PLAN_ALGORITHM) + if isinstance(value.get("generation"), Mapping) + else PLAN_ALGORITHM + ), + transcriptml_version=str( + value.get("generation", {}).get("transcriptml_version", "unknown") + if isinstance(value.get("generation"), Mapping) + else "unknown" + ), + ) + _validate_plan(plan) + expected_id = _plan_digest(plan._content_dict()) + if plan.plan_id != expected_id: + raise ValueError("chromosome CV plan_id does not match its contents") + return plan + + +@dataclass(frozen=True) +class ChromosomeCVResolution: + """Train/validation/test groups and row indices for one CV run.""" + + fold: int + validation_fold: int + groups: Mapping[str, tuple[str, ...]] + indices: Mapping[str, list[int]] + + +def _count_chromosomes( + metadata: Sequence[Mapping[str, object]], group_col: str +) -> dict[str, int]: + if not metadata: + raise ValueError("metadata must contain at least one example") + counts: dict[str, int] = {} + for index, row in enumerate(metadata): + value = row.get(group_col) + if value is None or not str(value).strip(): + raise ValueError( + f"metadata row {index} lacks chromosome grouping column {group_col!r}" + ) + chromosome = str(value) + counts[chromosome] = counts.get(chromosome, 0) + 1 + return counts + + +def _validate_plan(plan: ChromosomeCVPlan) -> None: + if plan.n_folds < 3: + raise ValueError("chromosome CV plans require at least three folds") + if not plan.group_col: + raise ValueError("chromosome CV group_col must be non-empty") + if len(plan.fold_groups) != plan.n_folds: + raise ValueError("chromosome CV fold count disagrees with n_folds") + if len(plan.fold_example_counts) != plan.n_folds: + raise ValueError("chromosome CV fold example counts disagree with n_folds") + flattened = [chrom for group in plan.fold_groups for chrom in group] + if len(flattened) != len(set(flattened)): + raise ValueError("a chromosome occurs in more than one fold group") + if set(flattened) != set(plan.chromosome_counts): + raise ValueError("fold groups do not partition chromosome_counts") + if any(int(count) <= 0 for count in plan.chromosome_counts.values()): + raise ValueError("chromosome example counts must be positive") + expected_fold_counts = tuple( + sum(int(plan.chromosome_counts[chrom]) for chrom in chromosomes) + for chromosomes in plan.fold_groups + ) + if expected_fold_counts != plan.fold_example_counts: + raise ValueError("fold example counts do not equal their chromosome totals") + if sum(expected_fold_counts) != plan.n_examples: + raise ValueError("chromosome CV plan example totals are inconsistent") + if any(not chromosomes for chromosomes in plan.fold_groups): + raise ValueError("every chromosome CV fold group must be non-empty") + + +def create_chromosome_cv_plan( + metadata: Sequence[Mapping[str, object]], + *, + n_folds: int, + group_col: str = "group_chromosome", +) -> ChromosomeCVPlan: + """Greedily balance complete chromosomes by their example counts.""" + + n_folds = int(n_folds) + if n_folds < 3: + raise ValueError("chromosome CV plans require at least three folds") + counts = _count_chromosomes(metadata, str(group_col)) + if len(counts) < n_folds: + raise ValueError( + f"cannot create {n_folds} chromosome folds from only {len(counts)} chromosomes" + ) + ordered = sorted(counts, key=lambda chrom: (-counts[chrom], _natural_key(chrom))) + fold_groups: list[list[str]] = [[] for _ in range(n_folds)] + fold_counts = [0] * n_folds + for chromosome in ordered: + fold = min(range(n_folds), key=lambda index: (fold_counts[index], index)) + fold_groups[fold].append(chromosome) + fold_counts[fold] += counts[chromosome] + normalized_groups = tuple( + tuple(sorted(group, key=_natural_key)) for group in fold_groups + ) + provisional = ChromosomeCVPlan( + n_folds=n_folds, + group_col=str(group_col), + n_examples=len(metadata), + chromosome_counts=dict(counts), + fold_groups=normalized_groups, + fold_example_counts=tuple(fold_counts), + plan_id="", + ) + _validate_plan(provisional) + return ChromosomeCVPlan( + **{ + **provisional.__dict__, + "plan_id": _plan_digest(provisional._content_dict()), + } + ) + + +def save_chromosome_cv_plan( + plan: ChromosomeCVPlan, path: str | Path +) -> Path: + """Write a stable human-readable chromosome CV plan JSON file.""" + + _validate_plan(plan) + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(plan.to_dict(), indent=2) + "\n", encoding="utf-8") + return output + + +def load_chromosome_cv_plan(path: str | Path) -> ChromosomeCVPlan: + """Load and validate an immutable chromosome CV plan.""" + + return ChromosomeCVPlan.from_dict( + json.loads(Path(path).read_text(encoding="utf-8")) + ) + + +def resolve_chromosome_cv_plan( + plan: ChromosomeCVPlan, + metadata: Sequence[Mapping[str, object]], + *, + fold: int, +) -> ChromosomeCVResolution: + """Resolve one test fold, the following validation fold, and training rows.""" + + fold = int(fold) + if fold < 0 or fold >= plan.n_folds: + raise ValueError(f"fold must be in [0, {plan.n_folds})") + observed_counts = _count_chromosomes(metadata, plan.group_col) + if observed_counts != dict(plan.chromosome_counts): + raise ValueError( + "dataset chromosome membership/counts differ from the saved CV plan" + ) + validation_fold = (fold + 1) % plan.n_folds + test_groups = plan.fold_groups[fold] + validation_groups = plan.fold_groups[validation_fold] + train_groups = tuple( + chromosome + for group_index, chromosomes in enumerate(plan.fold_groups) + if group_index not in {fold, validation_fold} + for chromosome in chromosomes + ) + owner = { + **{chrom: "train" for chrom in train_groups}, + **{chrom: "val" for chrom in validation_groups}, + **{chrom: "test" for chrom in test_groups}, + } + indices: dict[str, list[int]] = {"train": [], "val": [], "test": []} + for index, row in enumerate(metadata): + indices[owner[str(row[plan.group_col])]].append(index) + if sum(len(values) for values in indices.values()) != len(metadata): + raise RuntimeError("chromosome CV resolution did not assign every example") + return ChromosomeCVResolution( + fold=fold, + validation_fold=validation_fold, + groups={ + "train": train_groups, + "val": validation_groups, + "test": test_groups, + }, + indices=indices, + ) diff --git a/tests/test_chromosome_cv.py b/tests/test_chromosome_cv.py new file mode 100644 index 0000000..1164948 --- /dev/null +++ b/tests/test_chromosome_cv.py @@ -0,0 +1,184 @@ +import json + +import numpy as np +import pytest + +from transcriptml.cli.main import build_parser, main +from transcriptml.data.bundle import DatasetBundle, save_bundle +from transcriptml.training.trainer import TrainConfig, _select_splits +from transcriptml.workflows.chromosome_cv import ( + create_chromosome_cv_plan, + load_chromosome_cv_plan, + resolve_chromosome_cv_plan, + save_chromosome_cv_plan, +) + + +def _metadata(): + counts = { + "chr1": 23, + "chr2": 17, + "chr3": 15, + "chr4": 11, + "chr5": 9, + "chr6": 8, + "chr7": 6, + "chr8": 5, + "chr9": 4, + "chr10": 2, + } + return [ + {"group_chromosome": chromosome, "row": row} + for chromosome, count in counts.items() + for row in range(count) + ] + + +def test_chromosome_plan_is_deterministic_balanced_and_partitions_groups(): + metadata = _metadata() + first = create_chromosome_cv_plan(metadata, n_folds=5) + second = create_chromosome_cv_plan(metadata, n_folds=5) + assert first == second + assert first.plan_id == second.plan_id + flattened = [chrom for group in first.fold_groups for chrom in group] + assert len(flattened) == len(set(flattened)) == 10 + assert set(flattened) == {row["group_chromosome"] for row in metadata} + assert sum(first.fold_example_counts) == len(metadata) + # Largest-first greedy balancing cannot guarantee equal folds, but this + # realistic skew remains substantially tighter than one largest group. + assert max(first.fold_example_counts) - min(first.fold_example_counts) <= 6 + + +def test_chromosome_plan_rotation_has_no_leakage_and_complete_cv_coverage(): + metadata = _metadata() + plan = create_chromosome_cv_plan(metadata, n_folds=5) + test_occurrences = {chrom: 0 for chrom in plan.chromosome_counts} + val_occurrences = {chrom: 0 for chrom in plan.chromosome_counts} + for fold in range(plan.n_folds): + resolution = resolve_chromosome_cv_plan(plan, metadata, fold=fold) + split_groups = {name: set(values) for name, values in resolution.groups.items()} + assert split_groups["train"].isdisjoint(split_groups["val"]) + assert split_groups["train"].isdisjoint(split_groups["test"]) + assert split_groups["val"].isdisjoint(split_groups["test"]) + assigned = [index for values in resolution.indices.values() for index in values] + assert sorted(assigned) == list(range(len(metadata))) + assert len(assigned) == len(set(assigned)) + for chromosome in split_groups["test"]: + test_occurrences[chromosome] += 1 + for chromosome in split_groups["val"]: + val_occurrences[chromosome] += 1 + assert set(test_occurrences.values()) == {1} + assert set(val_occurrences.values()) == {1} + + +def test_chromosome_plan_roundtrip_and_dataset_mismatch_detection(tmp_path): + metadata = _metadata() + plan = create_chromosome_cv_plan(metadata, n_folds=5) + path = save_chromosome_cv_plan(plan, tmp_path / "cv5.json") + loaded = load_chromosome_cv_plan(path) + assert loaded == plan + assert loaded.to_dict() == json.loads(path.read_text(encoding="utf-8")) + changed = list(metadata) + changed[0] = {**changed[0], "group_chromosome": "chrX"} + with pytest.raises(ValueError, match="differ from the saved CV plan"): + resolve_chromosome_cv_plan(loaded, changed, fold=0) + + +def test_chromosome_plan_integrates_with_training_split_selection(tmp_path): + metadata = _metadata() + plan_path = save_chromosome_cv_plan( + create_chromosome_cv_plan(metadata, n_folds=5), tmp_path / "cv5.json" + ) + bundle = DatasetBundle( + X=np.zeros((len(metadata), 4, 8), dtype=np.float32), + y=np.zeros(len(metadata), dtype=np.float32), + metadata=metadata, + ) + cfg = TrainConfig( + dataset="unused", + output_dir=str(tmp_path / "model"), + cv_plan=str(plan_path), + fold=3, + ) + splits, source = _select_splits(bundle, cfg) + expected = resolve_chromosome_cv_plan( + load_chromosome_cv_plan(plan_path), metadata, fold=3 + ) + assert source == "cv_plan" + assert splits == expected.indices + + +def test_chromosome_plan_cli_create_and_resolve(tmp_path, capsys): + metadata = _metadata() + bundle_dir = tmp_path / "bundle" + save_bundle( + DatasetBundle( + X=np.zeros((len(metadata), 4, 8), dtype=np.uint8), + ids=[f"id{i}" for i in range(len(metadata))], + metadata=metadata, + ), + bundle_dir, + ) + plan_path = tmp_path / "cv5.json" + main( + [ + "cv", + "create-chromosome-plan", + "--dataset", + str(bundle_dir), + "--output", + str(plan_path), + "--n-folds", + "5", + ] + ) + assert capsys.readouterr().out.strip() == str(plan_path) + splits_path = tmp_path / "fold2.json" + main( + [ + "cv", + "resolve-plan", + "--dataset", + str(bundle_dir), + "--cv-plan", + str(plan_path), + "--fold", + "2", + "--output", + str(splits_path), + ] + ) + assert capsys.readouterr().out.strip() == str(splits_path) + resolved = json.loads(splits_path.read_text(encoding="utf-8")) + assert resolved["fold"] == 2 + assert resolved["validation_fold"] == 3 + assert sum(len(values) for values in resolved["indices"].values()) == len(metadata) + + +def test_chromosome_plan_requires_enough_chromosomes(): + with pytest.raises(ValueError, match="only 1 chromosomes"): + create_chromosome_cv_plan( + [{"group_chromosome": "chr21"}] * 10, + n_folds=5, + ) + + +def test_train_cli_accepts_plan_and_job_array_overrides(): + args = build_parser().parse_args( + [ + "train", + "train.json", + "--cv-plan", + "cv5.json", + "--fold", + "3", + "--dataset", + "bundle", + "--output-dir", + "runs/fold3", + ] + ) + assert args.cv_plan == "cv5.json" + assert args.fold == 3 + assert args.dataset == "bundle" + assert args.output_dir == "runs/fold3" diff --git a/tests/test_rbpnet.py b/tests/test_rbpnet.py index dbce0cd..5069022 100644 --- a/tests/test_rbpnet.py +++ b/tests/test_rbpnet.py @@ -31,6 +31,7 @@ extract_bam_to_store, five_prime_reference_position, read1_rna_strand, + write_sparse_signal_chunkwise, ) from transcriptml.rbpnet.windows import WindowScanConfig, generate_window_bounds, scan_windows @@ -188,6 +189,56 @@ def test_bam_assignment_reports_transcript_incompatibility(tmp_path): assert counts.tolist() == [1] +@pytest.mark.parametrize( + ("compression", "level"), + [("gzip", 1), ("gzip", 4), ("lzf", None), (None, None)], +) +def test_chunkwise_sparse_signal_write_matches_fancy_indexing_and_skips_empty_chunks( + tmp_path, compression, level +): + positions = np.asarray([1, 7, 8, 15, 33, 39], dtype=np.int64) + values = np.asarray([2, 4, 1, 8, 3, 9], dtype=np.uint64) + expected = np.zeros(40, dtype=np.uint32) + expected[positions] = values.astype(np.uint32) + path = tmp_path / f"signal_{compression or 'none'}_{level}.h5" + with h5py.File(path, "w") as store: + kwargs = {"compression": compression, "shuffle": compression is not None} + if compression == "gzip": + kwargs["compression_opts"] = level + dataset = store.create_dataset( + "counts", + shape=(1, 40), + dtype=np.uint32, + chunks=(1, 8), + fillvalue=0, + **kwargs, + ) + batches = [ + (positions[:3], values[:3]), + (positions[3:5], values[3:5]), + (positions[5:], values[5:]), + ] + assert write_sparse_signal_chunkwise(dataset, 0, batches) == len(positions) + np.testing.assert_array_equal(dataset[0], expected) + # Positions touch chunks 0, 1, and 4. Chunks 2 and 3 stay at the + # HDF5 fill value and are never explicitly allocated. + if hasattr(dataset.id, "get_num_chunks"): + assert dataset.id.get_num_chunks() == 3 + + +def test_chunkwise_sparse_signal_write_validates_sorted_unique_positions(tmp_path): + with h5py.File(tmp_path / "invalid.h5", "w") as store: + dataset = store.create_dataset( + "counts", shape=(1, 16), dtype=np.uint32, chunks=(1, 8), fillvalue=0 + ) + with pytest.raises(ValueError, match="strictly increasing"): + write_sparse_signal_chunkwise( + dataset, + 0, + [(np.asarray([4, 3]), np.asarray([1, 1]))], + ) + + def test_gene_space_bam_assignment_retains_intronic_and_spliced_reads(tmp_path): gene = _gene_tx("+") mature = _junction_tx("+") diff --git a/tests/test_rbpnet_model.py b/tests/test_rbpnet_model.py index 39b4220..f7b7880 100644 --- a/tests/test_rbpnet_model.py +++ b/tests/test_rbpnet_model.py @@ -21,6 +21,11 @@ from transcriptml.training.splits import group_split_indices, validate_group_disjoint from transcriptml.training.evaluation import evaluate_checkpoint from transcriptml.training.trainer import train_model +from transcriptml.workflows.chromosome_cv import ( + create_chromosome_cv_plan, + load_chromosome_cv_plan, + save_chromosome_cv_plan, +) def _synthetic_bundle(n=9, length=16, jitter=2): @@ -66,7 +71,7 @@ def _synthetic_bundle(n=9, length=16, jitter=2): "profile_materialized_end": materialized_start + width, "group_gene_id": f"g{index // 3}", "group_transcript_id": f"tx{index // 3}", - "group_chromosome": "chr1", + "group_chromosome": f"chr{index // 3 + 1}", } ) arrays = { @@ -429,3 +434,50 @@ def step(update): step(True) final = step(False) assert final < initial + + +def test_rbpnet_training_consumes_saved_chromosome_cv_plan(tmp_path): + bundle = _synthetic_bundle() + plan_path = save_chromosome_cv_plan( + create_chromosome_cv_plan( + bundle.metadata, n_folds=3, group_col="group_chromosome" + ), + tmp_path / "cv3.json", + ) + result = train_model( + bundle, + { + "dataset": "unused", + "output_dir": str(tmp_path / "fold1"), + "batch_size": 3, + "epochs": 1, + "patience": 0, + "progress": False, + "learning_rate": 0.005, + "model": { + "name": "rbpnet", + "params": { + "n_filters": 8, + "initial_kernel_size": 3, + "n_residual_blocks": 1, + "residual_kernel_size": 3, + "dilations": [1], + "normalization": "none", + "dropout": 0.0, + "profile_head_kernel_size": 3, + "profile_length": 16, + "enrichment_head_type": "none", + }, + }, + "loss": {"name": "rbpnet"}, + "cv_plan": str(plan_path), + "fold": 1, + }, + ) + assert result["summary"]["split_source_used"] == "cv_plan" + assert result["summary"]["cv_plan_id"] == load_chromosome_cv_plan( + plan_path + ).plan_id + assert result["summary"]["fold"] == 1 + assert result["summary"]["split_group_col"] == "group_chromosome" + assert result["summary"]["split_counts"] == {"train": 3, "val": 3, "test": 3} From 0951a7adb24afbe01e27a31e83fb9efb1fcb9b6d Mon Sep 17 00:00:00 2001 From: isvock Date: Wed, 12 Aug 2026 10:40:44 -0700 Subject: [PATCH 06/12] Add optional region-type filtering --- docs/rbpnet.md | 27 ++++++ scripts/rbpnet/README.md | 2 + scripts/rbpnet/rbpnet_config.sh | 2 + scripts/rbpnet/scan_select_bundle.sh | 3 + src/transcriptml/rbpnet/cli.py | 18 ++++ src/transcriptml/rbpnet/selection.py | 100 +++++++++++++++++++- tests/test_rbpnet.py | 132 +++++++++++++++++++++++++++ 7 files changed, 281 insertions(+), 3 deletions(-) diff --git a/docs/rbpnet.md b/docs/rbpnet.md index 2385c4c..097e93e 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -208,6 +208,33 @@ Selection asks which experimental loci are eligible and why. It writes a versioned `*.parquet` manifest, equivalent `*.tsv.gz`, and a `*.selection.json` provenance sidecar. +Every selector can restrict its candidate universe to exact scanner +annotations without regenerating the descriptive window table: + +```bash +transcriptml rbpnet select-regions \ + --processed-dir processed/chr21 \ + --windows processed/chr21_windows_100nt.parquet \ + --strategy peak_gray_negative \ + --region-types 5putr,cds,3putr \ + --output-prefix processed/chr21_exonic_selection +``` + +Valid values are `5putr`, `cds`, `3putr`, `noncoding_exon`, `intron`, and +`mixed`. The default is all types. Filtering is exact: for example, +`--region-types 3putr` accepts only windows wholly contained in 3' UTR and +does not accept a boundary-crossing `mixed` window. Include `mixed` +explicitly when desired. + +The restriction is applied before each strategy's signal/statistical rules. +Consequently, excluded windows do not affect coverage eligibility, the +original selector's testing/50-nt advance, or the `peak_gray_negative` BH +correction universe. The published IP locus-density Poisson null remains based +on the complete locus; this option restricts which windows are tested, not how +that published null is defined. The selection provenance records the requested +types and source/eligible window counts, while the sidecar also reports +selected-example counts by region type. + ### Published RBPNet v1 First make the published 100-nt, stride-1 descriptive scan, then select: diff --git a/scripts/rbpnet/README.md b/scripts/rbpnet/README.md index 6b71985..8e18bf6 100644 --- a/scripts/rbpnet/README.md +++ b/scripts/rbpnet/README.md @@ -36,6 +36,8 @@ The data-construction script chooses stride 1 automatically for `original_rbpnet` and stride 50 for the other selectors unless `WINDOW_STRIDE` is explicitly set. Its default selector is `peak_gray_negative`; all thresholds remain editable in the config. +Set `REGION_TYPES` to an exact comma-separated selection universe such as +`3putr` or `cds,3putr`; leave it empty to preserve all region types. The CV stage counts examples per chromosome, greedily balances whole chromosomes across `N_FOLDS`, and writes `CV_PLAN` once. Every fold job loads diff --git a/scripts/rbpnet/rbpnet_config.sh b/scripts/rbpnet/rbpnet_config.sh index 0c24a79..a48935a 100644 --- a/scripts/rbpnet/rbpnet_config.sh +++ b/scripts/rbpnet/rbpnet_config.sh @@ -55,6 +55,8 @@ OVERWRITE="${OVERWRITE:-0}" WINDOW_SIZE="${WINDOW_SIZE:-100}" WINDOW_STRIDE="${WINDOW_STRIDE:-}" MIN_SMINPUT_TPM="${MIN_SMINPUT_TPM:-0}" +# Optional comma-separated exact annotations, e.g. "3putr" or "cds,3putr". +REGION_TYPES="${REGION_TYPES:-}" SELECTION_STRATEGY="${SELECTION_STRATEGY:-peak_gray_negative}" MIN_TOTAL_COUNT="${MIN_TOTAL_COUNT:-8}" MIN_IP_COUNT="${MIN_IP_COUNT:-0}" diff --git a/scripts/rbpnet/scan_select_bundle.sh b/scripts/rbpnet/scan_select_bundle.sh index de2e4d0..915aeb3 100644 --- a/scripts/rbpnet/scan_select_bundle.sh +++ b/scripts/rbpnet/scan_select_bundle.sh @@ -54,6 +54,9 @@ selection_args=( --min-sminput-tpm "${MIN_SMINPUT_TPM}" --output-prefix "${SELECTION_PREFIX}" ) +if [[ -n "${REGION_TYPES}" ]]; then + selection_args+=(--region-types "${REGION_TYPES}") +fi case "${SELECTION_STRATEGY}" in original_rbpnet) selection_args+=(--poisson-null "${POISSON_NULL}") diff --git a/src/transcriptml/rbpnet/cli.py b/src/transcriptml/rbpnet/cli.py index fd6f4eb..5ff0b63 100644 --- a/src/transcriptml/rbpnet/cli.py +++ b/src/transcriptml/rbpnet/cli.py @@ -10,6 +10,13 @@ _VALID_SAMPLE = re.compile(r"^[A-Za-z][A-Za-z0-9_.-]*$") +def _csv_tuple(value: str) -> tuple[str, ...]: + values = tuple(part.strip() for part in value.split(",") if part.strip()) + if not values: + raise argparse.ArgumentTypeError("expected one or more comma-separated values") + return values + + def add_rbpnet_parser(subparsers) -> None: """Add the nested ``transcriptml rbpnet`` command family.""" @@ -112,6 +119,16 @@ def add_rbpnet_parser(subparsers) -> None: help="optional pooled/per-replicate IP minimum (default: 0)", ) select.add_argument("--min-sminput-tpm", type=float, default=0.0) + select.add_argument( + "--region-types", + type=_csv_tuple, + default=None, + metavar="TYPE[,TYPE...]", + help=( + "restrict selection to exact window annotations: 5putr, cds, 3putr, " + "noncoding_exon, intron, mixed (default: all)" + ), + ) select.add_argument( "--replicate-mode", choices=("combined", "per_ip"), default="per_ip", help="broad_coverage eligibility mode (default: per_ip)", @@ -224,6 +241,7 @@ def run_rbpnet_command(args: argparse.Namespace, parser: argparse.ArgumentParser min_sminput_count=args.min_sminput_count, min_ip_count=args.min_ip_count, min_sminput_tpm=args.min_sminput_tpm, + region_types=args.region_types, replicate_mode=args.replicate_mode, peak_fdr=args.peak_fdr, peak_min_log2_ratio=args.peak_min_log2_ratio, diff --git a/src/transcriptml/rbpnet/selection.py b/src/transcriptml/rbpnet/selection.py index 7dd3034..27a5362 100644 --- a/src/transcriptml/rbpnet/selection.py +++ b/src/transcriptml/rbpnet/selection.py @@ -55,6 +55,8 @@ class SelectionConfig: min_ip_count: int = 0 min_sminput_tpm: float = 0.0 replicate_mode: str = "per_ip" + # Optional exact window-annotation universe. None preserves all regions. + region_types: tuple[str, ...] | None = None # Peak / gray / confident-negative selector. peak_fdr: float = 0.05 peak_min_log2_ratio: float = 1.0 @@ -99,6 +101,55 @@ def _iter_window_rows(path: Path, *, batch_size: int) -> Iterator[dict]: yield from batch.to_pylist() +def _normalize_region_types( + region_types: str | Iterable[str] | None, +) -> tuple[str, ...] | None: + """Normalize exact region-type filters while preserving user order.""" + + if region_types is None: + return None + raw_values = (region_types,) if isinstance(region_types, str) else region_types + values: list[str] = [] + for raw in raw_values: + values.extend(part.strip().lower() for part in str(raw).split(",")) + values = [value for value in values if value] + if not values: + raise ValueError("region_types must contain at least one region type") + return tuple(dict.fromkeys(values)) + + +def _region_is_eligible(config: SelectionConfig, row: dict) -> bool: + return config.region_types is None or str(row["region_type"]) in config.region_types + + +def _window_region_counts( + path: Path, + region_types: tuple[str, ...] | None, + *, + batch_size: int, +) -> tuple[dict[str, int], dict[str, int]]: + """Count source and region-filter-eligible descriptive windows.""" + + source: Counter[str] = Counter() + eligible: Counter[str] = Counter() + parquet = pq.ParquetFile(path) + for batch in parquet.iter_batches(batch_size=batch_size, columns=["region_type"]): + for value in batch.column(0).to_pylist(): + region_type = str(value) + source[region_type] += 1 + if region_types is None or region_type in region_types: + eligible[region_type] += 1 + order = {name: index for index, name in enumerate((*REGION_TYPES, "mixed"))} + + def sort_key(item: tuple[str, int]) -> tuple[int, str]: + return order.get(item[0], len(order)), item[0] + + return ( + dict(sorted(source.items(), key=sort_key)), + dict(sorted(eligible.items(), key=sort_key)), + ) + + def _stable_example_id( strategy: str, transcript_id: str, @@ -326,6 +377,8 @@ def _original_rows( try: for row in _iter_window_rows(windows_path, batch_size=config.batch_size): reporter.update() + if not _region_is_eligible(config, row): + continue tx_id = row["transcript_id"] if tx_id != current_tx: tx = ds.get_transcript(tx_id) @@ -386,6 +439,8 @@ def _broad_coverage_rows( try: for row in _iter_window_rows(windows_path, batch_size=config.batch_size): reporter.update() + if not _region_is_eligible(config, row): + continue if float(row["sminput_tpm"]) < config.min_sminput_tpm: continue input_count = int(row[f"{input_name}_count"]) @@ -433,15 +488,23 @@ def _peak_statistics( windows_path, columns=[ f"{input_name}_count", "ip_pooled_count", "total_ip_sminput_count", - "sminput_tpm", + "sminput_tpm", "region_type", ], ) input_counts = table[f"{input_name}_count"].to_numpy(zero_copy_only=False).astype(np.int64) ip_counts = table["ip_pooled_count"].to_numpy(zero_copy_only=False).astype(np.int64) totals = table["total_ip_sminput_count"].to_numpy(zero_copy_only=False).astype(np.int64) tpm = table["sminput_tpm"].to_numpy(zero_copy_only=False).astype(np.float64) + if config.region_types is None: + eligible_region = np.ones(len(totals), dtype=bool) + else: + region_types = np.asarray( + table["region_type"].to_pylist(), dtype=object + ) + eligible_region = np.isin(region_types, config.region_types) adequate = ( - (totals >= config.min_total_count) + eligible_region + & (totals >= config.min_total_count) & (input_counts >= config.min_sminput_count) & (ip_counts >= config.min_ip_count) & (tpm >= config.min_sminput_tpm) @@ -458,7 +521,8 @@ def _peak_statistics( ip_counts[adequate], totals[adequate], null_ip_probability ) log_progress( - f"rbpnet select-regions: {int(adequate.sum()):,}/{len(adequate):,} windows adequately measured", + f"rbpnet select-regions: {int(adequate.sum()):,}/{int(eligible_region.sum()):,} " + "region-eligible windows adequately measured", enabled=config.progress, ) return ( @@ -607,6 +671,14 @@ def _validate_config(config: SelectionConfig) -> None: raise ValueError("sminput_poisson_pseudocount must be positive") if config.replicate_mode not in {"combined", "per_ip"}: raise ValueError("replicate_mode must be combined or per_ip") + valid_region_types = set(REGION_TYPES) | {"mixed"} + if config.region_types is not None: + invalid = sorted(set(config.region_types) - valid_region_types) + if invalid: + raise ValueError( + "unsupported region_types: " + f"{', '.join(invalid)}; choose from {', '.join((*REGION_TYPES, 'mixed'))}" + ) def _validate_scan_dataset( @@ -649,6 +721,10 @@ def _validate_scan_dataset( def select_regions(config: SelectionConfig) -> dict: """Select biological loci and write a versioned lightweight manifest.""" + config = replace( + config, + region_types=_normalize_region_types(config.region_types), + ) if config.min_total_count is None: config = replace( config, @@ -673,6 +749,11 @@ def select_regions(config: SelectionConfig) -> dict: if any(sample.effective_library_size is None or sample.effective_library_size <= 0 for sample in ds.samples): raise ValueError("all samples need positive effective_library_size values for selection") _validate_scan_dataset(ds, windows_path, scan_metadata) + source_region_counts, eligible_region_counts = _window_region_counts( + windows_path, + config.region_types, + batch_size=config.batch_size, + ) provenance = { "format": "transcriptml-rbpnet-selection", "format_version": "1", @@ -681,6 +762,14 @@ def select_regions(config: SelectionConfig) -> dict: "coordinate_space": ds.coordinate_space, "source_windows": str(windows_path.resolve()), "window_scan": scan_metadata, + "region_filter": { + "mode": "all" if config.region_types is None else "exact_region_type", + "allowed_region_types": ( + None if config.region_types is None else list(config.region_types) + ), + "source_window_counts": source_region_counts, + "eligible_window_counts": eligible_region_counts, + }, "configuration": { key: (str(value) if isinstance(value, Path) else value) for key, value in config.__dict__.items() @@ -724,6 +813,7 @@ def select_regions(config: SelectionConfig) -> dict: selected = 0 state_counts: Counter[str] = Counter() + selected_region_counts: Counter[str] = Counter() transcript_ids: set[str] = set() batch: list[dict] = [] reporter = ProgressReporter( @@ -751,6 +841,7 @@ def flush() -> None: batch.append(row) selected += 1 state_counts[row["selection_state"]] += 1 + selected_region_counts[row["region_type"]] += 1 transcript_ids.add(row["transcript_id"]) reporter.update() if len(batch) >= config.batch_size: @@ -762,6 +853,9 @@ def flush() -> None: "n_examples": selected, "n_transcripts": len(transcript_ids), "state_counts": dict(sorted(state_counts.items())), + "selected_example_region_counts": dict( + sorted(selected_region_counts.items()) + ), "parquet": str(parquet_path), "tsv": str(tsv_path), } diff --git a/tests/test_rbpnet.py b/tests/test_rbpnet.py index 5069022..ae7f3dd 100644 --- a/tests/test_rbpnet.py +++ b/tests/test_rbpnet.py @@ -10,6 +10,7 @@ import pytest from scipy.stats import poisson +from transcriptml.cli.main import build_parser from transcriptml.data.encoding import encode_rna_sequence from transcriptml.rbpnet.bundle import ( RBPNetBundleConfig, @@ -594,6 +595,109 @@ def test_selection_strategies_ids_serialization_and_stitching(tmp_path): assert any(row["source_window_count"] > 1 for row in load_selection_manifest(classified).rows) +def test_exact_region_type_filtering_provenance_and_peak_bh_universe(tmp_path): + root = tmp_path / "processed" + _write_processed_fixture(root) + windows = tmp_path / "windows" + scan_windows(WindowScanConfig( + processed_dir=root, output_prefix=windows, window_size=4, stride=2, + progress=False, + )) + + broad = tmp_path / "broad_regions" + broad_summary = select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=broad, + strategy="broad_coverage", + replicate_mode="combined", + min_total_count=0, + region_types=("5putr", "cds"), + progress=False, + )) + broad_manifest = load_selection_manifest(broad) + assert [row["region_type"] for row in broad_manifest.rows] == ["5putr", "cds"] + assert broad_summary["region_filter"] == { + "mode": "exact_region_type", + "allowed_region_types": ["5putr", "cds"], + "source_window_counts": {"5putr": 1, "cds": 1, "mixed": 3}, + "eligible_window_counts": {"5putr": 1, "cds": 1}, + } + assert broad_summary["selected_example_region_counts"] == { + "5putr": 1, + "cds": 1, + } + assert broad_manifest.metadata["configuration"]["region_types"] == [ + "5putr", + "cds", + ] + + mixed = tmp_path / "broad_mixed" + mixed_summary = select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=mixed, + strategy="broad_coverage", + replicate_mode="combined", + min_total_count=0, + region_types="mixed", + progress=False, + )) + assert mixed_summary["n_examples"] == 3 + assert {row["region_type"] for row in load_selection_manifest(mixed).rows} == { + "mixed" + } + + classified = tmp_path / "classified_cds" + select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=classified, + strategy="peak_gray_negative", + min_total_count=0, + peak_fdr=1.0, + negative_fdr=1.0, + peak_min_log2_ratio=100.0, + negative_max_log2_ratio=-100.0, + region_types=("cds",), + progress=False, + )) + classified_rows = load_selection_manifest(classified).rows + assert len(classified_rows) == 1 + assert classified_rows[0]["region_type"] == "cds" + # CDS contributes one adequately measured test, so BH adjustment over the + # requested region universe leaves each exact-tail p-value unchanged. + assert classified_rows[0]["source_min_enrichment_qvalue"] == pytest.approx( + classified_rows[0]["source_min_enrichment_pvalue"] + ) + assert classified_rows[0]["source_min_depletion_qvalue"] == pytest.approx( + classified_rows[0]["source_min_depletion_pvalue"] + ) + + with pytest.raises(ValueError, match="unsupported region_types: promoter"): + select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=tmp_path / "invalid_region", + strategy="broad_coverage", + region_types=("promoter",), + progress=False, + )) + + +def test_select_regions_cli_parses_comma_separated_region_types(): + args = build_parser().parse_args([ + "rbpnet", + "select-regions", + "--processed-dir", "processed", + "--windows", "windows.parquet", + "--output-prefix", "selected", + "--strategy", "broad_coverage", + "--region-types", "3putr,cds", + ]) + assert args.region_types == ("3putr", "cds") + + def test_zero_count_peak_negative_edges_and_broad_coverage_defaults(tmp_path): profiles = np.zeros((3, 12), dtype=np.uint32) profiles[0, 0] = 30 # informative input-only window @@ -677,6 +781,34 @@ def test_original_rbpnet_poisson_selection_and_50nt_advance(tmp_path): [row["selection_pvalue"] for row in rows], ) + matching_region = tmp_path / "original_noncoding" + matching_summary = select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=matching_region, + strategy="original_rbpnet", + region_types=("noncoding_exon",), + progress=False, + )) + assert [row["example_id"] for row in load_selection_manifest(matching_region).rows] == [ + row["example_id"] for row in rows + ] + assert matching_summary["region_filter"]["eligible_window_counts"] == { + "noncoding_exon": 401 + } + + excluded_region = tmp_path / "original_cds" + excluded_summary = select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=excluded_region, + strategy="original_rbpnet", + region_types=("cds",), + progress=False, + )) + assert excluded_summary["n_examples"] == 0 + assert excluded_summary["region_filter"]["eligible_window_counts"] == {} + sminput_null = tmp_path / "original_sminput" sminput_summary = select_regions(SelectionConfig( processed_dir=root, windows=windows, output_prefix=sminput_null, From 4d1e08e77f2e252182acec024f54b36753a34430 Mon Sep 17 00:00:00 2001 From: isvock Date: Wed, 12 Aug 2026 11:43:28 -0700 Subject: [PATCH 07/12] Better handling of mixed region windows --- docs/rbpnet.md | 29 ++++--- scripts/rbpnet/README.md | 6 +- scripts/rbpnet/rbpnet_config.sh | 4 +- scripts/rbpnet/scan_select_bundle.sh | 8 ++ src/transcriptml/rbpnet/cli.py | 17 ++++- src/transcriptml/rbpnet/selection.py | 73 ++++++++++++++---- tests/test_rbpnet.py | 109 +++++++++++++++++++++++---- 7 files changed, 203 insertions(+), 43 deletions(-) diff --git a/docs/rbpnet.md b/docs/rbpnet.md index 097e93e..5ee708e 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -208,8 +208,8 @@ Selection asks which experimental loci are eligible and why. It writes a versioned `*.parquet` manifest, equivalent `*.tsv.gz`, and a `*.selection.json` provenance sidecar. -Every selector can restrict its candidate universe to exact scanner -annotations without regenerating the descriptive window table: +Every selector can restrict its candidate universe by scanner annotation +without regenerating the descriptive window table: ```bash transcriptml rbpnet select-regions \ @@ -220,14 +220,25 @@ transcriptml rbpnet select-regions \ --output-prefix processed/chr21_exonic_selection ``` -Valid values are `5putr`, `cds`, `3putr`, `noncoding_exon`, `intron`, and -`mixed`. The default is all types. Filtering is exact: for example, -`--region-types 3putr` accepts only windows wholly contained in 3' UTR and -does not accept a boundary-crossing `mixed` window. Include `mixed` -explicitly when desired. +Valid values are the biological classes `5putr`, `cds`, `3putr`, +`noncoding_exon`, and `intron`. `mixed` is an annotation status, not a region +class, and therefore is not a valid value for `--region-types`. -The restriction is applied before each strategy's signal/statistical rules. -Consequently, excluded windows do not affect coverage eligibility, the +By default, a requested class includes both pure windows and mixed windows +with a positive overlap with that class. Thus `--region-types 3putr` includes +windows wholly contained in 3' UTR and CDS/3' UTR boundary windows. Two flags +control boundary-crossing windows: + +- `--discard-mixed` retains only pure windows of the requested classes. +- `--only-mixed` retains only boundary-crossing windows overlapping the + requested classes. It requires `--region-types`. + +The flags are mutually exclusive. Omitting `--region-types` preserves the +complete window universe; `--discard-mixed` may still be used to remove all +boundary-crossing windows. + +The overlap restriction and mixed-window policy are applied before each +strategy's signal/statistical rules. Consequently, excluded windows do not affect coverage eligibility, the original selector's testing/50-nt advance, or the `peak_gray_negative` BH correction universe. The published IP locus-density Poisson null remains based on the complete locus; this option restricts which windows are tested, not how diff --git a/scripts/rbpnet/README.md b/scripts/rbpnet/README.md index 8e18bf6..f6d4474 100644 --- a/scripts/rbpnet/README.md +++ b/scripts/rbpnet/README.md @@ -36,8 +36,10 @@ The data-construction script chooses stride 1 automatically for `original_rbpnet` and stride 50 for the other selectors unless `WINDOW_STRIDE` is explicitly set. Its default selector is `peak_gray_negative`; all thresholds remain editable in the config. -Set `REGION_TYPES` to an exact comma-separated selection universe such as -`3putr` or `cds,3putr`; leave it empty to preserve all region types. +Set `REGION_TYPES` to a comma-separated selection universe such as `3putr` or +`cds,3putr`; matching boundary-crossing windows are included by default. Set +`DISCARD_MIXED=1` to keep only pure windows or `ONLY_MIXED=1` to keep only +matching boundary-crossing windows. The CV stage counts examples per chromosome, greedily balances whole chromosomes across `N_FOLDS`, and writes `CV_PLAN` once. Every fold job loads diff --git a/scripts/rbpnet/rbpnet_config.sh b/scripts/rbpnet/rbpnet_config.sh index a48935a..03ecd46 100644 --- a/scripts/rbpnet/rbpnet_config.sh +++ b/scripts/rbpnet/rbpnet_config.sh @@ -55,8 +55,10 @@ OVERWRITE="${OVERWRITE:-0}" WINDOW_SIZE="${WINDOW_SIZE:-100}" WINDOW_STRIDE="${WINDOW_STRIDE:-}" MIN_SMINPUT_TPM="${MIN_SMINPUT_TPM:-0}" -# Optional comma-separated exact annotations, e.g. "3putr" or "cds,3putr". +# Optional comma-separated biological annotations, e.g. "3putr" or "cds,3putr". REGION_TYPES="${REGION_TYPES:-}" +DISCARD_MIXED="${DISCARD_MIXED:-0}" +ONLY_MIXED="${ONLY_MIXED:-0}" SELECTION_STRATEGY="${SELECTION_STRATEGY:-peak_gray_negative}" MIN_TOTAL_COUNT="${MIN_TOTAL_COUNT:-8}" MIN_IP_COUNT="${MIN_IP_COUNT:-0}" diff --git a/scripts/rbpnet/scan_select_bundle.sh b/scripts/rbpnet/scan_select_bundle.sh index 915aeb3..ac4523c 100644 --- a/scripts/rbpnet/scan_select_bundle.sh +++ b/scripts/rbpnet/scan_select_bundle.sh @@ -57,6 +57,14 @@ selection_args=( if [[ -n "${REGION_TYPES}" ]]; then selection_args+=(--region-types "${REGION_TYPES}") fi +if [[ "${DISCARD_MIXED}" == "1" && "${ONLY_MIXED}" == "1" ]]; then + echo "DISCARD_MIXED and ONLY_MIXED cannot both be enabled." >&2 + exit 1 +elif [[ "${DISCARD_MIXED}" == "1" ]]; then + selection_args+=(--discard-mixed) +elif [[ "${ONLY_MIXED}" == "1" ]]; then + selection_args+=(--only-mixed) +fi case "${SELECTION_STRATEGY}" in original_rbpnet) selection_args+=(--poisson-null "${POISSON_NULL}") diff --git a/src/transcriptml/rbpnet/cli.py b/src/transcriptml/rbpnet/cli.py index 5ff0b63..5e78e77 100644 --- a/src/transcriptml/rbpnet/cli.py +++ b/src/transcriptml/rbpnet/cli.py @@ -125,10 +125,21 @@ def add_rbpnet_parser(subparsers) -> None: default=None, metavar="TYPE[,TYPE...]", help=( - "restrict selection to exact window annotations: 5putr, cds, 3putr, " - "noncoding_exon, intron, mixed (default: all)" + "restrict selection by overlap with: 5putr, cds, 3putr, " + "noncoding_exon, intron; matching mixed windows are included by default" ), ) + mixed = select.add_mutually_exclusive_group() + mixed.add_argument( + "--discard-mixed", + action="store_true", + help="exclude boundary-crossing windows from the requested region types", + ) + mixed.add_argument( + "--only-mixed", + action="store_true", + help="select only boundary-crossing windows overlapping --region-types", + ) select.add_argument( "--replicate-mode", choices=("combined", "per_ip"), default="per_ip", help="broad_coverage eligibility mode (default: per_ip)", @@ -242,6 +253,8 @@ def run_rbpnet_command(args: argparse.Namespace, parser: argparse.ArgumentParser min_ip_count=args.min_ip_count, min_sminput_tpm=args.min_sminput_tpm, region_types=args.region_types, + discard_mixed=args.discard_mixed, + only_mixed=args.only_mixed, replicate_mode=args.replicate_mode, peak_fdr=args.peak_fdr, peak_min_log2_ratio=args.peak_min_log2_ratio, diff --git a/src/transcriptml/rbpnet/selection.py b/src/transcriptml/rbpnet/selection.py index 27a5362..7e8a7e3 100644 --- a/src/transcriptml/rbpnet/selection.py +++ b/src/transcriptml/rbpnet/selection.py @@ -56,7 +56,9 @@ class SelectionConfig: min_sminput_tpm: float = 0.0 replicate_mode: str = "per_ip" # Optional exact window-annotation universe. None preserves all regions. - region_types: tuple[str, ...] | None = None + region_types: tuple[str, ...] | str | None = None + discard_mixed: bool = False + only_mixed: bool = False # Peak / gray / confident-negative selector. peak_fdr: float = 0.05 peak_min_log2_ratio: float = 1.0 @@ -119,12 +121,27 @@ def _normalize_region_types( def _region_is_eligible(config: SelectionConfig, row: dict) -> bool: - return config.region_types is None or str(row["region_type"]) in config.region_types + """Match a pure annotation or a mixed window overlapping requested types.""" + + region_type = str(row["region_type"]) + is_mixed = region_type == "mixed" + if config.only_mixed and not is_mixed: + return False + if config.discard_mixed and is_mixed: + return False + if config.region_types is None: + return True + if not is_mixed: + return region_type in config.region_types + return any( + int(row.get(f"region_{requested}_nt", 0)) > 0 + for requested in config.region_types + ) def _window_region_counts( path: Path, - region_types: tuple[str, ...] | None, + config: SelectionConfig, *, batch_size: int, ) -> tuple[dict[str, int], dict[str, int]]: @@ -133,11 +150,12 @@ def _window_region_counts( source: Counter[str] = Counter() eligible: Counter[str] = Counter() parquet = pq.ParquetFile(path) - for batch in parquet.iter_batches(batch_size=batch_size, columns=["region_type"]): - for value in batch.column(0).to_pylist(): - region_type = str(value) + columns = ["region_type", *(f"region_{name}_nt" for name in REGION_TYPES)] + for batch in parquet.iter_batches(batch_size=batch_size, columns=columns): + for row in batch.to_pylist(): + region_type = str(row["region_type"]) source[region_type] += 1 - if region_types is None or region_type in region_types: + if _region_is_eligible(config, row): eligible[region_type] += 1 order = {name: index for index, name in enumerate((*REGION_TYPES, "mixed"))} @@ -489,19 +507,30 @@ def _peak_statistics( columns=[ f"{input_name}_count", "ip_pooled_count", "total_ip_sminput_count", "sminput_tpm", "region_type", + *(f"region_{name}_nt" for name in REGION_TYPES), ], ) input_counts = table[f"{input_name}_count"].to_numpy(zero_copy_only=False).astype(np.int64) ip_counts = table["ip_pooled_count"].to_numpy(zero_copy_only=False).astype(np.int64) totals = table["total_ip_sminput_count"].to_numpy(zero_copy_only=False).astype(np.int64) tpm = table["sminput_tpm"].to_numpy(zero_copy_only=False).astype(np.float64) + region_labels = np.asarray(table["region_type"].to_pylist(), dtype=object) + mixed = region_labels == "mixed" if config.region_types is None: eligible_region = np.ones(len(totals), dtype=bool) else: - region_types = np.asarray( - table["region_type"].to_pylist(), dtype=object - ) - eligible_region = np.isin(region_types, config.region_types) + pure_match = np.isin(region_labels, config.region_types) + mixed_match = np.zeros(len(totals), dtype=bool) + for region_type in config.region_types: + overlap = table[f"region_{region_type}_nt"].to_numpy( + zero_copy_only=False + ) + mixed_match |= mixed & (overlap > 0) + eligible_region = pure_match | mixed_match + if config.discard_mixed: + eligible_region &= ~mixed + if config.only_mixed: + eligible_region &= mixed adequate = ( eligible_region & (totals >= config.min_total_count) @@ -671,14 +700,18 @@ def _validate_config(config: SelectionConfig) -> None: raise ValueError("sminput_poisson_pseudocount must be positive") if config.replicate_mode not in {"combined", "per_ip"}: raise ValueError("replicate_mode must be combined or per_ip") - valid_region_types = set(REGION_TYPES) | {"mixed"} + valid_region_types = set(REGION_TYPES) if config.region_types is not None: invalid = sorted(set(config.region_types) - valid_region_types) if invalid: raise ValueError( "unsupported region_types: " - f"{', '.join(invalid)}; choose from {', '.join((*REGION_TYPES, 'mixed'))}" + f"{', '.join(invalid)}; choose from {', '.join(REGION_TYPES)}" ) + if config.discard_mixed and config.only_mixed: + raise ValueError("discard_mixed and only_mixed are mutually exclusive") + if config.only_mixed and config.region_types is None: + raise ValueError("only_mixed requires one or more region_types") def _validate_scan_dataset( @@ -709,6 +742,7 @@ def _validate_scan_dataset( "total_ip_sminput_count", "log2_ip_pooled_vs_sminput", "max_ip_pooled_5pend", "genomic_blocks", } + required_columns.update(f"region_{name}_nt" for name in REGION_TYPES) for sample in ds.samples: required_columns.update({ f"{sample.name}_count", f"{sample.name}_cpm", f"max_{sample.name}_5pend" @@ -751,7 +785,7 @@ def select_regions(config: SelectionConfig) -> dict: _validate_scan_dataset(ds, windows_path, scan_metadata) source_region_counts, eligible_region_counts = _window_region_counts( windows_path, - config.region_types, + config, batch_size=config.batch_size, ) provenance = { @@ -763,10 +797,17 @@ def select_regions(config: SelectionConfig) -> dict: "source_windows": str(windows_path.resolve()), "window_scan": scan_metadata, "region_filter": { - "mode": "all" if config.region_types is None else "exact_region_type", - "allowed_region_types": ( + "mode": "all" if config.region_types is None else "overlap", + "requested_region_types": ( None if config.region_types is None else list(config.region_types) ), + "mixed_policy": ( + "only" + if config.only_mixed + else "discard" + if config.discard_mixed + else "include_matching" + ), "source_window_counts": source_region_counts, "eligible_window_counts": eligible_region_counts, }, diff --git a/tests/test_rbpnet.py b/tests/test_rbpnet.py index ae7f3dd..4084179 100644 --- a/tests/test_rbpnet.py +++ b/tests/test_rbpnet.py @@ -595,7 +595,7 @@ def test_selection_strategies_ids_serialization_and_stitching(tmp_path): assert any(row["source_window_count"] > 1 for row in load_selection_manifest(classified).rows) -def test_exact_region_type_filtering_provenance_and_peak_bh_universe(tmp_path): +def test_region_type_filtering_includes_matching_mixed_by_default(tmp_path): root = tmp_path / "processed" _write_processed_fixture(root) windows = tmp_path / "windows" @@ -616,35 +616,74 @@ def test_exact_region_type_filtering_provenance_and_peak_bh_universe(tmp_path): progress=False, )) broad_manifest = load_selection_manifest(broad) - assert [row["region_type"] for row in broad_manifest.rows] == ["5putr", "cds"] + assert [row["region_type"] for row in broad_manifest.rows] == [ + "5putr", "mixed", "cds", "mixed", "mixed" + ] assert broad_summary["region_filter"] == { - "mode": "exact_region_type", - "allowed_region_types": ["5putr", "cds"], + "mode": "overlap", + "requested_region_types": ["5putr", "cds"], + "mixed_policy": "include_matching", "source_window_counts": {"5putr": 1, "cds": 1, "mixed": 3}, - "eligible_window_counts": {"5putr": 1, "cds": 1}, + "eligible_window_counts": {"5putr": 1, "cds": 1, "mixed": 3}, } assert broad_summary["selected_example_region_counts"] == { "5putr": 1, "cds": 1, + "mixed": 3, } assert broad_manifest.metadata["configuration"]["region_types"] == [ "5putr", "cds", ] - mixed = tmp_path / "broad_mixed" - mixed_summary = select_regions(SelectionConfig( + matching_3putr = tmp_path / "broad_3putr" + matching_summary = select_regions(SelectionConfig( processed_dir=root, windows=windows, - output_prefix=mixed, + output_prefix=matching_3putr, strategy="broad_coverage", replicate_mode="combined", min_total_count=0, - region_types="mixed", + region_types="3putr", progress=False, )) - assert mixed_summary["n_examples"] == 3 - assert {row["region_type"] for row in load_selection_manifest(mixed).rows} == { + assert matching_summary["n_examples"] == 2 + assert {row["region_type"] for row in load_selection_manifest(matching_3putr).rows} == { + "mixed" + } + + pure = tmp_path / "broad_pure" + pure_summary = select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=pure, + strategy="broad_coverage", + replicate_mode="combined", + min_total_count=0, + region_types=("5putr", "cds"), + discard_mixed=True, + progress=False, + )) + assert [row["region_type"] for row in load_selection_manifest(pure).rows] == [ + "5putr", "cds" + ] + assert pure_summary["region_filter"]["mixed_policy"] == "discard" + + only_mixed = tmp_path / "broad_only_mixed" + only_mixed_summary = select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=only_mixed, + strategy="broad_coverage", + replicate_mode="combined", + min_total_count=0, + region_types=("3putr",), + only_mixed=True, + progress=False, + )) + assert only_mixed_summary["n_examples"] == 2 + assert only_mixed_summary["region_filter"]["mixed_policy"] == "only" + assert {row["region_type"] for row in load_selection_manifest(only_mixed).rows} == { "mixed" } @@ -660,6 +699,7 @@ def test_exact_region_type_filtering_provenance_and_peak_bh_universe(tmp_path): peak_min_log2_ratio=100.0, negative_max_log2_ratio=-100.0, region_types=("cds",), + discard_mixed=True, progress=False, )) classified_rows = load_selection_manifest(classified).rows @@ -674,13 +714,53 @@ def test_exact_region_type_filtering_provenance_and_peak_bh_universe(tmp_path): classified_rows[0]["source_min_depletion_pvalue"] ) - with pytest.raises(ValueError, match="unsupported region_types: promoter"): + classified_boundary = tmp_path / "classified_3putr_boundary" + select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=classified_boundary, + strategy="peak_gray_negative", + min_total_count=0, + peak_fdr=1.0, + negative_fdr=1.0, + peak_min_log2_ratio=100.0, + negative_max_log2_ratio=-100.0, + region_types=("3putr",), + only_mixed=True, + progress=False, + )) + boundary_rows = load_selection_manifest(classified_boundary).rows + assert boundary_rows + assert all(row["region_type"] == "mixed" for row in boundary_rows) + assert all(row["region_3putr_nt"] > 0 for row in boundary_rows) + + with pytest.raises(ValueError, match="unsupported region_types: mixed"): select_regions(SelectionConfig( processed_dir=root, windows=windows, output_prefix=tmp_path / "invalid_region", strategy="broad_coverage", - region_types=("promoter",), + region_types=("mixed",), + progress=False, + )) + with pytest.raises(ValueError, match="only_mixed requires"): + select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=tmp_path / "missing_regions", + strategy="broad_coverage", + only_mixed=True, + progress=False, + )) + with pytest.raises(ValueError, match="mutually exclusive"): + select_regions(SelectionConfig( + processed_dir=root, + windows=windows, + output_prefix=tmp_path / "conflicting_mixed", + strategy="broad_coverage", + region_types=("cds",), + discard_mixed=True, + only_mixed=True, progress=False, )) @@ -694,8 +774,11 @@ def test_select_regions_cli_parses_comma_separated_region_types(): "--output-prefix", "selected", "--strategy", "broad_coverage", "--region-types", "3putr,cds", + "--only-mixed", ]) assert args.region_types == ("3putr", "cds") + assert args.only_mixed is True + assert args.discard_mixed is False def test_zero_count_peak_negative_edges_and_broad_coverage_defaults(tmp_path): From ad12d3ca13f568b3d5efa4c4e5553bcf86f78e23 Mon Sep 17 00:00:00 2001 From: isvock Date: Wed, 12 Aug 2026 15:33:28 -0700 Subject: [PATCH 08/12] Add rbpnet evals --- docs/api.rst | 8 + docs/rbpnet.md | 106 +- src/transcriptml/cli/main.py | 53 +- .../plotting/rbpnet_evaluation.py | 476 ++++++++ src/transcriptml/rbpnet/__init__.py | 5 + src/transcriptml/rbpnet/evaluation.py | 1001 +++++++++++++++++ src/transcriptml/rbpnet/evaluation_metrics.py | 546 +++++++++ src/transcriptml/training/evaluation.py | 63 +- tests/test_cli_analysis.py | 20 + tests/test_rbpnet_evaluation.py | 183 +++ tests/test_rbpnet_model.py | 112 +- 11 files changed, 2558 insertions(+), 15 deletions(-) create mode 100644 src/transcriptml/plotting/rbpnet_evaluation.py create mode 100644 src/transcriptml/rbpnet/evaluation.py create mode 100644 src/transcriptml/rbpnet/evaluation_metrics.py create mode 100644 tests/test_rbpnet_evaluation.py diff --git a/docs/api.rst b/docs/api.rst index 6439632..9119704 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -78,6 +78,14 @@ RBPNet/eCLIP data :members: train_rbpnet_model, evaluate_rbpnet_model, write_rbpnet_predictions :member-order: bysource +.. automodule:: transcriptml.rbpnet.evaluation + :members: evaluate_rbpnet_report, resolve_rbpnet_checkpoint_indices + :member-order: bysource + +.. automodule:: transcriptml.rbpnet.evaluation_metrics + :members: profile_metrics, enrichment_metrics, replicate_ceiling_metrics, aggregate_observations, calibration_rows, select_representative_examples + :member-order: bysource + Models ------ diff --git a/docs/rbpnet.md b/docs/rbpnet.md index 5ee708e..f757875 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -468,7 +468,9 @@ transcriptml train configs/rbpnet/train_config.json transcriptml evaluate \ --checkpoint runs/rbpnet/model/best.pt \ --dataset data/rbpnet_chr21 \ - --out-csv runs/rbpnet/predictions.csv + --out-dir runs/rbpnet/evaluation \ + --split test \ + --save-profiles ``` `transcriptml models show rbpnet --json` prints every architectural default. @@ -628,6 +630,108 @@ provenance. Evaluation CSVs contain `pi`, optional `eta`, and each replicate's depth-adjusted predicted IP fraction. The raw structured tensors remain available through the Python model output for future attribution work. +### Scientific evaluation reports + +For an RBPNet checkpoint, `transcriptml evaluate` writes a structured report +when given `--out-dir`. The default split is `test`; `--split` accepts `train`, +`val`, `test`, or `all`. These indices are resolved **only** from `splits` stored +in the checkpoint used for training. RBPNet evaluation deliberately does not +fall back to `bundle.splits`, because a reused or edited bundle must not change +which observations are considered held out. Evaluation always uses jitter +shift zero and is deterministic apart from platform-level floating-point +details. + +```text +evaluation/ + summary.json + examples.parquet + stratified_metrics.parquet + calibration.parquet + plots/ + predicted_target_profiles.npy # only with --save-profiles + predicted_control_profiles.npy + predicted_ip_profiles.npy +``` + +The optional profile arrays are normalized positional probabilities (`float32`) +with shape `(N_evaluated, profile_length)`; each row sums to one over valid +positions. Their first axis is exactly `evaluation_row` in +`examples.parquet`; they can be opened without loading them into memory using +`np.load(path, mmap_mode="r")`. `predicted_ip_profiles.npy` stores the final +target/control mixture, not just the latent target component. + +For observed counts `y`, total `N`, empirical distribution `q=y/N`, predicted +distribution `p`, uniform distribution `u` over valid positions, and predicted +control distribution `p_control`, profile metrics use natural logarithms: + +| Metric | Definition | +| --- | --- | +| complete multinomial NLL | `-log Multinomial(y | N, p)`, including the count combinatorial constant | +| KL/read (saturated gap) | `(NLL_model - NLL_saturated) / N = KL(q || p)` | +| JSD | `0.5 KL(q || m) + 0.5 KL(p || m)`, where `m=(q+p)/2` | +| information gain over uniform/read | `(LL_model - LL_uniform) / N` | +| pooled-IP information gain over control/read | `(LL_predicted_IP - LL_predicted_control) / N` | +| Wasserstein | one-dimensional earth-mover distance between `q` and `p`, in nucleotides | + +Empirical-profile metrics are undefined and recorded as null/NaN when the +observed profile total is zero. Counts outside a validity mask are rejected; +all probability distributions are restricted and normalized over valid +positions. The report always exposes the scientifically comparable complete +NLL, while its checkpoint-objective reconstruction honors the checkpoint's +`include_multinomial_constant` and `include_binomial_constant` settings. + +With the enrichment head enabled, the report retains the complete +replicate-aware binomial NLL and compares it with the depth-only null +`eta=0`. Its information gain is +`(LL_model - LL_eta=0) / (IP+SMInput)`. The descriptive empirical enrichment is + +```text +eta_hat = log((IP + c) / (SMInput + c)) - log(L_IP / L_SM) +``` + +where `c=0.5` by default and is configurable with +`--enrichment-pseudocount`. This pseudocount is used only for Pearson/Spearman +diagnostics, never for either likelihood. `calibration.parquet` contains both +locus rows and fixed-width predicted-probability bins. Each bin uses +read-weighted summaries: + +```text +predicted_bin = sum(N * p) / sum(N) +observed_bin = sum(IP) / sum(N) +``` + +With multiple IP replicates, each observed replicate is also compared with the +pooled profile of all other replicates. The resulting leave-one-replicate-out +JSD and Wasserstein values are an experimental reproducibility reference, not +a guaranteed upper bound on every model metric. + +`stratified_metrics.parquet` is long-form. It reports locus macro (replicates +within a locus averaged first), gene macro (loci within a gene averaged first), +and read micro summaries. Read micro likelihood and information values sum the +appropriate likelihood/information numerators and divide by contributing +reads; JSD and Wasserstein use read-count-weighted means. Every row records its +informative observation, locus, gene, and read counts. Summaries include +pooled-IP and SMInput read-depth bins, chromosome, and optional +`selection_state` and `region_type` strata. + +`examples.parquet` is the inspectable per-locus table. It includes bundle and +evaluation indices, stable example and biological identifiers, selection +metadata where available, profile and selection-interval counts, effective +library sizes/depth offsets, `pi`, optional `eta` and replicate predicted IP +fractions, every per-example profile metric, and replicate-ceiling metrics. +`summary.json` records metric definitions, sample depths, aggregate metrics, +correlations, representative-example sampling, plot provenance, and all output +paths. + +The plot collection includes performance versus depth, metric distributions, +eta agreement, read-weighted calibration, fixed-seed representative profiles +(default: at least 10 pooled-IP and 10 SMInput profile reads), +selection/region stratification, and `pi` diagnostics. Missing optional +metadata, a disabled enrichment head, or a single IP replicate causes only the +inapplicable plot/metric to be skipped; `summary.json` records why. The legacy +`--out-csv` route remains available for compact `pi`/`eta` predictions, but the +structured report is preferred for scientific evaluation. + Profile-only model block: ```json diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index cdbde26..b084ff1 100644 --- a/src/transcriptml/cli/main.py +++ b/src/transcriptml/cli/main.py @@ -57,16 +57,30 @@ def _resolve_named_or_positional_args( def _resolve_evaluate_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> dict[str, str]: """Resolve evaluate paths from named flags or legacy positional arguments.""" - return _resolve_named_or_positional_args( + resolved = _resolve_named_or_positional_args( args, parser, command="evaluate", specs=[ ("checkpoint", "checkpoint_flag", "--checkpoint", "CHECKPOINT"), ("dataset", "dataset_flag", "--dataset", "DATASET"), - ("out_csv", "out_csv_flag", "--out-csv", "OUT_CSV"), ], ) + positional_csv = getattr(args, "out_csv", None) + flagged_csv = getattr(args, "out_csv_flag", None) + out_dir = getattr(args, "out_dir_flag", None) + if positional_csv is not None and flagged_csv is not None and str(positional_csv) != str(flagged_csv): + parser.error("evaluate got both --out-csv and positional OUT_CSV; use only one") + out_csv = flagged_csv if flagged_csv is not None else positional_csv + if out_csv is not None and out_dir is not None: + parser.error("evaluate accepts either --out-csv or --out-dir, not both") + if out_csv is None and out_dir is None: + parser.error("evaluate requires --out-dir, --out-csv, or legacy positional OUT_CSV") + if out_dir is not None: + resolved["out_dir"] = out_dir + else: + resolved["out_csv"] = out_csv + return resolved def _resolve_interpret_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> dict[str, str]: @@ -203,9 +217,30 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument("--checkpoint", dest="checkpoint_flag", help="Checkpoint path") p.add_argument("--dataset", dest="dataset_flag", help="Dataset bundle directory") p.add_argument("--out-csv", dest="out_csv_flag", help="Prediction CSV output path") - p.add_argument("--split") + p.add_argument( + "--out-dir", + dest="out_dir_flag", + help="Structured RBPNet evaluation report directory", + ) + p.add_argument( + "--split", + help=( + "Split to evaluate; RBPNet accepts train/val/test/all and defaults " + "to checkpoint-recorded test indices" + ), + ) p.add_argument("--batch-size", type=int, default=128) p.add_argument("--device", default="cpu") + p.add_argument( + "--save-profiles", + action="store_true", + help="Save RBPNet target/control/IP predicted profiles as memory-mappable .npy arrays", + ) + p.add_argument("--calibration-bins", type=int, default=10) + p.add_argument("--enrichment-pseudocount", type=float, default=0.5) + p.add_argument("--representative-seed", type=int, default=123) + p.add_argument("--representative-per-tier", type=int, default=3) + p.add_argument("--representative-min-profile-count", type=int, default=10) for name, help_text in [ ("ism", "Run single-nucleotide ISM"), @@ -615,11 +650,21 @@ def main(argv: list[str] | None = None) -> None: result = evaluate_checkpoint( evaluate_paths["checkpoint"], evaluate_paths["dataset"], - evaluate_paths["out_csv"], + evaluate_paths.get("out_csv"), + out_dir=evaluate_paths.get("out_dir"), split=args.split, batch_size=args.batch_size, device=args.device, + save_profiles=args.save_profiles, + calibration_bins=args.calibration_bins, + enrichment_pseudocount=args.enrichment_pseudocount, + representative_seed=args.representative_seed, + representative_per_tier=args.representative_per_tier, + representative_min_profile_count=args.representative_min_profile_count, ) + if "report_dir" in result: + log_progress(f"evaluate: wrote RBPNet report to {result['report_dir']}") + return non_summary_fields = { "predictions", "targets", "indices", "example_ids", "pi", "enrichment_logit", "depth_offsets", "replicate_names", diff --git a/src/transcriptml/plotting/rbpnet_evaluation.py b/src/transcriptml/plotting/rbpnet_evaluation.py new file mode 100644 index 0000000..3fa6a8c --- /dev/null +++ b/src/transcriptml/plotting/rbpnet_evaluation.py @@ -0,0 +1,476 @@ +"""Robust, deterministic plots for structured RBPNet evaluation reports.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import Mapping, Sequence + +import numpy as np +import torch + +if "MPLCONFIGDIR" not in os.environ: + cache = Path(os.environ.get("TMPDIR", "/tmp")) / "matplotlib-rbpnet-evaluation" + cache.mkdir(parents=True, exist_ok=True) + os.environ["MPLCONFIGDIR"] = str(cache) + +import matplotlib + +if ( + "matplotlib.pyplot" not in sys.modules + and "MPLBACKEND" not in os.environ + and not os.environ.get("DISPLAY") +): + matplotlib.use("Agg") + +import matplotlib.pyplot as plt + +from transcriptml.rbpnet.dataset import RBPNetDataset, collate_rbpnet +from transcriptml.rbpnet.evaluation_metrics import safe_correlations + + +_TRACK_LABELS = {"pooled_ip": "Pooled IP", "sminput": "SMInput"} +_METRIC_LABELS = { + "kl_per_read": "KL / read (nats)", + "jsd": "JSD (nats)", + "wasserstein_nt": "Wasserstein (nt)", + "information_gain_uniform_per_read": "Information gain / read (nats)", +} + + +def _values(rows: Sequence[Mapping[str, object]], column: str) -> np.ndarray: + return np.asarray([row.get(column, np.nan) for row in rows], dtype=np.float64) + + +def _finite_xy(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + keep = np.isfinite(x) & np.isfinite(y) + return x[keep], y[keep] + + +def _scatter_or_hexbin(ax, x: np.ndarray, y: np.ndarray, *, seed: int = 123) -> None: + x, y = _finite_xy(np.asarray(x), np.asarray(y)) + if x.size == 0: + ax.text(0.5, 0.5, "No informative observations", ha="center", va="center") + return + if x.size > 2_500: + image = ax.hexbin(x, y, gridsize=45, mincnt=1, bins="log", cmap="viridis") + ax.figure.colorbar(image, ax=ax, label="log10 bin count") + else: + rng = np.random.default_rng(seed) + order = rng.permutation(x.size) + ax.scatter(x[order], y[order], s=10, alpha=0.35, edgecolors="none") + + +def _save(fig, path: Path) -> None: + fig.tight_layout() + fig.savefig(path, dpi=160, bbox_inches="tight") + plt.close(fig) + + +def _plot_profile_depth(rows: Sequence[Mapping[str, object]], path: Path) -> None: + metrics = ("kl_per_read", "jsd", "wasserstein_nt") + fig, axes = plt.subplots(2, 3, figsize=(14, 8), squeeze=False) + for row_index, track in enumerate(("pooled_ip", "sminput")): + depth = _values(rows, f"{track}_profile_count") + x = np.log10(depth + 1.0) + for column_index, metric in enumerate(metrics): + ax = axes[row_index, column_index] + _scatter_or_hexbin(ax, x, _values(rows, f"{track}_{metric}")) + ax.set_xlabel("log10(observed profile reads + 1)") + ax.set_ylabel(_METRIC_LABELS[metric]) + ax.set_title(f"{_TRACK_LABELS[track]}: {_METRIC_LABELS[metric]}") + _save(fig, path) + + +def _plot_profile_distributions( + rows: Sequence[Mapping[str, object]], path: Path +) -> None: + metrics = ( + "kl_per_read", + "jsd", + "wasserstein_nt", + "information_gain_uniform_per_read", + ) + fig, axes = plt.subplots(2, 4, figsize=(16, 7), squeeze=False) + for row_index, track in enumerate(("pooled_ip", "sminput")): + for column_index, metric in enumerate(metrics): + ax = axes[row_index, column_index] + values = _values(rows, f"{track}_{metric}") + values = values[np.isfinite(values)] + if values.size: + ax.hist(values, bins=min(50, max(10, int(np.sqrt(values.size)))), alpha=0.7) + else: + ax.text(0.5, 0.5, "No informative observations", ha="center", va="center") + if track == "pooled_ip" and metric in {"jsd", "wasserstein_nt"}: + ceiling = _values(rows, f"ip_replicate_ceiling_{metric}") + ceiling = ceiling[np.isfinite(ceiling)] + if ceiling.size: + ax.hist( + ceiling, + bins=min(50, max(10, int(np.sqrt(ceiling.size)))), + histtype="step", + linewidth=1.8, + label="replicate ceiling", + ) + ax.legend(fontsize=8) + ax.set_xlabel(_METRIC_LABELS[metric]) + ax.set_ylabel("Loci") + ax.set_title(_TRACK_LABELS[track]) + _save(fig, path) + + +def _locus_calibration( + calibration: Sequence[Mapping[str, object]], replicate: str +) -> list[Mapping[str, object]]: + return [ + row + for row in calibration + if row.get("row_type") == "locus" and row.get("replicate") == replicate + ] + + +def _bin_calibration( + calibration: Sequence[Mapping[str, object]], replicate: str +) -> list[Mapping[str, object]]: + return sorted( + [ + row + for row in calibration + if row.get("row_type") == "bin" and row.get("replicate") == replicate + ], + key=lambda row: int(row.get("bin", 0)), + ) + + +def _plot_eta( + calibration: Sequence[Mapping[str, object]], + replicate_names: Sequence[str], + path: Path, +) -> bool: + panels = [name for name in replicate_names if _locus_calibration(calibration, name)] + if not panels: + return False + fig, axes = plt.subplots(1, len(panels), figsize=(5.5 * len(panels), 4.5), squeeze=False) + for ax, replicate in zip(axes[0], panels): + data = _locus_calibration(calibration, replicate) + eta = np.asarray([row["eta"] for row in data], dtype=float) + empirical = np.asarray([row["empirical_eta"] for row in data], dtype=float) + _scatter_or_hexbin(ax, eta, empirical) + correlations = safe_correlations(eta, empirical) + ax.text( + 0.02, + 0.98, + ( + f"Pearson={correlations['pearson']:.3g}\n" + f"Spearman={correlations['spearman']:.3g}\n" + f"n={correlations['n']}" + ), + transform=ax.transAxes, + ha="left", + va="top", + fontsize=9, + ) + ax.set_xlabel("Predicted eta") + ax.set_ylabel("Stabilized empirical eta") + ax.set_title(str(replicate)) + _save(fig, path) + return True + + +def _plot_calibration( + calibration: Sequence[Mapping[str, object]], + replicate_names: Sequence[str], + path: Path, +) -> bool: + panels = [name for name in replicate_names if _locus_calibration(calibration, name)] + if not panels: + return False + fig, axes = plt.subplots(1, len(panels), figsize=(5.5 * len(panels), 4.8), squeeze=False) + for ax, replicate in zip(axes[0], panels): + data = _locus_calibration(calibration, replicate) + predicted = np.asarray([row["predicted_probability"] for row in data], dtype=float) + observed = np.asarray([row["observed_fraction"] for row in data], dtype=float) + total = np.asarray([row["total_reads"] for row in data], dtype=float) + keep = np.isfinite(predicted) & np.isfinite(observed) & (total > 0) + if keep.sum() > 3_000: + image = ax.hexbin( + predicted[keep], + observed[keep], + C=np.log10(total[keep] + 1), + reduce_C_function=np.mean, + gridsize=42, + mincnt=1, + cmap="viridis", + ) + fig.colorbar(image, ax=ax, label="mean log10(reads + 1)") + else: + image = ax.scatter( + predicted[keep], + observed[keep], + c=np.log10(total[keep] + 1), + s=14, + alpha=0.45, + cmap="viridis", + edgecolors="none", + ) + if keep.any(): + fig.colorbar(image, ax=ax, label="log10(reads + 1)") + binned = _bin_calibration(calibration, replicate) + if binned: + ax.plot( + [float(row["predicted_bin"]) for row in binned], + [float(row["observed_bin"]) for row in binned], + marker="o", + color="#d62728", + linewidth=2, + label="read-weighted bins", + ) + ax.plot([0, 1], [0, 1], "--", color="black", linewidth=1, label="perfect") + ax.set_xlim(0, 1) + ax.set_ylim(0, 1) + ax.set_xlabel("Predicted IP fraction") + ax.set_ylabel("Observed IP / (IP + SMInput)") + ax.set_title(str(replicate)) + ax.legend(fontsize=8) + _save(fig, path) + return True + + +def _normalized(values: np.ndarray, mask: np.ndarray) -> np.ndarray: + result = np.asarray(values, dtype=np.float64).copy() + result[~mask] = np.nan + total = np.nansum(result) + if total > 0: + result /= total + return result + + +@torch.no_grad() +def _predict_example( + dataset: RBPNetDataset, + model: torch.nn.Module, + index: int, + device: torch.device, +) -> tuple[dict[str, object], dict[str, np.ndarray]]: + item = dataset.item_for_shift(index, 0) + batch = collate_rbpnet([item]).to(device) + output = model( + batch.sequence, + measurement_mask=batch.measurement_mask, + profile_mask=batch.profile_valid_mask, + ) + predicted = { + "target": output.target_probs[0].detach().cpu().numpy(), + "control": output.control_probs[0].detach().cpu().numpy(), + "ip": output.ip_probs[0].detach().cpu().numpy(), + } + return item, predicted + + +def _plot_representatives( + representative: Sequence[Mapping[str, object]], + dataset: RBPNetDataset, + model: torch.nn.Module, + device: torch.device, + path: Path, + saved_profile_paths: Mapping[str, Path] | None, +) -> bool: + if not representative: + return False + saved = ( + { + name: np.load(profile_path, mmap_mode="r", allow_pickle=False) + for name, profile_path in saved_profile_paths.items() + } + if saved_profile_paths is not None + else None + ) + fig, axes = plt.subplots( + len(representative), 2, figsize=(13, max(3.0, 2.6 * len(representative))), squeeze=False + ) + for row_index, selected in enumerate(representative): + original_index = int(selected["index"]) + item = dataset.item_for_shift(original_index, 0) + if saved is None: + _, predicted = _predict_example(dataset, model, original_index, device) + else: + evaluation_row = int(selected["evaluation_row"]) + predicted = { + name: np.asarray(array[evaluation_row]) for name, array in saved.items() + } + mask = np.asarray(item["profile_valid_mask"], dtype=bool) + x = np.arange(mask.size) + ip_ax, sm_ax = axes[row_index] + observed_ip = _normalized(np.asarray(item["pooled_ip_profile"]), mask) + observed_sm = _normalized(np.asarray(item["sminput_profile"]), mask) + ip_ax.plot(x, observed_ip, color="black", linewidth=1.2, label="observed pooled IP") + ip_ax.plot(x, _normalized(predicted["ip"], mask), label="predicted IP") + ip_ax.plot( + x, + _normalized(predicted["target"], mask), + linestyle=":", + label="latent target", + ) + sm_ax.plot(x, observed_sm, color="black", linewidth=1.2, label="observed SMInput") + sm_ax.plot(x, _normalized(predicted["control"], mask), label="predicted control") + title = ( + f"{selected['tier']}: {selected['example_id']} | " + f"IP KL/read={float(selected['metric_value']):.3g}" + ) + ip_ax.set_title(title, fontsize=9) + sm_ax.set_title("Control profile", fontsize=9) + for ax in (ip_ax, sm_ax): + ax.set_xlabel("Profile position (nt)") + ax.set_ylabel("Normalized probability") + ax.legend(fontsize=7, loc="upper right") + _save(fig, path) + return True + + +def _plot_stratified(rows: Sequence[Mapping[str, object]], path: Path) -> bool: + fields = [ + field + for field in ("selection_state", "region_type") + if any(row.get(field) not in {None, ""} for row in rows) + ] + if not fields: + return False + metrics = ("kl_per_read", "jsd", "wasserstein_nt") + fig, axes = plt.subplots( + len(fields), 3, figsize=(16, 4.2 * len(fields)), squeeze=False + ) + for row_index, field in enumerate(fields): + categories = sorted({str(row[field]) for row in rows if row.get(field) not in {None, ""}}) + positions = np.arange(len(categories), dtype=float) + for column_index, metric in enumerate(metrics): + ax = axes[row_index, column_index] + for track_index, track in enumerate(("pooled_ip", "sminput")): + means = [] + counts = [] + for category in categories: + values = np.asarray( + [ + row.get(f"{track}_{metric}", np.nan) + for row in rows + if str(row.get(field, "")) == category + ], + dtype=float, + ) + values = values[np.isfinite(values)] + means.append(float(values.mean()) if values.size else np.nan) + counts.append(int(values.size)) + offset = (-0.19, 0.19)[track_index] + bars = ax.bar( + positions + offset, + means, + width=0.36, + label=_TRACK_LABELS[track], + ) + for bar, count in zip(bars, counts): + if np.isfinite(bar.get_height()): + ax.text( + bar.get_x() + bar.get_width() / 2, + bar.get_height(), + f"{count}", + ha="center", + va="bottom", + fontsize=7, + ) + ax.set_xticks(positions, categories, rotation=30, ha="right") + ax.set_ylabel(f"Mean {_METRIC_LABELS[metric]}") + ax.set_title(f"{field}: {_METRIC_LABELS[metric]} (labels are n)") + ax.legend(fontsize=8) + _save(fig, path) + return True + + +def _plot_pi(rows: Sequence[Mapping[str, object]], path: Path) -> bool: + pi = _values(rows, "pi") + if not np.isfinite(pi).any(): + return False + info = _values(rows, "pooled_ip_information_gain_control_per_read") + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + finite_pi = pi[np.isfinite(pi)] + axes[0].hist(finite_pi, bins=min(50, max(10, int(np.sqrt(finite_pi.size))))) + axes[0].set_xlabel("pi") + axes[0].set_ylabel("Loci") + axes[0].set_title("Latent target mixture weight") + _scatter_or_hexbin(axes[1], pi, info) + axes[1].set_xlabel("pi") + axes[1].set_ylabel("IP information gain over control / read") + axes[1].set_title("Mixture diagnostic") + _save(fig, path) + return True + + +def create_rbpnet_evaluation_plots( + rows: Sequence[Mapping[str, object]], + calibration: Sequence[Mapping[str, object]], + plots_dir: str | Path, + *, + replicate_names: Sequence[str], + representative: Sequence[Mapping[str, object]], + dataset: RBPNetDataset, + model: torch.nn.Module, + device: torch.device, + saved_profile_paths: Mapping[str, Path] | None = None, +) -> dict[str, object]: + """Write the standard RBPNet diagnostic plot collection.""" + + out = Path(plots_dir) + out.mkdir(parents=True, exist_ok=True) + written: list[str] = [] + skipped: dict[str, str] = {} + + always = ( + ("profile_performance_vs_read_depth.png", _plot_profile_depth), + ("profile_metric_distributions.png", _plot_profile_distributions), + ) + for name, function in always: + function(rows, out / name) + written.append(name) + + optional = ( + ( + "eta_vs_empirical_enrichment.png", + lambda path: _plot_eta(calibration, replicate_names, path), + "checkpoint has no informative enrichment-head observations", + ), + ( + "enrichment_calibration.png", + lambda path: _plot_calibration(calibration, replicate_names, path), + "checkpoint has no informative enrichment-head observations", + ), + ( + "representative_profile_examples.png", + lambda path: _plot_representatives( + representative, + dataset, + model, + device, + path, + saved_profile_paths, + ), + "no examples have both informative IP and SMInput profiles", + ), + ( + "stratified_performance_summaries.png", + lambda path: _plot_stratified(rows, path), + "selection_state and region_type metadata are absent", + ), + ( + "pi_diagnostics.png", + lambda path: _plot_pi(rows, path), + "pi predictions are absent", + ), + ) + for name, function, reason in optional: + path = out / name + if function(path): + written.append(name) + else: + if path.is_file(): + path.unlink() + skipped[name] = reason + return {"written": written, "skipped": skipped} diff --git a/src/transcriptml/rbpnet/__init__.py b/src/transcriptml/rbpnet/__init__.py index 1bccbfa..1e9d863 100644 --- a/src/transcriptml/rbpnet/__init__.py +++ b/src/transcriptml/rbpnet/__init__.py @@ -12,6 +12,7 @@ "SelectionConfig", "WindowScanConfig", "make_rbpnet_bundle", + "evaluate_rbpnet_report", "evaluate_rbpnet_model", "preprocess_eclip", "scan_windows", @@ -64,4 +65,8 @@ def __getattr__(name: str): "train_rbpnet_model": train_rbpnet_model, "write_rbpnet_predictions": write_rbpnet_predictions, }[name] + if name == "evaluate_rbpnet_report": + from transcriptml.rbpnet.evaluation import evaluate_rbpnet_report + + return evaluate_rbpnet_report raise AttributeError(f"module 'transcriptml.rbpnet' has no attribute {name!r}") diff --git a/src/transcriptml/rbpnet/evaluation.py b/src/transcriptml/rbpnet/evaluation.py new file mode 100644 index 0000000..272d414 --- /dev/null +++ b/src/transcriptml/rbpnet/evaluation.py @@ -0,0 +1,1001 @@ +"""Structured, checkpoint-split-aware evaluation reports for RBPNet.""" + +from __future__ import annotations + +import json +import math +from collections import defaultdict +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import torch +from torch.utils.data import DataLoader, Subset + +from transcriptml.data.bundle import DatasetBundle +from transcriptml.devices import resolve_device +from transcriptml.models.rbpnet import RBPNet +from transcriptml.progress import ProgressReporter, log_progress +from transcriptml.rbpnet.dataset import RBPNetDataset, collate_rbpnet +from transcriptml.rbpnet.evaluation_metrics import ( + ENRICHMENT_METRIC_DEFINITIONS, + PROFILE_METRIC_DEFINITIONS, + aggregate_observations, + calibration_rows, + enrichment_metrics, + profile_metrics, + replicate_ceiling_metrics, + safe_correlations, + select_representative_examples, +) +from transcriptml.rbpnet.losses import RBPNetLossConfig + + +PROFILE_TRACKS = ("pooled_ip", "sminput") +PROFILE_METRICS = ( + "multinomial_nll", + "kl_per_read", + "jsd", + "information_gain_uniform_per_read", + "wasserstein_nt", +) + + +def resolve_rbpnet_checkpoint_indices( + checkpoint: Mapping[str, object], + *, + split: str | None, + n_examples: int, +) -> tuple[str, list[int]]: + """Resolve a named RBPNet split exclusively from checkpoint artifacts.""" + + requested = "test" if split is None else str(split).strip().lower() + if requested not in {"train", "val", "test", "all"}: + raise ValueError("RBPNet --split must be one of: train, val, test, all") + raw = checkpoint.get("splits") + if not isinstance(raw, Mapping): + raise ValueError( + "RBPNet checkpoint has no recorded training splits; evaluation will not " + "fall back to bundle.splits" + ) + normalized: dict[str, list[int]] = {} + owners: dict[int, str] = {} + for name in ("train", "val", "test"): + values = raw.get(name) + if values is None or isinstance(values, (str, bytes)): + raise ValueError(f"RBPNet checkpoint lacks recorded {name!r} indices") + indices = [int(value) for value in values] + if len(indices) != len(set(indices)): + raise ValueError(f"RBPNet checkpoint {name!r} split contains duplicates") + if any(index < 0 or index >= int(n_examples) for index in indices): + raise ValueError( + f"RBPNet checkpoint {name!r} split contains an out-of-range index" + ) + for index in indices: + previous = owners.setdefault(index, name) + if previous != name: + raise ValueError( + f"RBPNet checkpoint index {index} occurs in both {previous} and {name}" + ) + normalized[name] = indices + indices = ( + sorted(owners) + if requested == "all" + else list(normalized[requested]) + ) + if not indices: + raise ValueError(f"RBPNet checkpoint split {requested!r} is empty") + return requested, indices + + +def _loader( + dataset: RBPNetDataset, + indices: Sequence[int], + batch_size: int, +) -> DataLoader: + return DataLoader( + Subset(dataset, [int(index) for index in indices]), + batch_size=int(batch_size), + shuffle=False, + num_workers=0, + collate_fn=collate_rbpnet, + ) + + +def _mean_finite(values: np.ndarray) -> float: + array = np.asarray(values, dtype=np.float64) + finite = array[np.isfinite(array)] + return float(finite.mean()) if finite.size else float("nan") + + +def _weighted_row_mean( + values: np.ndarray, weights: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + valid = np.isfinite(values) & np.isfinite(weights) & (weights > 0) + numerators = np.where(valid, values * weights, 0.0).sum(axis=1) + denominators = np.where(valid, weights, 0.0).sum(axis=1) + means = np.divide( + numerators, + denominators, + out=np.full(denominators.shape, np.nan, dtype=np.float64), + where=denominators > 0, + ) + return means, numerators, denominators + + +def _metadata_value(row: Mapping[str, object], name: str) -> object | None: + value = row.get(name) + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, np.generic): + return value.item() + return str(value) + + +def _base_example_row( + bundle: DatasetBundle, + index: int, + evaluation_row: int, + split: str, +) -> dict[str, object]: + metadata = bundle.metadata[index] if bundle.metadata is not None else {} + row: dict[str, object] = { + "evaluation_row": int(evaluation_row), + "index": int(index), + "example_id": str(bundle.ids[index]), + "evaluated_split": split, + } + for name in ( + "gene_id", + "transcript_id", + "chromosome", + "strand", + "coordinate_space", + "locus_length", + "transcript_anchor", + "selection_start", + "selection_end", + "selection_state", + "selection_strategy", + "replicate_id", + "region_type", + "sequence_materialized_start", + "sequence_materialized_end", + "profile_materialized_start", + "profile_materialized_end", + "group_gene_id", + "group_transcript_id", + "group_chromosome", + ): + value = _metadata_value(metadata, name) + if value is not None: + row[name] = value + return row + + +def _profile_columns( + row: dict[str, object], + prefix: str, + metrics: Mapping[str, np.ndarray], + position: int, +) -> None: + row[f"{prefix}_profile_count"] = int(metrics["count"][position]) + row[f"{prefix}_valid_positions"] = int(metrics["valid_positions"][position]) + for name in ( + "multinomial_nll", + "multinomial_nll_without_constant", + "saturated_nll", + "uniform_nll", + "kl_per_read", + "jsd", + "information_gain_uniform_per_read", + "wasserstein_nt", + "information_gain_control_per_read", + ): + if name in metrics: + row[f"{prefix}_{name}"] = float(metrics[name][position]) + + +def _depth_bin(value: float) -> str: + count = int(value) + if count <= 0: + return "0" + for upper, label in ( + (2, "1-2"), + (5, "3-5"), + (10, "6-10"), + (20, "11-20"), + (50, "21-50"), + (100, "51-100"), + ): + if count <= upper: + return label + return ">100" + + +def _json_safe(value: object) -> object: + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if isinstance(value, np.ndarray): + return _json_safe(value.tolist()) + if isinstance(value, np.generic): + return _json_safe(value.item()) + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, Path): + return str(value) + return value + + +def _metric_unit(metric: str, aggregation: str) -> str: + if metric == "wasserstein_nt": + return "nt" + if metric == "multinomial_nll": + return "nats/read" if aggregation == "read_micro" else "nats/locus" + if metric == "binomial_nll": + return "nats/read" if aggregation == "read_micro" else "nats/observation" + return "nats/read" if "information_gain" in metric or "kl_per_read" in metric else "nats" + + +def _stratified_metric_rows( + examples: Sequence[Mapping[str, object]], + replicate_names: Sequence[str], +) -> list[dict[str, object]]: + """Create long-form overall and stratified metric summaries.""" + + rows = list(examples) + result: list[dict[str, object]] = [] + if not rows: + return result + locus_ids = np.asarray([row["example_id"] for row in rows], dtype=object) + gene_ids = np.asarray( + [row.get("gene_id", row.get("transcript_id", row["example_id"])) for row in rows], + dtype=object, + ) + + base_strata: list[tuple[str, str, np.ndarray]] = [ + ("overall", "all", np.ones(len(rows), dtype=bool)) + ] + for field, dimension in ( + ("selection_state", "selection_state"), + ("region_type", "region_type"), + ("chromosome", "chromosome"), + ): + values = np.asarray([row.get(field) for row in rows], dtype=object) + for value in sorted({str(value) for value in values if value not in {None, ""}}): + base_strata.append((dimension, value, values == value)) + + for track in PROFILE_TRACKS: + track_strata = list(base_strata) + depth = np.asarray([row[f"{track}_profile_count"] for row in rows], dtype=float) + depth_labels = np.asarray([_depth_bin(value) for value in depth], dtype=object) + for label in ("0", "1-2", "3-5", "6-10", "11-20", "21-50", "51-100", ">100"): + if np.any(depth_labels == label): + track_strata.append((f"{track}_read_depth", label, depth_labels == label)) + + metrics = list(PROFILE_METRICS) + if track == "pooled_ip": + metrics.append("information_gain_control_per_read") + for dimension, stratum, selected in track_strata: + for metric in metrics: + column = f"{track}_{metric}" + values = np.asarray([row.get(column, np.nan) for row in rows], dtype=float) + nll = metric == "multinomial_nll" + summary = aggregate_observations( + values[selected], + locus_ids=locus_ids[selected], + gene_ids=gene_ids[selected], + read_weights=depth[selected], + micro_numerators=values[selected] if nll else None, + ) + for aggregation, details in summary.items(): + result.append( + { + "dimension": dimension, + "stratum": stratum, + "track": track, + "metric": metric, + "aggregation": aggregation, + "value": details["value"], + "unit": _metric_unit(metric, aggregation), + "n_observations": details["n_observations"], + "n_loci": details["n_loci"], + "n_genes": details["n_genes"], + "total_reads": details["total_reads"], + } + ) + + if "enrichment_binomial_nll" in rows[0]: + n_replicates = len(replicate_names) + counts = np.asarray( + [ + [row.get(f"replicate_{name}_enrichment_count", 0) for name in replicate_names] + for row in rows + ], + dtype=float, + ).reshape(-1) + enrichment_loci = np.repeat(locus_ids, n_replicates) + enrichment_genes = np.repeat(gene_ids, n_replicates) + enrichment_strata = [ + ( + dimension, + stratum, + np.repeat(selected[:, None], n_replicates, axis=1).reshape(-1), + ) + for dimension, stratum, selected in base_strata + ] + depth_labels = np.asarray([_depth_bin(value) for value in counts], dtype=object) + for label in ("0", "1-2", "3-5", "6-10", "11-20", "21-50", "51-100", ">100"): + if np.any(depth_labels == label): + enrichment_strata.append( + ("enrichment_read_depth", label, depth_labels == label) + ) + for dimension, stratum, selected in enrichment_strata: + for metric in ( + "binomial_nll", + "information_gain_depth_null_per_read", + ): + values = np.asarray( + [ + row.get(f"replicate_{name}_{metric}", np.nan) + for row in rows + for name in replicate_names + ], + dtype=float, + ) + summary = aggregate_observations( + values[selected], + locus_ids=enrichment_loci[selected], + gene_ids=enrichment_genes[selected], + read_weights=counts[selected], + micro_numerators=values[selected] if metric == "binomial_nll" else None, + ) + for aggregation, details in summary.items(): + result.append( + { + "dimension": dimension, + "stratum": stratum, + "track": "enrichment", + "metric": metric, + "aggregation": aggregation, + "value": details["value"], + "unit": _metric_unit(metric, aggregation), + "n_observations": details["n_observations"], + "n_loci": details["n_loci"], + "n_genes": details["n_genes"], + "total_reads": details["total_reads"], + } + ) + + for metric in ("jsd", "wasserstein_nt"): + if not replicate_names or not any( + f"replicate_{name}_ceiling_{metric}" in rows[0] + for name in replicate_names + ): + continue + values = np.asarray( + [ + row.get(f"replicate_{name}_ceiling_{metric}", np.nan) + for row in rows + for name in replicate_names + ], + dtype=float, + ) + weights = np.asarray( + [ + row.get(f"replicate_{name}_profile_count", 0) + for row in rows + for name in replicate_names + ], + dtype=float, + ) + summary = aggregate_observations( + values, + locus_ids=np.repeat(locus_ids, len(replicate_names)), + gene_ids=np.repeat(gene_ids, len(replicate_names)), + read_weights=weights, + ) + for aggregation, details in summary.items(): + result.append( + { + "dimension": "overall", + "stratum": "all", + "track": "replicate_ceiling", + "metric": metric, + "aggregation": aggregation, + "value": details["value"], + "unit": _metric_unit(metric, aggregation), + "n_observations": details["n_observations"], + "n_loci": details["n_loci"], + "n_genes": details["n_genes"], + "total_reads": details["total_reads"], + } + ) + return result + + +def _calibration_schema() -> pa.Schema: + return pa.schema( + [ + pa.field("row_type", pa.string()), + pa.field("replicate", pa.string()), + pa.field("evaluation_row", pa.int64()), + pa.field("index", pa.int64()), + pa.field("example_id", pa.string()), + pa.field("eta", pa.float64()), + pa.field("empirical_eta", pa.float64()), + pa.field("predicted_probability", pa.float64()), + pa.field("observed_fraction", pa.float64()), + pa.field("ip_count", pa.int64()), + pa.field("sminput_count", pa.int64()), + pa.field("total_reads", pa.int64()), + pa.field("bin", pa.int64()), + pa.field("bin_left", pa.float64()), + pa.field("bin_right", pa.float64()), + pa.field("n_observations", pa.int64()), + pa.field("predicted_bin", pa.float64()), + pa.field("observed_bin", pa.float64()), + ] + ) + + +def _write_table(rows: Sequence[Mapping[str, object]], path: Path, schema: pa.Schema | None = None) -> None: + table = ( + pa.Table.from_pylist(list(rows), schema=schema) + if rows or schema is not None + else pa.table({}) + ) + pq.write_table(table, path, compression="zstd") + + +def evaluate_rbpnet_report( + model: RBPNet, + checkpoint: Mapping[str, object], + bundle: DatasetBundle, + output_dir: str | Path, + *, + split: str | None = None, + batch_size: int = 128, + device: str | torch.device = "cpu", + save_profiles: bool = False, + calibration_bins: int = 10, + enrichment_pseudocount: float = 0.5, + representative_seed: int = 123, + representative_per_tier: int = 3, + representative_min_profile_count: int = 10, + checkpoint_path: str | Path | None = None, + progress: bool = True, +) -> dict[str, object]: + """Create a complete deterministic RBPNet evaluation report directory.""" + + if bundle.config.get("bundle_format") != "transcriptml-rbpnet-bundle": + raise ValueError("RBPNet evaluation requires a TranscriptML RBPNet bundle") + if bundle.metadata is None: + raise ValueError("RBPNet evaluation requires bundle metadata") + if int(batch_size) <= 0: + raise ValueError("batch_size must be positive") + if int(calibration_bins) <= 0: + raise ValueError("calibration_bins must be positive") + if float(enrichment_pseudocount) <= 0: + raise ValueError("enrichment_pseudocount must be positive") + if int(representative_min_profile_count) < 0: + raise ValueError("representative_min_profile_count must be non-negative") + resolved_split, indices = resolve_rbpnet_checkpoint_indices( + checkpoint, + split=split, + n_examples=int(bundle.X.shape[0]), + ) + out = Path(output_dir) + out.mkdir(parents=True, exist_ok=True) + plots_dir = out / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + resolved_device = resolve_device(device) + model = model.to(resolved_device) + model.eval() + dataset = RBPNetDataset( + bundle, + max_train_jitter=0, + training=False, + require_full_measurement_interval=model.enrichment_enabled, + ) + loader = _loader(dataset, indices, batch_size) + n_evaluated = len(indices) + profile_length = dataset.crop_length + profile_memmaps: dict[str, np.memmap] = {} + if save_profiles: + for name in ("target", "control", "ip"): + profile_memmaps[name] = np.lib.format.open_memmap( + out / f"predicted_{name}_profiles.npy", + mode="w+", + dtype=np.float32, + shape=(n_evaluated, profile_length), + ) + else: + for name in ("target", "control", "ip"): + stale = out / f"predicted_{name}_profiles.npy" + if stale.is_file(): + stale.unlink() + + log_progress( + f"RBPNet evaluate: split={resolved_split}, examples={n_evaluated:,}, device={resolved_device}", + enabled=progress, + ) + rows: list[dict[str, object]] = [] + calibration_locus_rows: list[dict[str, object]] = [] + replicate_names = dataset.replicate_names + depth_offsets = dataset.depth_offsets.astype(np.float64) + enrichment_eta: list[float] = [] + enrichment_empirical: list[np.ndarray] = [] + enrichment_predicted: list[np.ndarray] = [] + enrichment_ip_counts: list[np.ndarray] = [] + enrichment_totals: list[np.ndarray] = [] + output_position = 0 + reporter = ProgressReporter( + "RBPNet evaluate", + total=len(loader), + unit="batches", + enabled=progress, + ) + with torch.no_grad(): + for batch in loader: + batch = batch.to(resolved_device) + output = model( + batch.sequence, + measurement_mask=batch.measurement_mask, + profile_mask=batch.profile_valid_mask, + ) + pooled_counts = batch.pooled_ip_profile.detach().cpu().numpy().astype(np.float64) + sm_counts = batch.sminput_profile.detach().cpu().numpy().astype(np.float64) + individual_counts = ( + batch.individual_ip_profiles.detach().cpu().numpy().astype(np.float64) + ) + valid = batch.profile_valid_mask.detach().cpu().numpy().astype(bool) + predicted_ip = output.ip_probs.detach().cpu().numpy().astype(np.float64) + predicted_control = output.control_probs.detach().cpu().numpy().astype(np.float64) + predicted_target = output.target_probs.detach().cpu().numpy().astype(np.float64) + ip_metrics = profile_metrics( + pooled_counts, + predicted_ip, + valid_mask=valid, + control_probabilities=predicted_control, + ) + sm_metrics = profile_metrics( + sm_counts, + predicted_control, + valid_mask=valid, + ) + ceiling = replicate_ceiling_metrics( + individual_counts, + valid_mask=valid, + ) + eta_values = ( + output.enrichment_logit.detach().cpu().numpy().astype(np.float64) + if output.enrichment_logit is not None + else None + ) + enrichment = None + selection_ip = batch.ip_measurement_counts.detach().cpu().numpy().astype(np.float64) + selection_sm = ( + batch.sminput_measurement_counts.detach().cpu().numpy().astype(np.float64) + ) + if eta_values is not None: + enrichment = enrichment_metrics( + eta_values, + selection_ip, + selection_sm, + depth_offsets, + pseudocount=enrichment_pseudocount, + ) + enrichment_eta.extend(float(value) for value in eta_values) + enrichment_empirical.append(enrichment["empirical_eta"]) + enrichment_predicted.append(enrichment["predicted_probability"]) + enrichment_ip_counts.append(selection_ip) + enrichment_totals.append(enrichment["count"]) + + batch_size_actual = pooled_counts.shape[0] + if save_profiles: + end = output_position + batch_size_actual + profile_memmaps["target"][output_position:end] = predicted_target.astype(np.float32) + profile_memmaps["control"][output_position:end] = predicted_control.astype(np.float32) + profile_memmaps["ip"][output_position:end] = predicted_ip.astype(np.float32) + + for local in range(batch_size_actual): + original_index = int(batch.indices[local].detach().cpu()) + row = _base_example_row( + bundle, + original_index, + output_position + local, + resolved_split, + ) + row["pi"] = float(output.pi[local].detach().cpu()) + row["zero_jitter_crop_start"] = int( + batch.crop_start[local].detach().cpu() + ) + row["zero_jitter_crop_end"] = ( + row["zero_jitter_crop_start"] + profile_length + ) + _profile_columns(row, "pooled_ip", ip_metrics, local) + _profile_columns(row, "sminput", sm_metrics, local) + row["selection_sminput_count"] = int(selection_sm[local]) + row["selection_ip_pooled_count"] = int(selection_ip[local].sum()) + row["sminput_effective_library_size"] = int( + dataset.sminput_library_size + ) + + if ceiling["jsd"].shape[1] > 0: + ceiling_jsd = ceiling["jsd"][local] + ceiling_wasserstein = ceiling["wasserstein_nt"][local] + ceiling_count = ceiling["count"][local] + jsd_mean, jsd_num, jsd_reads = _weighted_row_mean( + ceiling_jsd[None, :], ceiling_count[None, :] + ) + wass_mean, wass_num, _ = _weighted_row_mean( + ceiling_wasserstein[None, :], ceiling_count[None, :] + ) + row["ip_replicate_ceiling_jsd"] = _mean_finite(ceiling_jsd) + row["ip_replicate_ceiling_wasserstein_nt"] = _mean_finite( + ceiling_wasserstein + ) + row["ip_replicate_ceiling_count"] = float(jsd_reads[0]) + row["ip_replicate_ceiling_jsd_read_numerator"] = float(jsd_num[0]) + row["ip_replicate_ceiling_wasserstein_nt_read_numerator"] = float( + wass_num[0] + ) + for replicate, name in enumerate(replicate_names): + row[f"replicate_{name}_ceiling_jsd"] = float(ceiling_jsd[replicate]) + row[f"replicate_{name}_ceiling_wasserstein_nt"] = float( + ceiling_wasserstein[replicate] + ) + + for replicate, name in enumerate(replicate_names): + row[f"replicate_{name}_profile_count"] = int( + individual_counts[local, replicate].sum() + ) + row[f"replicate_{name}_selection_ip_count"] = int( + selection_ip[local, replicate] + ) + row[f"replicate_{name}_effective_library_size"] = int( + dataset.ip_library_sizes[replicate] + ) + row[f"replicate_{name}_depth_offset"] = float( + depth_offsets[replicate] + ) + + if enrichment is not None and eta_values is not None: + row["eta"] = float(eta_values[local]) + valid_enrichment = enrichment["count"][local] > 0 + nll_values = enrichment["binomial_nll"][local] + total_values = enrichment["count"][local] + info_values = enrichment[ + "information_gain_depth_null_per_read" + ][local] + row["enrichment_binomial_nll"] = _mean_finite(nll_values) + row["enrichment_binomial_nll_numerator"] = float( + np.nansum(nll_values) + ) + nll_without_constant = enrichment[ + "binomial_nll_without_constant" + ][local] + row["enrichment_binomial_nll_without_constant"] = _mean_finite( + nll_without_constant + ) + row[ + "enrichment_binomial_nll_without_constant_numerator" + ] = float(np.nansum(nll_without_constant)) + row["enrichment_n_observations"] = int( + np.count_nonzero(np.isfinite(nll_values)) + ) + info_mean, info_num, info_reads = _weighted_row_mean( + info_values[None, :], total_values[None, :] + ) + row["enrichment_information_gain_depth_null_per_read"] = float( + info_mean[0] + ) + row["enrichment_count"] = float(info_reads[0]) + row["enrichment_information_gain_depth_null_numerator"] = float( + info_num[0] + ) + for replicate, name in enumerate(replicate_names): + total = int(enrichment["count"][local, replicate]) + row[f"replicate_{name}_enrichment_count"] = total + row[f"replicate_{name}_predicted_ip_fraction"] = float( + enrichment["predicted_probability"][local, replicate] + ) + row[f"replicate_{name}_observed_ip_fraction"] = float( + enrichment["observed_fraction"][local, replicate] + ) + row[f"replicate_{name}_empirical_eta"] = float( + enrichment["empirical_eta"][local, replicate] + ) + row[f"replicate_{name}_binomial_nll"] = float( + enrichment["binomial_nll"][local, replicate] + ) + row[f"replicate_{name}_binomial_nll_without_constant"] = float( + enrichment["binomial_nll_without_constant"][ + local, replicate + ] + ) + row[f"replicate_{name}_depth_null_nll"] = float( + enrichment["depth_null_nll"][local, replicate] + ) + row[ + f"replicate_{name}_information_gain_depth_null_per_read" + ] = float(info_values[replicate]) + if valid_enrichment[replicate]: + calibration_locus_rows.append( + { + "row_type": "locus", + "replicate": name, + "evaluation_row": output_position + local, + "index": original_index, + "example_id": str(bundle.ids[original_index]), + "eta": float(eta_values[local]), + "empirical_eta": float( + enrichment["empirical_eta"][local, replicate] + ), + "predicted_probability": float( + enrichment["predicted_probability"][local, replicate] + ), + "observed_fraction": float( + enrichment["observed_fraction"][local, replicate] + ), + "ip_count": int(selection_ip[local, replicate]), + "sminput_count": int(selection_sm[local]), + "total_reads": total, + } + ) + rows.append(row) + output_position += batch_size_actual + reporter.update() + reporter.close() + for array in profile_memmaps.values(): + array.flush() + if output_position != n_evaluated: + raise RuntimeError("RBPNet evaluation did not emit every requested example") + + stratified = _stratified_metric_rows(rows, replicate_names) + calibration = list(calibration_locus_rows) + enrichment_correlations: dict[str, object] | None = None + if enrichment_predicted: + predicted = np.concatenate(enrichment_predicted, axis=0) + empirical = np.concatenate(enrichment_empirical, axis=0) + ip_counts = np.concatenate(enrichment_ip_counts, axis=0) + totals = np.concatenate(enrichment_totals, axis=0) + eta_array = np.asarray(enrichment_eta, dtype=np.float64) + binned = calibration_rows( + predicted, + ip_counts, + totals, + replicate_names, + n_bins=calibration_bins, + ) + calibration.extend({"row_type": "bin", **row} for row in binned) + enrichment_correlations = { + "overall_locus_replicate": safe_correlations( + np.repeat(eta_array[:, None], len(replicate_names), axis=1), + empirical, + ), + "by_replicate": { + name: safe_correlations(eta_array, empirical[:, replicate]) + for replicate, name in enumerate(replicate_names) + }, + "empirical_eta_pseudocount": float(enrichment_pseudocount), + } + + _write_table(rows, out / "examples.parquet") + _write_table(stratified, out / "stratified_metrics.parquet") + _write_table(calibration, out / "calibration.parquet", _calibration_schema()) + + ip_kl = np.asarray([row["pooled_ip_kl_per_read"] for row in rows], dtype=float) + representative = select_representative_examples( + ip_kl, + eligible=np.asarray( + [ + row["pooled_ip_profile_count"] >= representative_min_profile_count + and row["sminput_profile_count"] >= representative_min_profile_count + for row in rows + ], + dtype=bool, + ), + seed=representative_seed, + per_tier=representative_per_tier, + ) + for selection in representative: + evaluation_row = int(selection["index"]) + selection["evaluation_row"] = evaluation_row + selection["index"] = int(rows[evaluation_row]["index"]) + selection["example_id"] = str(rows[evaluation_row]["example_id"]) + + from transcriptml.plotting.rbpnet_evaluation import create_rbpnet_evaluation_plots + + plot_result = create_rbpnet_evaluation_plots( + rows, + calibration, + plots_dir, + replicate_names=replicate_names, + representative=representative, + dataset=dataset, + model=model, + device=resolved_device, + saved_profile_paths=( + { + name: out / f"predicted_{name}_profiles.npy" + for name in ("target", "control", "ip") + } + if save_profiles + else None + ), + ) + + overall_metrics = [ + row + for row in stratified + if row["dimension"] == "overall" and row["stratum"] == "all" + ] + loss_config = RBPNetLossConfig.from_config(checkpoint.get("loss_config")) + enrichment_observations = sum( + int(row.get("enrichment_n_observations", 0)) for row in rows + ) + enrichment_nll_sum = sum( + float(row.get("enrichment_binomial_nll_numerator", 0.0)) for row in rows + ) + enrichment_nll_without_constant_sum = sum( + float( + row.get("enrichment_binomial_nll_without_constant_numerator", 0.0) + ) + for row in rows + ) + ip_nll_column = ( + "pooled_ip_multinomial_nll" + if loss_config.include_multinomial_constant + else "pooled_ip_multinomial_nll_without_constant" + ) + sm_nll_column = ( + "sminput_multinomial_nll" + if loss_config.include_multinomial_constant + else "sminput_multinomial_nll_without_constant" + ) + ip_nll = np.asarray( + [row.get(ip_nll_column, np.nan) for row in rows], dtype=float + ) + sm_nll = np.asarray( + [row.get(sm_nll_column, np.nan) for row in rows], dtype=float + ) + objective_components = { + "pooled_ip_multinomial_nll": ( + float(ip_nll[np.isfinite(ip_nll)].mean()) + if np.isfinite(ip_nll).any() + else 0.0 + ), + "sminput_multinomial_nll": ( + float(sm_nll[np.isfinite(sm_nll)].mean()) + if np.isfinite(sm_nll).any() + else 0.0 + ), + "enrichment_binomial_nll": ( + ( + enrichment_nll_sum + if loss_config.include_binomial_constant + else enrichment_nll_without_constant_sum + ) + / enrichment_observations + if enrichment_observations > 0 + else 0.0 + ), + } + objective_loss = ( + loss_config.lambda_ip_profile + * objective_components["pooled_ip_multinomial_nll"] + + loss_config.lambda_sm_profile + * objective_components["sminput_multinomial_nll"] + + ( + loss_config.lambda_enrichment + * objective_components["enrichment_binomial_nll"] + if model.enrichment_enabled + else 0.0 + ) + ) + summary = { + "format": "transcriptml-rbpnet-evaluation", + "format_version": "1", + "checkpoint": str(checkpoint_path) if checkpoint_path is not None else None, + "evaluated_split": resolved_split, + "split_source": "checkpoint.splits", + "n_examples": n_evaluated, + "zero_jitter": True, + "replicate_names": list(replicate_names), + "sample_depths": { + "sminput": { + "name": dataset.sminput_name, + "effective_library_size": int(dataset.sminput_library_size), + }, + "ip": [ + { + "name": name, + "effective_library_size": int(dataset.ip_library_sizes[index]), + "log_library_size_ratio_vs_sminput": float(depth_offsets[index]), + } + for index, name in enumerate(replicate_names) + ], + }, + "enrichment_head_enabled": bool(model.enrichment_enabled), + "enrichment_empirical_pseudocount": float(enrichment_pseudocount), + "calibration_bins": int(calibration_bins), + "profile_metric_definitions": PROFILE_METRIC_DEFINITIONS, + "enrichment_metric_definitions": ENRICHMENT_METRIC_DEFINITIONS, + "aggregation_definitions": { + "locus_macro": "average replicate observations within locus, then loci equally", + "gene_macro": "average loci within gene, then genes equally", + "read_micro": "sum likelihood/information numerators divided by contributing reads", + }, + "overall_metrics": overall_metrics, + "training_objective": { + "loss_config": loss_config.to_dict( + enrichment_enabled=model.enrichment_enabled + ), + "component_columns": { + "pooled_ip": ip_nll_column, + "sminput": sm_nll_column, + "enrichment": ( + "enrichment_binomial_nll" + if loss_config.include_binomial_constant + else "enrichment_binomial_nll_without_constant" + ), + }, + "components": objective_components, + "weighted_loss": objective_loss, + }, + "enrichment_correlations": enrichment_correlations, + "representative_examples": representative, + "representative_sampling": { + "seed": int(representative_seed), + "per_performance_tertile": int(representative_per_tier), + "minimum_pooled_ip_profile_count": int(representative_min_profile_count), + "minimum_sminput_profile_count": int(representative_min_profile_count), + "ranking_metric": "pooled_ip_kl_per_read", + }, + "outputs": { + "examples": "examples.parquet", + "stratified_metrics": "stratified_metrics.parquet", + "calibration": "calibration.parquet", + "plots": "plots", + "predicted_profiles": ( + { + name: f"predicted_{name}_profiles.npy" + for name in ("target", "control", "ip") + } + if save_profiles + else None + ), + }, + "predicted_profile_contract": { + "saved": bool(save_profiles), + "dtype": "float32", + "shape": [n_evaluated, profile_length], + "axis_0": "evaluation_row in examples.parquet", + "axis_1": "zero-jitter model profile position", + "values": ( + "normalized positional probabilities over valid positions; " + "target is latent, control predicts SMInput, and ip is the " + "target/control mixture" + ), + }, + "plots": plot_result, + } + summary_path = out / "summary.json" + summary_path.write_text( + json.dumps(_json_safe(summary), indent=2) + "\n", + encoding="utf-8", + ) + log_progress(f"RBPNet evaluate: wrote report {out}", enabled=progress) + return { + "report_dir": str(out), + "summary": summary, + "indices": indices, + "example_ids": [str(bundle.ids[index]) for index in indices], + } diff --git a/src/transcriptml/rbpnet/evaluation_metrics.py b/src/transcriptml/rbpnet/evaluation_metrics.py new file mode 100644 index 0000000..95f4128 --- /dev/null +++ b/src/transcriptml/rbpnet/evaluation_metrics.py @@ -0,0 +1,546 @@ +"""Numerically explicit metrics for structured RBPNet evaluation.""" + +from __future__ import annotations + +import math +from collections import defaultdict +from typing import Mapping, Sequence + +import numpy as np +from scipy.special import expit, gammaln +from scipy.stats import pearsonr, spearmanr + + +PROFILE_METRIC_DEFINITIONS = { + "multinomial_nll": ( + "Complete multinomial negative log likelihood, including the count " + "combinatorial constant; lower is better." + ), + "multinomial_nll_without_constant": ( + "Multinomial cross-entropy term -sum(count * log probability), used " + "only to reconstruct checkpoints trained with the optional count " + "combinatorial constant disabled." + ), + "kl_per_read": ( + "(NLL_model - NLL_saturated) / observed profile reads, equal to " + "KL(empirical || model) in natural-log units; zero is ideal." + ), + "jsd": ( + "Jensen-Shannon divergence between empirical and predicted normalized " + "profiles in natural-log units; bounded by ln(2), zero is ideal." + ), + "information_gain_uniform_per_read": ( + "(LL_model - LL_uniform) / observed profile reads in natural-log units; " + "larger is better." + ), + "information_gain_control_per_read": ( + "For pooled IP only, (LL_predicted_IP - LL_predicted_control) / observed " + "IP reads in natural-log units; larger is better." + ), + "wasserstein_nt": ( + "One-dimensional earth-mover/Wasserstein-1 distance between empirical " + "and predicted normalized profiles, in nucleotide units; zero is ideal." + ), +} + +ENRICHMENT_METRIC_DEFINITIONS = { + "binomial_nll": ( + "Complete replicate-aware binomial negative log likelihood using " + "logit(p_ij)=eta_i+log(L_IP,j/L_SM)." + ), + "binomial_nll_without_constant": ( + "Binomial cross-entropy term without log(N choose IP), used only to " + "reconstruct checkpoints trained with that optional constant disabled." + ), + "information_gain_depth_null_per_read": ( + "(LL_model - LL_eta=0_depth_only) / (IP+SMInput) in natural-log units; " + "larger is better." + ), + "empirical_eta": ( + "log((IP+c)/(SMInput+c))-log(L_IP/L_SM), used only for descriptive " + "correlation and plotting, never as a likelihood target." + ), +} + + +def _as_2d(value: np.ndarray, name: str) -> np.ndarray: + result = np.asarray(value) + if result.ndim == 1: + result = result[None, :] + if result.ndim != 2: + raise ValueError(f"{name} must have shape (N, L) or (L,)") + return result + + +def profile_metrics( + counts: np.ndarray, + probabilities: np.ndarray, + *, + valid_mask: np.ndarray | None = None, + control_probabilities: np.ndarray | None = None, +) -> dict[str, np.ndarray]: + """Calculate complete-likelihood and normalized profile-shape metrics. + + Rows with zero observed counts retain their count and valid-position count + but receive ``NaN`` for empirical-profile metrics. Probabilities are + normalized over valid positions; counts outside the mask are rejected. + Natural logarithms are used throughout. + """ + + observed = _as_2d(np.asarray(counts, dtype=np.float64), "counts") + predicted = _as_2d( + np.asarray(probabilities, dtype=np.float64), "probabilities" + ) + if observed.shape != predicted.shape: + raise ValueError("counts and probabilities must have matching shapes") + if np.any(observed < 0) or not np.all(np.isfinite(observed)): + raise ValueError("profile counts must be finite and non-negative") + if np.any(predicted < 0) or not np.all(np.isfinite(predicted)): + raise ValueError("profile probabilities must be finite and non-negative") + if valid_mask is None: + valid = np.ones(observed.shape, dtype=bool) + else: + valid = _as_2d(np.asarray(valid_mask, dtype=bool), "valid_mask") + if valid.shape != observed.shape: + raise ValueError("valid_mask must match profile shape") + if np.any((~valid) & (observed != 0)): + raise ValueError("profile counts occur outside the validity mask") + if np.any(valid.sum(axis=1) == 0): + raise ValueError("every profile requires at least one valid position") + + control = None + if control_probabilities is not None: + control = _as_2d( + np.asarray(control_probabilities, dtype=np.float64), + "control_probabilities", + ) + if control.shape != observed.shape: + raise ValueError("control_probabilities must match profile shape") + if np.any(control < 0) or not np.all(np.isfinite(control)): + raise ValueError( + "control profile probabilities must be finite and non-negative" + ) + + n_rows = observed.shape[0] + names = ( + "multinomial_nll", + "multinomial_nll_without_constant", + "saturated_nll", + "uniform_nll", + "kl_per_read", + "jsd", + "information_gain_uniform_per_read", + "wasserstein_nt", + ) + result = { + "count": observed.sum(axis=1), + "valid_positions": valid.sum(axis=1).astype(np.int64), + **{name: np.full(n_rows, np.nan, dtype=np.float64) for name in names}, + } + if control is not None: + result["information_gain_control_per_read"] = np.full( + n_rows, np.nan, dtype=np.float64 + ) + + for index in range(n_rows): + mask = valid[index] + row_counts = observed[index, mask] + total = float(row_counts.sum()) + if total <= 0: + continue + model = predicted[index, mask] + model_sum = float(model.sum()) + if model_sum <= 0: + raise ValueError("predicted profile has zero mass on valid positions") + model = model / model_sum + empirical = row_counts / total + constant = float( + gammaln(total + 1.0) - np.sum(gammaln(row_counts + 1.0)) + ) + + positive = row_counts > 0 + with np.errstate(divide="ignore"): + model_log_likelihood = constant + float( + np.sum(row_counts[positive] * np.log(model[positive])) + ) + saturated_log_likelihood = constant + float( + np.sum(row_counts[positive] * np.log(empirical[positive])) + ) + uniform_log_likelihood = constant - total * math.log(len(row_counts)) + mixture = 0.5 * (empirical + model) + empirical_kl_mixture = float( + np.sum(empirical[positive] * np.log(empirical[positive] / mixture[positive])) + ) + model_positive = model > 0 + model_kl_mixture = float( + np.sum(model[model_positive] * np.log(model[model_positive] / mixture[model_positive])) + ) + + result["multinomial_nll"][index] = -model_log_likelihood + result["multinomial_nll_without_constant"][index] = -( + model_log_likelihood - constant + ) + result["saturated_nll"][index] = -saturated_log_likelihood + result["uniform_nll"][index] = -uniform_log_likelihood + result["kl_per_read"][index] = ( + saturated_log_likelihood - model_log_likelihood + ) / total + result["jsd"][index] = 0.5 * ( + empirical_kl_mixture + model_kl_mixture + ) + result["information_gain_uniform_per_read"][index] = ( + model_log_likelihood - uniform_log_likelihood + ) / total + result["wasserstein_nt"][index] = float( + np.abs(np.cumsum(empirical) - np.cumsum(model)).sum() + ) + + if control is not None: + control_row = control[index, mask] + control_sum = float(control_row.sum()) + if control_sum <= 0: + raise ValueError( + "control profile has zero mass on valid positions" + ) + control_row = control_row / control_sum + with np.errstate(divide="ignore"): + control_log_likelihood = constant + float( + np.sum(row_counts[positive] * np.log(control_row[positive])) + ) + result["information_gain_control_per_read"][index] = ( + model_log_likelihood - control_log_likelihood + ) / total + return result + + +def enrichment_metrics( + eta: np.ndarray, + ip_counts: np.ndarray, + sminput_counts: np.ndarray, + depth_offsets: np.ndarray, + *, + pseudocount: float = 0.5, +) -> dict[str, np.ndarray]: + """Calculate replicate-aware binomial metrics and descriptive enrichment.""" + + eta = np.asarray(eta, dtype=np.float64).reshape(-1) + ip = np.asarray(ip_counts, dtype=np.float64) + if ip.ndim == 1: + ip = ip[:, None] + sm = np.asarray(sminput_counts, dtype=np.float64).reshape(-1) + offsets = np.asarray(depth_offsets, dtype=np.float64) + if ip.ndim != 2 or ip.shape[0] != eta.size or sm.shape != eta.shape: + raise ValueError("eta, IP counts, and SMInput counts are not aligned") + if offsets.ndim == 1: + if offsets.shape[0] != ip.shape[1]: + raise ValueError("depth_offsets must have one value per IP replicate") + offsets = np.broadcast_to(offsets[None, :], ip.shape) + elif offsets.shape != ip.shape: + raise ValueError("depth_offsets must have shape (R,) or (N, R)") + if pseudocount <= 0: + raise ValueError("enrichment pseudocount must be positive") + if np.any(ip < 0) or np.any(sm < 0): + raise ValueError("enrichment counts must be non-negative") + + failures = np.broadcast_to(sm[:, None], ip.shape) + total = ip + failures + logits = eta[:, None] + offsets + predicted = expit(logits) + depth_only = expit(offsets) + log_choose = gammaln(total + 1.0) - gammaln(ip + 1.0) - gammaln( + failures + 1.0 + ) + + def nll_for_logits(value: np.ndarray) -> np.ndarray: + return total * np.logaddexp(0.0, value) - ip * value - log_choose + + model_nll_without_constant = ( + total * np.logaddexp(0.0, logits) - ip * logits + ) + model_nll = nll_for_logits(logits) + null_nll = nll_for_logits(offsets) + informative = total > 0 + model_nll = np.where(informative, model_nll, np.nan) + model_nll_without_constant = np.where( + informative, model_nll_without_constant, np.nan + ) + null_nll = np.where(informative, null_nll, np.nan) + information_gain = np.divide( + null_nll - model_nll, + total, + out=np.full(total.shape, np.nan, dtype=np.float64), + where=informative, + ) + observed_fraction = np.divide( + ip, + total, + out=np.full(ip.shape, np.nan, dtype=np.float64), + where=informative, + ) + empirical_eta = np.where( + informative, + np.log((ip + pseudocount) / (failures + pseudocount)) - offsets, + np.nan, + ) + return { + "count": total, + "predicted_probability": predicted, + "observed_fraction": observed_fraction, + "binomial_nll": model_nll, + "binomial_nll_without_constant": model_nll_without_constant, + "depth_null_nll": null_nll, + "information_gain_depth_null_per_read": information_gain, + "empirical_eta": empirical_eta, + } + + +def replicate_ceiling_metrics( + replicate_profiles: np.ndarray, + *, + valid_mask: np.ndarray | None = None, +) -> dict[str, np.ndarray]: + """Compare each IP replicate with the pooled profile of all other replicates.""" + + profiles = np.asarray(replicate_profiles, dtype=np.float64) + if profiles.ndim != 3: + raise ValueError("replicate_profiles must have shape (N, R, L)") + n_examples, n_replicates, length = profiles.shape + if valid_mask is None: + valid = np.ones((n_examples, length), dtype=bool) + else: + valid = _as_2d(np.asarray(valid_mask, dtype=bool), "valid_mask") + if valid.shape != (n_examples, length): + raise ValueError("valid_mask must have shape (N, L)") + if n_replicates < 2: + empty = np.empty((n_examples, 0), dtype=np.float64) + return {"count": empty, "jsd": empty, "wasserstein_nt": empty} + + counts = profiles.sum(axis=2) + jsd = np.full((n_examples, n_replicates), np.nan, dtype=np.float64) + wasserstein = np.full_like(jsd, np.nan) + pooled = profiles.sum(axis=1) + for replicate in range(n_replicates): + observed = profiles[:, replicate, :] + leave_one_out = pooled - observed + leave_one_out_counts = leave_one_out.sum(axis=1) + informative = (counts[:, replicate] > 0) & (leave_one_out_counts > 0) + if np.any(informative): + metrics = profile_metrics( + observed[informative], + leave_one_out[informative], + valid_mask=valid[informative], + ) + jsd[informative, replicate] = metrics["jsd"] + wasserstein[informative, replicate] = metrics["wasserstein_nt"] + return {"count": counts, "jsd": jsd, "wasserstein_nt": wasserstein} + + +def safe_correlations(x: np.ndarray, y: np.ndarray) -> dict[str, float | int]: + """Return finite-pair Pearson/Spearman correlations without warnings.""" + + left = np.asarray(x, dtype=np.float64).reshape(-1) + right = np.asarray(y, dtype=np.float64).reshape(-1) + keep = np.isfinite(left) & np.isfinite(right) + left = left[keep] + right = right[keep] + result: dict[str, float | int] = {"n": int(left.size)} + if left.size < 2 or np.ptp(left) == 0 or np.ptp(right) == 0: + result.update({"pearson": float("nan"), "spearman": float("nan")}) + return result + result.update( + { + "pearson": float(pearsonr(left, right).statistic), + "spearman": float(spearmanr(left, right).statistic), + } + ) + return result + + +def calibration_rows( + predicted_probability: np.ndarray, + ip_counts: np.ndarray, + total_counts: np.ndarray, + replicate_names: Sequence[str], + *, + n_bins: int = 10, +) -> list[dict[str, object]]: + """Build read-weighted fixed-width predicted-probability calibration bins.""" + + predicted = np.asarray(predicted_probability, dtype=np.float64) + successes = np.asarray(ip_counts, dtype=np.float64) + totals = np.asarray(total_counts, dtype=np.float64) + if predicted.ndim == 1: + predicted = predicted[:, None] + if predicted.shape != successes.shape or predicted.shape != totals.shape: + raise ValueError("calibration arrays must have matching (N, R) shapes") + if predicted.shape[1] != len(replicate_names): + raise ValueError("replicate_names does not match calibration arrays") + if n_bins <= 0: + raise ValueError("n_bins must be positive") + rows: list[dict[str, object]] = [] + for replicate, name in enumerate(replicate_names): + values = predicted[:, replicate] + valid = ( + np.isfinite(values) + & (values >= 0) + & (values <= 1) + & (totals[:, replicate] > 0) + ) + indices = np.zeros(values.shape, dtype=np.int64) + indices[valid] = np.minimum( + (values[valid] * n_bins).astype(np.int64), n_bins - 1 + ) + for bin_index in range(n_bins): + selected = valid & (indices == bin_index) + if not np.any(selected): + continue + weight = float(totals[selected, replicate].sum()) + rows.append( + { + "replicate": str(name), + "bin": bin_index, + "bin_left": bin_index / n_bins, + "bin_right": (bin_index + 1) / n_bins, + "n_observations": int(selected.sum()), + "total_reads": int(round(weight)), + "predicted_bin": float( + np.sum(totals[selected, replicate] * values[selected]) + / weight + ), + "observed_bin": float( + successes[selected, replicate].sum() / weight + ), + } + ) + return rows + + +def aggregate_observations( + values: np.ndarray, + *, + locus_ids: Sequence[object], + gene_ids: Sequence[object], + read_weights: np.ndarray | None = None, + micro_numerators: np.ndarray | None = None, +) -> dict[str, dict[str, float | int]]: + """Calculate locus-macro, gene-macro, and read-micro summaries. + + Replicate observations sharing a locus are averaged before locus- and + gene-macro aggregation. Read micro is a read-weighted mean, or equivalently + ``sum(micro_numerators)/sum(read_weights)`` when explicit numerators are + supplied (for example complete NLL rather than NLL times read count). + """ + + metric = np.asarray(values, dtype=np.float64).reshape(-1) + loci = np.asarray(locus_ids, dtype=object).reshape(-1) + genes = np.asarray(gene_ids, dtype=object).reshape(-1) + if metric.shape != loci.shape or metric.shape != genes.shape: + raise ValueError("values, locus_ids, and gene_ids must be aligned") + weights = ( + np.ones(metric.shape, dtype=np.float64) + if read_weights is None + else np.asarray(read_weights, dtype=np.float64).reshape(-1) + ) + if weights.shape != metric.shape: + raise ValueError("read_weights must align with values") + numerators = ( + metric * weights + if micro_numerators is None + else np.asarray(micro_numerators, dtype=np.float64).reshape(-1) + ) + if numerators.shape != metric.shape: + raise ValueError("micro_numerators must align with values") + keep = ( + np.isfinite(metric) + & np.isfinite(weights) + & np.isfinite(numerators) + & (weights > 0) + ) + metric = metric[keep] + weights = weights[keep] + numerators = numerators[keep] + loci = loci[keep] + genes = genes[keep] + if metric.size == 0: + empty = { + "value": float("nan"), + "n_observations": 0, + "n_loci": 0, + "n_genes": 0, + "total_reads": 0.0, + } + return { + "locus_macro": dict(empty), + "gene_macro": dict(empty), + "read_micro": dict(empty), + } + + locus_values: dict[object, list[float]] = defaultdict(list) + locus_gene: dict[object, object] = {} + for value, locus, gene in zip(metric, loci, genes): + locus_values[locus].append(float(value)) + locus_gene.setdefault(locus, gene) + locus_means = {key: float(np.mean(value)) for key, value in locus_values.items()} + gene_values: dict[object, list[float]] = defaultdict(list) + for locus, value in locus_means.items(): + gene_values[locus_gene[locus]].append(value) + gene_means = [float(np.mean(value)) for value in gene_values.values()] + common = { + "n_observations": int(metric.size), + "n_loci": len(locus_means), + "n_genes": len(gene_values), + "total_reads": float(weights.sum()), + } + return { + "locus_macro": { + **common, + "value": float(np.mean(list(locus_means.values()))), + }, + "gene_macro": {**common, "value": float(np.mean(gene_means))}, + "read_micro": { + **common, + "value": float(numerators.sum() / weights.sum()), + }, + } + + +def select_representative_examples( + metric: np.ndarray, + *, + eligible: np.ndarray | None = None, + seed: int = 123, + per_tier: int = 3, +) -> list[dict[str, object]]: + """Reproducibly sample good/middle/poor examples from rank tertiles.""" + + values = np.asarray(metric, dtype=np.float64).reshape(-1) + allowed = np.isfinite(values) + if eligible is not None: + requested = np.asarray(eligible, dtype=bool).reshape(-1) + if requested.shape != values.shape: + raise ValueError("eligible must align with metric") + allowed &= requested + if per_tier < 0: + raise ValueError("per_tier must be non-negative") + ordered = np.flatnonzero(allowed) + ordered = ordered[np.lexsort((ordered, values[ordered]))] + tiers = np.array_split(ordered, 3) + labels = ("good", "intermediate", "poor") + rng = np.random.default_rng(int(seed)) + selected: list[dict[str, object]] = [] + for tier_index, (label, candidates) in enumerate(zip(labels, tiers)): + size = min(int(per_tier), len(candidates)) + if size == 0: + continue + chosen = np.sort(rng.choice(candidates, size=size, replace=False)) + for index in chosen: + selected.append( + { + "index": int(index), + "tier": label, + "quantile_left": tier_index / 3, + "quantile_right": (tier_index + 1) / 3, + "metric_value": float(values[index]), + } + ) + return selected diff --git a/src/transcriptml/training/evaluation.py b/src/transcriptml/training/evaluation.py index 4fb70f9..5ec0f7d 100644 --- a/src/transcriptml/training/evaluation.py +++ b/src/transcriptml/training/evaluation.py @@ -479,9 +479,16 @@ def evaluate_checkpoint( dataset_path: str | Path, out_csv: str | Path | None = None, *, + out_dir: str | Path | None = None, split: str | None = None, batch_size: int = 128, device: str | torch.device = "cpu", + save_profiles: bool = False, + calibration_bins: int = 10, + enrichment_pseudocount: float = 0.5, + representative_seed: int = 123, + representative_per_tier: int = 3, + representative_min_profile_count: int = 10, progress: bool = True, ) -> dict[str, object]: """Load a checkpoint and evaluate it on a dataset bundle. @@ -489,25 +496,35 @@ def evaluate_checkpoint( Args: checkpoint_path: TranscriptML checkpoint path to load. dataset_path: Processed dataset bundle directory. - out_csv: Optional destination CSV path for predictions. - split: Optional named split from the dataset bundle to evaluate. + out_csv: Optional legacy destination CSV path for predictions. + out_dir: Structured report directory for RBPNet checkpoints. + split: Named split to evaluate. RBPNet resolves this exclusively from + checkpoint artifacts and defaults to ``test``; scalar models retain + the existing dataset-bundle behavior. batch_size: Number of examples to score per prediction batch. device: Torch device used for model execution. progress: Whether to emit progress messages while evaluating. """ + if out_csv is not None and out_dir is not None: + raise ValueError("provide either out_csv or out_dir, not both") device = resolve_device(device) log_progress(f"evaluate: loading checkpoint {checkpoint_path}", enabled=progress) model, checkpoint = load_checkpoint(checkpoint_path, map_location=device) log_progress(f"evaluate: loading dataset {dataset_path}", enabled=progress) bundle = load_bundle(dataset_path, mmap_mode="r") - indices = None - if split is not None: - if not bundle.splits or split not in bundle.splits: - raise ValueError(f"Dataset has no split '{split}'") - indices = [int(i) for i in bundle.splits[split]] if checkpoint.get("model_config", {}).get("name") == "rbpnet": from transcriptml.models.rbpnet import RBPNet + try: + from transcriptml.rbpnet.evaluation import ( + evaluate_rbpnet_report, + resolve_rbpnet_checkpoint_indices, + ) + except ImportError as exc: + raise ImportError( + "RBPNet evaluation requires optional dependencies; install " + "TranscriptML[rbpnet]" + ) from exc from transcriptml.rbpnet.training import ( evaluate_rbpnet_model, write_rbpnet_predictions, @@ -515,6 +532,29 @@ def evaluate_checkpoint( if not isinstance(model, RBPNet): raise TypeError("rbpnet checkpoint did not reconstruct an RBPNet model") + if out_dir is not None: + return evaluate_rbpnet_report( + model, + checkpoint, + bundle, + out_dir, + split=split, + batch_size=batch_size, + device=device, + save_profiles=save_profiles, + calibration_bins=calibration_bins, + enrichment_pseudocount=enrichment_pseudocount, + representative_seed=representative_seed, + representative_per_tier=representative_per_tier, + representative_min_profile_count=representative_min_profile_count, + checkpoint_path=checkpoint_path, + progress=progress, + ) + _, indices = resolve_rbpnet_checkpoint_indices( + checkpoint, + split=split, + n_examples=int(bundle.X.shape[0]), + ) result = evaluate_rbpnet_model( model, bundle, @@ -536,6 +576,15 @@ def evaluate_checkpoint( ) result["targets"] = None return result + if out_dir is not None: + raise ValueError("--out-dir structured reports are currently specific to RBPNet checkpoints") + if save_profiles: + raise ValueError("--save-profiles is only supported for RBPNet checkpoints") + indices = None + if split is not None: + if not bundle.splits or split not in bundle.splits: + raise ValueError(f"Dataset has no split '{split}'") + indices = [int(i) for i in bundle.splits[split]] log_progress( f"evaluate: running on {len(indices) if indices is not None else bundle.X.shape[0]} examples", enabled=progress, diff --git a/tests/test_cli_analysis.py b/tests/test_cli_analysis.py index e116536..00c96b1 100644 --- a/tests/test_cli_analysis.py +++ b/tests/test_cli_analysis.py @@ -92,6 +92,26 @@ def test_evaluate_cli_resolves_named_and_legacy_positional_args(): mixed = parser.parse_args(["evaluate", "model/best.pt", "data/saluki", "--out-csv", "eval/predictions.csv"]) assert _resolve_evaluate_args(mixed, parser)["out_csv"] == "eval/predictions.csv" + report = parser.parse_args( + [ + "evaluate", + "--checkpoint", + "model/best.pt", + "--dataset", + "data/rbpnet", + "--out-dir", + "eval/report", + "--save-profiles", + ] + ) + assert _resolve_evaluate_args(report, parser) == { + "checkpoint": "model/best.pt", + "dataset": "data/rbpnet", + "out_dir": "eval/report", + } + assert report.split is None + assert report.save_profiles is True + def test_evaluate_cli_rejects_conflicting_named_and_positional_args(): parser = build_parser() diff --git a/tests/test_rbpnet_evaluation.py b/tests/test_rbpnet_evaluation.py new file mode 100644 index 0000000..40d32d0 --- /dev/null +++ b/tests/test_rbpnet_evaluation.py @@ -0,0 +1,183 @@ +import math + +import numpy as np +import pytest + +from transcriptml.rbpnet.evaluation import resolve_rbpnet_checkpoint_indices +from transcriptml.rbpnet.evaluation_metrics import ( + aggregate_observations, + calibration_rows, + enrichment_metrics, + profile_metrics, + replicate_ceiling_metrics, + select_representative_examples, +) + + +def test_profile_metrics_hand_computable_uniform_control_and_saturated(): + result = profile_metrics( + np.asarray([[1, 0], [1, 1]], dtype=float), + np.asarray([[0.5, 0.5], [0.5, 0.5]], dtype=float), + control_probabilities=np.asarray([[1.0, 0.0], [0.5, 0.5]]), + ) + assert result["multinomial_nll"][0] == pytest.approx(math.log(2)) + assert result["saturated_nll"][0] == pytest.approx(0.0) + assert result["uniform_nll"][0] == pytest.approx(math.log(2)) + assert result["kl_per_read"][0] == pytest.approx(math.log(2)) + assert result["jsd"][0] == pytest.approx(0.75 * math.log(4 / 3)) + assert result["information_gain_uniform_per_read"][0] == pytest.approx(0.0) + assert result["information_gain_control_per_read"][0] == pytest.approx( + -math.log(2) + ) + assert result["wasserstein_nt"][0] == pytest.approx(0.5) + + # The empirical [0.5, 0.5] profile is saturated by the model. Its complete + # multinomial NLL is still ln(2), the negative log probability of counts + # [1,1] under n=2 and p=[0.5,0.5]. + assert result["multinomial_nll"][1] == pytest.approx(math.log(2)) + assert result["multinomial_nll_without_constant"][1] == pytest.approx( + 2 * math.log(2) + ) + assert result["saturated_nll"][1] == pytest.approx(math.log(2)) + assert result["kl_per_read"][1] == pytest.approx(0.0) + + +def test_profile_metrics_validity_mask_and_zero_count_behavior(): + result = profile_metrics( + np.asarray([[0, 1, 0], [0, 0, 0]]), + np.asarray([[0.25, 0.25, 0.5], [0.1, 0.4, 0.5]]), + valid_mask=np.asarray([[False, True, True], [False, True, True]]), + ) + assert result["valid_positions"].tolist() == [2, 2] + assert result["multinomial_nll"][0] == pytest.approx(math.log(3)) + assert np.isnan(result["kl_per_read"][1]) + assert np.isnan(result["jsd"][1]) + with pytest.raises(ValueError, match="outside the validity mask"): + profile_metrics( + np.asarray([[1, 0]]), + np.asarray([[0.5, 0.5]]), + valid_mask=np.asarray([[False, True]]), + ) + + +def test_enrichment_depth_null_and_stabilized_empirical_eta(): + result = enrichment_metrics( + eta=np.asarray([0.0, math.log(2)]), + ip_counts=np.asarray([[1], [2]]), + sminput_counts=np.asarray([1, 0]), + depth_offsets=np.asarray([0.0]), + pseudocount=0.5, + ) + assert result["information_gain_depth_null_per_read"][0, 0] == pytest.approx(0.0) + assert result["binomial_nll_without_constant"][0, 0] == pytest.approx( + 2 * math.log(2) + ) + assert result["predicted_probability"][1, 0] == pytest.approx(2 / 3) + assert result["binomial_nll"][1, 0] == pytest.approx(-2 * math.log(2 / 3)) + assert result["depth_null_nll"][1, 0] == pytest.approx(2 * math.log(2)) + assert result["information_gain_depth_null_per_read"][1, 0] == pytest.approx( + math.log(4 / 3) + ) + assert result["empirical_eta"][1, 0] == pytest.approx(math.log(5)) + empty = enrichment_metrics( + eta=np.asarray([2.0]), + ip_counts=np.asarray([[0]]), + sminput_counts=np.asarray([0]), + depth_offsets=np.asarray([0.0]), + ) + assert np.isnan(empty["binomial_nll"][0, 0]) + assert np.isnan(empty["information_gain_depth_null_per_read"][0, 0]) + assert np.isnan(empty["empirical_eta"][0, 0]) + + +def test_jsd_wasserstein_and_replicate_ceiling(): + ceiling = replicate_ceiling_metrics( + np.asarray( + [ + [[1, 0, 0], [1, 0, 0]], + [[1, 0, 0], [0, 1, 0]], + ] + ) + ) + np.testing.assert_allclose(ceiling["jsd"][0], 0.0) + np.testing.assert_allclose(ceiling["wasserstein_nt"][0], 0.0) + np.testing.assert_allclose(ceiling["jsd"][1], math.log(2)) + np.testing.assert_allclose(ceiling["wasserstein_nt"][1], 1.0) + single = replicate_ceiling_metrics(np.ones((2, 1, 3))) + assert single["jsd"].shape == (2, 0) + one_empty = replicate_ceiling_metrics( + np.asarray([[[1, 0, 0], [0, 0, 0]]]) + ) + assert np.isnan(one_empty["jsd"]).all() + assert np.isnan(one_empty["wasserstein_nt"]).all() + + +def test_locus_gene_and_read_aggregation(): + result = aggregate_observations( + np.asarray([1.0, 3.0, 5.0, 9.0]), + locus_ids=["a", "a", "b", "c"], + gene_ids=["g1", "g1", "g1", "g2"], + read_weights=np.asarray([1.0, 1.0, 2.0, 6.0]), + ) + assert result["locus_macro"]["value"] == pytest.approx((2 + 5 + 9) / 3) + assert result["gene_macro"]["value"] == pytest.approx((3.5 + 9) / 2) + assert result["read_micro"]["value"] == pytest.approx(6.8) + explicit = aggregate_observations( + np.asarray([2.0, 4.0]), + locus_ids=["a", "b"], + gene_ids=["g1", "g2"], + read_weights=np.asarray([2.0, 3.0]), + micro_numerators=np.asarray([2.0, 4.0]), + ) + assert explicit["read_micro"]["value"] == pytest.approx(6 / 5) + + +def test_read_weighted_calibration(): + rows = calibration_rows( + predicted_probability=np.asarray([[0.1], [0.2], [0.9]]), + ip_counts=np.asarray([[1], [3], [4]]), + total_counts=np.asarray([[2], [6], [4]]), + replicate_names=["ip1"], + n_bins=2, + ) + assert len(rows) == 2 + low = rows[0] + assert low["n_observations"] == 2 + assert low["total_reads"] == 8 + assert low["predicted_bin"] == pytest.approx((2 * 0.1 + 6 * 0.2) / 8) + assert low["observed_bin"] == pytest.approx(4 / 8) + assert rows[1]["predicted_bin"] == pytest.approx(0.9) + assert rows[1]["observed_bin"] == pytest.approx(1.0) + + +def test_representative_selection_is_deterministic_and_tiered(): + metric = np.arange(30, dtype=float) + first = select_representative_examples(metric, seed=9, per_tier=2) + second = select_representative_examples(metric, seed=9, per_tier=2) + assert first == second + assert [row["tier"] for row in first] == [ + "good", "good", "intermediate", "intermediate", "poor", "poor" + ] + assert all(row["index"] < 10 for row in first[:2]) + assert all(10 <= row["index"] < 20 for row in first[2:4]) + assert all(row["index"] >= 20 for row in first[4:]) + + +def test_checkpoint_split_resolution_defaults_to_test_and_never_falls_back(): + checkpoint = {"splits": {"train": [0, 1], "val": [2], "test": [3, 4]}} + name, indices = resolve_rbpnet_checkpoint_indices( + checkpoint, split=None, n_examples=5 + ) + assert name == "test" + assert indices == [3, 4] + assert resolve_rbpnet_checkpoint_indices( + checkpoint, split="all", n_examples=5 + ) == ("all", [0, 1, 2, 3, 4]) + with pytest.raises(ValueError, match="will not fall back"): + resolve_rbpnet_checkpoint_indices({}, split="test", n_examples=5) + with pytest.raises(ValueError, match="both train and test"): + resolve_rbpnet_checkpoint_indices( + {"splits": {"train": [0], "val": [1], "test": [0]}}, + split="test", + n_examples=2, + ) diff --git a/tests/test_rbpnet_model.py b/tests/test_rbpnet_model.py index f7b7880..10a16bc 100644 --- a/tests/test_rbpnet_model.py +++ b/tests/test_rbpnet_model.py @@ -11,8 +11,9 @@ SameLengthConvTranspose1d, SamePadConv1d, ) -from transcriptml.models.registry import build_model +from transcriptml.models.registry import build_model, load_checkpoint from transcriptml.rbpnet.dataset import RBPNetDataset, collate_rbpnet +from transcriptml.rbpnet.evaluation import evaluate_rbpnet_report from transcriptml.rbpnet.losses import ( RBPNetObjective, multinomial_nll, @@ -65,6 +66,8 @@ def _synthetic_bundle(n=9, length=16, jitter=2): "transcript_anchor": anchor, "selection_start": selection[0], "selection_end": selection[1], + "selection_state": ("peak", "gray", "negative")[index % 3], + "region_type": ("cds", "3putr", "mixed")[index % 3], "sequence_materialized_start": materialized_start, "sequence_materialized_end": materialized_start + width, "profile_materialized_start": materialized_start, @@ -388,6 +391,13 @@ def test_group_split_and_structured_training_profile_only_and_enrichment(tmp_pat assert (tmp_path / "enrichment" / "test_predictions.csv").is_file() bundle_dir = tmp_path / "bundle" + # Deliberately disagree with the checkpoint: RBPNet evaluation must use + # the immutable training artifact and never silently use bundle.splits. + bundle.splits = { + "train": list(range(1, len(bundle.ids) - 1)), + "val": [len(bundle.ids) - 1], + "test": [0], + } save_bundle(bundle, bundle_dir) predictions_path = tmp_path / "checkpoint_predictions.csv" evaluated = evaluate_checkpoint( @@ -397,10 +407,58 @@ def test_group_split_and_structured_training_profile_only_and_enrichment(tmp_pat batch_size=3, progress=False, ) - assert evaluated["pi"].shape == (len(bundle.ids),) - assert evaluated["enrichment_logit"].shape == (len(bundle.ids),) + _, saved_checkpoint = load_checkpoint( + tmp_path / "enrichment" / "best.pt", map_location="cpu" + ) + checkpoint_test = saved_checkpoint["splits"]["test"] + assert evaluated["indices"] == checkpoint_test + assert evaluated["pi"].shape == (len(checkpoint_test),) + assert evaluated["enrichment_logit"].shape == (len(checkpoint_test),) assert predictions_path.is_file() + report_dir = tmp_path / "evaluation_report" + report = evaluate_checkpoint( + tmp_path / "enrichment" / "best.pt", + bundle_dir, + out_dir=report_dir, + batch_size=3, + save_profiles=True, + representative_per_tier=1, + progress=False, + ) + assert report["indices"] == checkpoint_test + assert report["indices"] != bundle.splits["test"] + for name in ( + "summary.json", + "examples.parquet", + "stratified_metrics.parquet", + "calibration.parquet", + ): + assert (report_dir / name).is_file() + assert (report_dir / "plots" / "profile_performance_vs_read_depth.png").is_file() + assert (report_dir / "plots" / "enrichment_calibration.png").is_file() + assert ( + report_dir / "plots" / "stratified_performance_summaries.png" + ).is_file() + predicted_ip = np.load( + report_dir / "predicted_ip_profiles.npy", mmap_mode="r" + ) + assert isinstance(predicted_ip, np.memmap) + assert predicted_ip.shape == (len(checkpoint_test), 16) + + profile_report_dir = tmp_path / "profile_evaluation_report" + evaluate_checkpoint( + tmp_path / "profile" / "best.pt", + bundle_dir, + out_dir=profile_report_dir, + representative_per_tier=0, + progress=False, + ) + import pyarrow.parquet as pq + + assert pq.read_table(profile_report_dir / "calibration.parquet").num_rows == 0 + assert not (profile_report_dir / "plots" / "enrichment_calibration.png").exists() + unsafe = dict(base_config) unsafe["output_dir"] = str(tmp_path / "unsafe") unsafe["split"] = {"method": "random", "val_frac": 0.2, "test_frac": 0.2} @@ -436,6 +494,54 @@ def step(update): assert final < initial +def test_rbpnet_report_gracefully_handles_one_replicate_and_optional_metadata( + tmp_path, +): + original = _synthetic_bundle(n=4) + arrays = dict(original.arrays) + arrays["ip_profiles"] = arrays["ip_profiles"][:, :1, :] + arrays["profile_ip_totals"] = arrays["profile_ip_totals"][:, :1] + arrays["selection_ip_counts"] = arrays["selection_ip_counts"][:, :1] + config = dict(original.config) + config["sample_metadata"] = { + "sminput": {"name": "sminput", "effective_library_size": 100}, + "ip": [{"name": "ip1", "effective_library_size": 50}], + "ip_axis_order": ["ip1"], + } + metadata = [ + { + key: value + for key, value in row.items() + if key not in {"selection_state", "region_type"} + } + for row in original.metadata + ] + bundle = DatasetBundle( + X=original.X, + ids=original.ids, + metadata=metadata, + arrays=arrays, + config=config, + ) + result = evaluate_rbpnet_report( + _small_model(enrichment="none"), + {"splits": {"train": [0, 1], "val": [2], "test": [3]}}, + bundle, + tmp_path / "report", + representative_per_tier=0, + progress=False, + ) + assert result["indices"] == [3] + summary = result["summary"] + assert summary["enrichment_head_enabled"] is False + assert "eta_vs_empirical_enrichment.png" in summary["plots"]["skipped"] + assert "stratified_performance_summaries.png" in summary["plots"]["skipped"] + assert not any( + row["track"] == "replicate_ceiling" + for row in summary["overall_metrics"] + ) + + def test_rbpnet_training_consumes_saved_chromosome_cv_plan(tmp_path): bundle = _synthetic_bundle() plan_path = save_chromosome_cv_plan( From 04919ac14833a457c3671aaa28ed45bfbefe1bb3 Mon Sep 17 00:00:00 2001 From: isvock Date: Wed, 12 Aug 2026 15:41:25 -0700 Subject: [PATCH 09/12] Add evaluation scripts --- scripts/README.md | 2 +- scripts/rbpnet/README.md | 37 +++++++++++++++ scripts/rbpnet/eval_cv_fold.sh | 77 ++++++++++++++++++++++++++++++++ scripts/rbpnet/rbpnet_config.sh | 14 ++++++ scripts/rbpnet/submit_eval_cv.sh | 47 +++++++++++++++++++ 5 files changed, 176 insertions(+), 1 deletion(-) create mode 100644 scripts/rbpnet/eval_cv_fold.sh create mode 100644 scripts/rbpnet/submit_eval_cv.sh diff --git a/scripts/README.md b/scripts/README.md index 085daeb..c574ec8 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -17,7 +17,7 @@ These scripts are intentionally Sherlock-specific and deliberately small. For a - `motif_ablation_by_fold.sh` and `motif_ablation_all_folds.sh`: motif ablations across the configured motif list. - `motif_epistasis_by_fold.sh` and `motif_epistasis_all_folds.sh`: motif epistasis across the configured motif-pair list. - `mpra/`: MPRA 3-prime UTR insert workflows for building 4-channel LegNet input, training LegNet, and running single-nucleotide ISM. See `mpra/README.md`. -- `rbpnet/`: eCLIP preprocessing, scan/selection/bundle construction, immutable balanced chromosome CV planning, and one RBPNet training job per chromosome fold. See `rbpnet/README.md`. +- `rbpnet/`: eCLIP preprocessing, scan/selection/bundle construction, immutable balanced chromosome CV planning, RBPNet training, and structured scientific evaluation per chromosome fold. See `rbpnet/README.md`. ## Configure A Run diff --git a/scripts/rbpnet/README.md b/scripts/rbpnet/README.md index f6d4474..44f0ad1 100644 --- a/scripts/rbpnet/README.md +++ b/scripts/rbpnet/README.md @@ -8,6 +8,7 @@ FASTA/GTF/IP BAMs/SMInput BAM -> window scan, selection, and RBPNet bundle -> immutable balanced chromosome CV plan -> one independent training job per fold + -> one deterministic scientific evaluation report per fold ``` Copy this directory to a writable run directory and edit `rbpnet_config.sh`. @@ -30,6 +31,7 @@ sbatch scripts/rbpnet/preprocess.sh sbatch scripts/rbpnet/scan_select_bundle.sh sbatch scripts/rbpnet/create_chromosome_cv_plan.sh bash scripts/rbpnet/submit_train_cv.sh +bash scripts/rbpnet/submit_eval_cv.sh ``` The data-construction script chooses stride 1 automatically for @@ -47,6 +49,41 @@ that same file. Fold `k` uses group `k` as test, group `(k+1) mod N` as validation, and all other groups for training. `train_cv_fold.sh 0` can be run interactively for one fold without Slurm. +After every fold has a `${CV_ROOT}/fold/model/best.pt`, run +`submit_eval_cv.sh`. It submits one `eval_cv_fold.sh` task per fold and writes: + +```text +${EVAL_ROOT}/fold0/test/ + summary.json + examples.parquet + stratified_metrics.parquet + calibration.parquet + plots/ +``` + +Set `EVAL_SAVE_PROFILES=1` to additionally save memory-mappable predicted +target, control, and IP profiles. `EVAL_SPLIT` accepts `train`, `val`, `test`, +or `all` and defaults to `test`. Evaluation is deterministic with zero jitter +and resolves indices exclusively from each checkpoint, so the shared bundle's +own `splits` value is never used. `eval_cv_fold.sh 0` runs one fold +interactively. Other evaluation controls—including batch size, device, +calibration bins, enrichment pseudocount, and representative-example sampling— +are documented directly in `rbpnet_config.sh`. + +## Interpretation status + +The generic `transcriptml ism` and `transcriptml motif-ablation` commands do +not currently support RBPNet checkpoints. Those tools assume one scalar model +prediction per sequence. RBPNet instead returns structured target, control, +and IP positional distributions, `pi`, and optional `eta`; an interpretation +run therefore needs an explicit scientific objective such as profile +log-likelihood, regional mass, `pi`, or `eta`. The generic predictor also feeds +the materialized bundle width directly to the model and does not apply +RBPNet's coordinate-derived zero-jitter crop and validity/measurement masks, +which is incorrect for jitter-margin bundles. No ISM or motif-ablation Slurm +scripts are provided until a structured RBPNet attribution API defines those +choices explicitly. + `example_train_config.json` enables the independent linear enrichment head and 32-nt training jitter. Change `enrichment_head_type` to `none` for profile-only RBPNet, and keep `profile_length`/`max_train_jitter` consistent with the bundle. diff --git a/scripts/rbpnet/eval_cv_fold.sh b/scripts/rbpnet/eval_cv_fold.sh new file mode 100644 index 0000000..8beb4ef --- /dev/null +++ b/scripts/rbpnet/eval_cv_fold.sh @@ -0,0 +1,77 @@ +#!/bin/bash +#SBATCH --partition=akundaje +#SBATCH --job-name=tml_rbp_eval +#SBATCH --cpus-per-task=4 +#SBATCH --gpus=1 +#SBATCH --mem=32G +#SBATCH --time=08:00:00 +#SBATCH -C GPU_MEM:48GB +#SBATCH --output=slurm_output/%x_%A_%a.out +#SBATCH --error=slurm_output/%x_%A_%a.err + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/rbpnet/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts/rbpnet" + elif [[ -f "${SLURM_SUBMIT_DIR}/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/rbpnet_config.sh" +setup_transcriptml_env + +case "${EVAL_SPLIT}" in + train|val|test|all) ;; + *) + echo "EVAL_SPLIT must be train, val, test, or all; got ${EVAL_SPLIT}." >&2 + exit 1 + ;; +esac +if [[ "${EVAL_SAVE_PROFILES}" != "0" && "${EVAL_SAVE_PROFILES}" != "1" ]]; then + echo "EVAL_SAVE_PROFILES must be 0 or 1; got ${EVAL_SAVE_PROFILES}." >&2 + exit 1 +fi + +FOLD="${SLURM_ARRAY_TASK_ID:-${1:-}}" +if [[ -z "${FOLD}" ]]; then + echo "Provide a fold argument or run as a Slurm array job." >&2 + exit 1 +fi +if [[ ! "${FOLD}" =~ ^[0-9]+$ || "${FOLD}" -ge "${N_FOLDS}" ]]; then + echo "Fold must be an integer in [0, $((N_FOLDS - 1))]; got ${FOLD}." >&2 + exit 1 +fi + +CHECKPOINT="${CV_ROOT}/fold${FOLD}/model/${EVAL_CHECKPOINT_NAME}" +OUT_DIR="${EVAL_ROOT}/fold${FOLD}/${EVAL_SPLIT}" +if [[ ! -f "${CHECKPOINT}" ]]; then + echo "Missing fold checkpoint ${CHECKPOINT}; run submit_train_cv.sh first." >&2 + exit 1 +fi +if [[ ! -f "${BUNDLE_DIR}/config.json" ]]; then + echo "Missing RBPNet bundle at ${BUNDLE_DIR}; run scan_select_bundle.sh first." >&2 + exit 1 +fi + +command=( + transcriptml evaluate + --checkpoint "${CHECKPOINT}" + --dataset "${BUNDLE_DIR}" + --out-dir "${OUT_DIR}" + --split "${EVAL_SPLIT}" + --batch-size "${EVAL_BATCH_SIZE}" + --device "${EVAL_DEVICE}" + --calibration-bins "${EVAL_CALIBRATION_BINS}" + --enrichment-pseudocount "${EVAL_ENRICHMENT_PSEUDOCOUNT}" + --representative-seed "${EVAL_REPRESENTATIVE_SEED}" + --representative-per-tier "${EVAL_REPRESENTATIVE_PER_TIER}" + --representative-min-profile-count "${EVAL_REPRESENTATIVE_MIN_PROFILE_COUNT}" +) +if [[ "${EVAL_SAVE_PROFILES}" == "1" ]]; then + command+=(--save-profiles) +fi + +mkdir -p "${OUT_DIR}" +"${command[@]}" diff --git a/scripts/rbpnet/rbpnet_config.sh b/scripts/rbpnet/rbpnet_config.sh index 03ecd46..e782858 100644 --- a/scripts/rbpnet/rbpnet_config.sh +++ b/scripts/rbpnet/rbpnet_config.sh @@ -81,6 +81,20 @@ CHROMOSOME_GROUP_COL="${CHROMOSOME_GROUP_COL:-group_chromosome}" BASE_TRAIN_CONFIG="${BASE_TRAIN_CONFIG:-${SCRIPT_CONFIG_DIR}/example_train_config.json}" DEVICE="${DEVICE:-cuda}" +# Scientific checkpoint evaluation. Each fold report is written under +# EVAL_ROOT/fold// and always uses checkpoint-recorded indices. +EVAL_ROOT="${EVAL_ROOT:-${CV_ROOT}/evaluation}" +EVAL_CHECKPOINT_NAME="${EVAL_CHECKPOINT_NAME:-best.pt}" +EVAL_SPLIT="${EVAL_SPLIT:-test}" +EVAL_BATCH_SIZE="${EVAL_BATCH_SIZE:-128}" +EVAL_DEVICE="${EVAL_DEVICE:-${DEVICE}}" +EVAL_SAVE_PROFILES="${EVAL_SAVE_PROFILES:-0}" +EVAL_CALIBRATION_BINS="${EVAL_CALIBRATION_BINS:-10}" +EVAL_ENRICHMENT_PSEUDOCOUNT="${EVAL_ENRICHMENT_PSEUDOCOUNT:-0.5}" +EVAL_REPRESENTATIVE_SEED="${EVAL_REPRESENTATIVE_SEED:-123}" +EVAL_REPRESENTATIVE_PER_TIER="${EVAL_REPRESENTATIVE_PER_TIER:-3}" +EVAL_REPRESENTATIVE_MIN_PROFILE_COUNT="${EVAL_REPRESENTATIVE_MIN_PROFILE_COUNT:-10}" + setup_transcriptml_env() { module load gcc/10.1.0 module load openblas/0.3.10 diff --git a/scripts/rbpnet/submit_eval_cv.sh b/scripts/rbpnet/submit_eval_cv.sh new file mode 100644 index 0000000..cc8c00d --- /dev/null +++ b/scripts/rbpnet/submit_eval_cv.sh @@ -0,0 +1,47 @@ +#!/bin/bash + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/rbpnet/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts/rbpnet" + elif [[ -f "${SLURM_SUBMIT_DIR}/rbpnet_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/rbpnet_config.sh" + +case "${EVAL_SPLIT}" in + train|val|test|all) ;; + *) + echo "EVAL_SPLIT must be train, val, test, or all; got ${EVAL_SPLIT}." >&2 + exit 1 + ;; +esac +if [[ "${EVAL_SAVE_PROFILES}" != "0" && "${EVAL_SAVE_PROFILES}" != "1" ]]; then + echo "EVAL_SAVE_PROFILES must be 0 or 1; got ${EVAL_SAVE_PROFILES}." >&2 + exit 1 +fi + +if [[ ! -f "${BUNDLE_DIR}/config.json" ]]; then + echo "Missing RBPNet bundle at ${BUNDLE_DIR}; run scan_select_bundle.sh first." >&2 + exit 1 +fi + +missing=() +for ((fold = 0; fold < N_FOLDS; fold++)); do + checkpoint="${CV_ROOT}/fold${fold}/model/${EVAL_CHECKPOINT_NAME}" + if [[ ! -f "${checkpoint}" ]]; then + missing+=("${checkpoint}") + fi +done +if (( ${#missing[@]} > 0 )); then + echo "Cannot submit evaluation: missing ${#missing[@]} fold checkpoint(s):" >&2 + printf ' %s\n' "${missing[@]}" >&2 + echo "Run submit_train_cv.sh first or change EVAL_CHECKPOINT_NAME." >&2 + exit 1 +fi + +mkdir -p "${EVAL_ROOT}" slurm_output +sbatch --array="0-$((N_FOLDS - 1))" "${SCRIPT_DIR}/eval_cv_fold.sh" From 73a14c8cf4928854472c275d0e9715e257712929 Mon Sep 17 00:00:00 2001 From: isvock Date: Mon, 17 Aug 2026 14:11:50 -0700 Subject: [PATCH 10/12] Implemented region ablation analysis --- README.md | 3 +- docs/api.rst | 4 + docs/index.rst | 2 +- docs/usage.md | 74 +- scripts/README.md | 26 + scripts/region_ablation_by_fold.sh | 46 + scripts/sherlock_config.sh | 10 + scripts/submit_region_ablation_by_fold.sh | 14 + src/transcriptml/cli/main.py | 119 ++- src/transcriptml/data/controls.py | 206 +++-- src/transcriptml/data/region_edits.py | 185 ++++ src/transcriptml/interpret/__init__.py | 2 + src/transcriptml/interpret/codon_ism.py | 21 + src/transcriptml/interpret/predictor.py | 5 +- src/transcriptml/interpret/region_ablation.py | 842 ++++++++++++++++++ tests/test_cli_analysis.py | 119 ++- tests/test_region_ablation.py | 304 +++++++ 17 files changed, 1891 insertions(+), 91 deletions(-) create mode 100644 scripts/region_ablation_by_fold.sh create mode 100644 scripts/submit_region_ablation_by_fold.sh create mode 100644 src/transcriptml/data/region_edits.py create mode 100644 src/transcriptml/interpret/region_ablation.py create mode 100644 tests/test_region_ablation.py diff --git a/README.md b/README.md index c44a6c7..c912872 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,8 @@ TranscriptML is a toolkit for training, evaluating, and interpreting RNA sequence-to-function models. It provides command-line tools and reusable Python APIs for preparing sequence datasets, training models, evaluating held-out predictions, and investigating learned sequence features with analyses such as -in silico mutagenesis, motif ablation, context scans, etc. +in silico mutagenesis, transcript-region and exon-junction ablation, motif +ablation, context scans, etc. TranscriptML currently supports three main workflows: diff --git a/docs/api.rst b/docs/api.rst index 9119704..89ce318 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -159,6 +159,10 @@ Interpretation :members: MotifAblationResult, motif_ablation, save_motif_ablation_result :member-order: bysource +.. automodule:: transcriptml.interpret.region_ablation + :members: RegionAblationConfig, RegionAblationInstance, RegionAblationResult, region_ablation, save_region_ablation_result + :member-order: bysource + .. automodule:: transcriptml.interpret.context :members: MotifContextResult, motif_context_scan, save_motif_context_result :member-order: bysource diff --git a/docs/index.rst b/docs/index.rst index 4d5c9f2..6ee6425 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -11,7 +11,7 @@ two common starting points: with LegNet. Interpretation tools include single-nucleotide in silico mutagenesis (ISM), -motif ablations, motif context scans, motif epistasis analyses, and +region/junction ablations, motif ablations, motif context scans, motif epistasis analyses, and Saluki-specific codon ISM. These analyses can expose learned regulatory sequence features as well as technical artifacts in the model or assay. diff --git a/docs/usage.md b/docs/usage.md index cef1315..76ad08a 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -484,7 +484,75 @@ The `--region` flag can be `5utr`, `cds`, or `3utr`. Omit it to analyze motif instances across the full transcript. Region-aware analyses require Saluki-style annotation channels. -### 6. Run Codon Analyses +### 6. Run Region Ablation + +Region ablation measures how Saluki responds when complete transcript regions +or exon-junction arrangements are replaced with matched controls. For a coding +transcript, the default scan runs nucleotide shuffle and IID A/C/G/U +replacement controls for both UTRs, nucleotide shuffle, codon shuffle, and IID +replacement controls for the CDS, and junction-density scans at 1, 5, 10, ..., +50 junctions in the CDS. Transcripts without represented CDS annotation receive +the junction scan across their full valid sequence. + +Run one result per model fold: + +```bash +transcriptml region-ablation \ + --checkpoint runs/saluki_cv10/fold0/model/best.pt \ + --dataset data/saluki \ + --out-dir interpret/region_ablation/fold0 \ + --n-ablations 100 \ + --junction-counts 1,5,10,15,20,25,30,35,40,45,50 \ + --junction-min-spacing 25 \ + --seed 123 \ + --device auto \ + --batch-size 128 \ + --mutation-batch-size 512 +``` + +Sequence perturbations leave the CDS and splice channels unchanged. Junction +perturbations leave sequence and CDS annotation unchanged, clear the eligible +splice channel, and place exactly the requested number of new junction marks. +Coding-transcript UTR junctions are preserved. The 25-nt minimum separation is +a soft target: when a region cannot fit N junctions at that spacing, the command +uses the largest feasible spacing down to one nucleotide. A condition is skipped +only when the region length is not greater than N. + +`--n-ablations` sets a universal replicate count. Use repeatable overrides when +one family needs different sampling or should be disabled: + +```bash +transcriptml region-ablation \ + --checkpoint runs/saluki_cv10/fold0/model/best.pt \ + --dataset data/saluki \ + --out-dir interpret/region_ablation_targeted/fold0 \ + --n-ablations 100 \ + --n-ablations-for cds_codon_shuffle=250 \ + --n-ablations-for 5utr_shuffle=0 +``` + +The accepted family names are `5utr_shuffle`, `5utr_random`, +`cds_nt_shuffle`, `cds_codon_shuffle`, `cds_random`, `3utr_shuffle`, +`3utr_random`, and `junction_scatter`. The sequence slicing and sharding flags +match codon ISM: use `--sequence-start/--sequence-end` or +`--sequence-shard-index/--sequence-shards` for large runs. + +`instances.csv` has one row per transcript-condition. Raw replicate predictions +and signed `ablation - reference` effects occupy corresponding rows of +`ablation_predictions.npy` and `effects.npy`; `replicate_mask.npy` distinguishes +real values from NaN padding when family counts differ. The result also includes +per-condition mean, mean-absolute, and standard-deviation arrays, an audited +`skipped.csv`, and complete reproducibility metadata in `summary.json`. + +For the standard 11-point junction grid, a coding transcript with non-empty +UTRs receives 18 condition rows and approximately `18 * n_ablations` mutant +predictions. The Sherlock launcher is: + +```bash +bash scripts/submit_region_ablation_by_fold.sh +``` + +### 7. Run Codon Analyses Lots of work has shown that the coding sequence of an mRNA strongly influences its stability. Codon analyses are designed to dissect Saluki's understanding of this influence. @@ -516,7 +584,8 @@ per fold; see the TranscriptML [scripts](https://github.com/kundajelab/Transcrip ### HPC-Optimized Saluki Workflow The `scripts/` directory contains Sherlock-oriented SLURM jobs for Saluki input -building, 10-fold CV, hyperparameter sweeps, ISM, motif analyses, and codon ISM. +building, 10-fold CV, hyperparameter sweeps, ISM, region ablation, motif +analyses, and codon ISM. The scripts are intentionally editable. Copy them to a run directory, edit the copied config, and leave the clean repository checkout alone. @@ -559,6 +628,7 @@ hyperparameters, then submit stages: sbatch scripts/build_saluki_gtf.sh bash scripts/submit_train_eval_cv.sh bash scripts/submit_ism_by_fold.sh +bash scripts/submit_region_ablation_by_fold.sh bash scripts/submit_motif_ablation_by_fold.sh bash scripts/submit_motif_epistasis_by_fold.sh bash scripts/submit_codon_ism_by_fold.sh diff --git a/scripts/README.md b/scripts/README.md index c574ec8..eeddec2 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -14,6 +14,7 @@ These scripts are intentionally Sherlock-specific and deliberately small. For a - `ism_by_fold.sh` and `submit_ism_by_fold.sh`: single-nucleotide ISM, one job per trained fold. - `codon_ism_by_fold.sh` and `submit_codon_ism_by_fold.sh`: synonymous codon ISM, one job per trained fold. - `all_codon_ism_shard_by_fold.sh` and `submit_all_codon_ism_shard_by_fold.sh`: all-codon ISM, one 10-task job array per fold by default. +- `region_ablation_by_fold.sh` and `submit_region_ablation_by_fold.sh`: Saluki region and exon-junction density ablations, one job per fold. - `motif_ablation_by_fold.sh` and `motif_ablation_all_folds.sh`: motif ablations across the configured motif list. - `motif_epistasis_by_fold.sh` and `motif_epistasis_all_folds.sh`: motif epistasis across the configured motif-pair list. - `mpra/`: MPRA 3-prime UTR insert workflows for building 4-channel LegNet input, training LegNet, and running single-nucleotide ISM. See `mpra/README.md`. @@ -88,6 +89,8 @@ What each group means: | `N_FOLDS`, `CV_SEED`, `CV_VAL_OFFSET`, `CV_MODEL`, `EVAL_SPLIT` | Change these if you do not want the default 10-fold CV behavior or model name. | | `MODEL_DIR`, `EVAL_DIR`, `GENERATED_TRAIN_CONFIG`, `TRAIN_SEED`, `REQUIRE_SPLIT_FILE` | Optional controls for `train_eval_split.sh`. By default it writes under `${TRAIN_OUTPUT_ROOT}` and requires `${DATASET_DIR}/splits.json`. | | `PRED_BATCH_SIZE`, `MUTATION_BATCH_SIZE`, `DEVICE` | Runtime controls for GPU/CPU and prediction/ISM batch sizes. | +| `REGION_ABLATION_N`, `REGION_JUNCTION_COUNTS`, `REGION_JUNCTION_MIN_SPACING`, `REGION_ABLATION_SEED` | Universal region-ablation replicate count, junction-density grid, soft spacing target, and deterministic seed. | +| `REGION_ABLATION_N_FOR` | Optional Bash array of `FAMILY=COUNT` overrides; a zero count disables that family. | | `MOTIF_REGION` | Region for motif ablation and epistasis jobs. Defaults to `3utr`; use `5utr`, `cds`, `3utr`, or leave empty for whole-transcript analyses. | | `MOTIF_ABLATION_SPECS`, `MOTIF_EPISTASIS_SPECS` | Edit only when running motif ablation or motif epistasis with a custom motif list. | @@ -333,6 +336,29 @@ ${INTERPRET_ROOT}/ism/fold1/ ... ``` +## Run Region Ablation + +```bash +bash scripts/submit_region_ablation_by_fold.sh +``` + +The defaults run 100 replicates for each of seven coding-region sequence +controls and each junction count in `1,5,10,...,50`. Junction placement uses a +soft 25-nt spacing target. Customize the copied `sherlock_config.sh`, for +example: + +```bash +REGION_ABLATION_N=100 +REGION_JUNCTION_COUNTS="1,5,10,15,20,25,30,35,40,45,50" +REGION_JUNCTION_MIN_SPACING=25 +REGION_ABLATION_N_FOR=( + "cds_codon_shuffle=250" + "5utr_shuffle=0" +) +``` + +Outputs go to `${INTERPRET_ROOT}/region_ablation/fold*/`. + ## Run Synonymous Codon ISM ```bash diff --git a/scripts/region_ablation_by_fold.sh b/scripts/region_ablation_by_fold.sh new file mode 100644 index 0000000..aa4ad59 --- /dev/null +++ b/scripts/region_ablation_by_fold.sh @@ -0,0 +1,46 @@ +#!/bin/bash +#SBATCH --partition=akundaje +#SBATCH --job-name=tml_region_ablate +#SBATCH --cpus-per-task=4 +#SBATCH --gpus=1 +#SBATCH --mem=32G +#SBATCH --time=24:00:00 +#SBATCH -C GPU_MEM:48GB +#SBATCH --output=slurm_output/%x_%A_%a.out +#SBATCH --error=slurm_output/%x_%A_%a.err + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/sherlock_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts" + elif [[ -f "${SLURM_SUBMIT_DIR}/sherlock_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/sherlock_config.sh" +setup_transcriptml_env + +FOLD="${SLURM_ARRAY_TASK_ID}" +CHECKPOINT="${CV_ROOT}/fold${FOLD}/model/best.pt" +OUT_DIR="${INTERPRET_ROOT}/region_ablation/fold${FOLD}" + +REGION_ARGS=( + --checkpoint "${CHECKPOINT}" + --dataset "${INTERPRET_DATASET_DIR}" + --out-dir "${OUT_DIR}" + --n-ablations "${REGION_ABLATION_N}" + --junction-counts "${REGION_JUNCTION_COUNTS}" + --junction-min-spacing "${REGION_JUNCTION_MIN_SPACING}" + --seed "${REGION_ABLATION_SEED}" + --device "${DEVICE}" + --batch-size "${PRED_BATCH_SIZE}" + --mutation-batch-size "${MUTATION_BATCH_SIZE}" +) + +for spec in "${REGION_ABLATION_N_FOR[@]}"; do + REGION_ARGS+=(--n-ablations-for "${spec}") +done + +transcriptml region-ablation "${REGION_ARGS[@]}" diff --git a/scripts/sherlock_config.sh b/scripts/sherlock_config.sh index 6e55412..dda0ea6 100644 --- a/scripts/sherlock_config.sh +++ b/scripts/sherlock_config.sh @@ -89,6 +89,16 @@ PRED_BATCH_SIZE="${PRED_BATCH_SIZE:-128}" MUTATION_BATCH_SIZE="${MUTATION_BATCH_SIZE:-512}" DEVICE="${DEVICE:-cuda}" +# Region-ablation settings. Each FAMILY=COUNT entry overrides the universal N; +# use COUNT=0 to disable a family for a targeted run. +REGION_ABLATION_N="${REGION_ABLATION_N:-100}" +REGION_JUNCTION_COUNTS="${REGION_JUNCTION_COUNTS:-1,5,10,15,20,25,30,35,40,45,50}" +REGION_JUNCTION_MIN_SPACING="${REGION_JUNCTION_MIN_SPACING:-25}" +REGION_ABLATION_SEED="${REGION_ABLATION_SEED:-123}" +if ! declare -p REGION_ABLATION_N_FOR >/dev/null 2>&1; then + REGION_ABLATION_N_FOR=() +fi + # Set to all or transcript if you want to do full-transcript analysis MOTIF_REGION="${MOTIF_REGION:-3utr}" diff --git a/scripts/submit_region_ablation_by_fold.sh b/scripts/submit_region_ablation_by_fold.sh new file mode 100644 index 0000000..9515c4d --- /dev/null +++ b/scripts/submit_region_ablation_by_fold.sh @@ -0,0 +1,14 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + if [[ -f "${SLURM_SUBMIT_DIR}/scripts/sherlock_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}/scripts" + elif [[ -f "${SLURM_SUBMIT_DIR}/sherlock_config.sh" ]]; then + SCRIPT_DIR="${SLURM_SUBMIT_DIR}" + fi +fi +source "${SCRIPT_DIR}/sherlock_config.sh" + +sbatch --array=0-$((N_FOLDS - 1)) "${SCRIPT_DIR}/region_ablation_by_fold.sh" diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index b084ff1..602c071 100644 --- a/src/transcriptml/cli/main.py +++ b/src/transcriptml/cli/main.py @@ -5,6 +5,16 @@ from pathlib import Path DEFAULT_SALUKI_LENGTH = 12288 +_REGION_ABLATION_FAMILIES = ( + "5utr_shuffle", + "5utr_random", + "cds_nt_shuffle", + "cds_codon_shuffle", + "cds_random", + "3utr_shuffle", + "3utr_random", + "junction_scatter", +) def _csv_list(value: str | None) -> list[str] | None: @@ -28,6 +38,42 @@ def _maybe_int(value: str | None) -> str | int | None: return int(value) if value.isdigit() else value +def _positive_int_csv(value: str) -> tuple[int, ...]: + """Parse a non-empty comma-separated list of unique positive integers.""" + + try: + values = tuple(int(token.strip()) for token in value.split(",") if token.strip()) + except ValueError as exc: + raise argparse.ArgumentTypeError("expected comma-separated integers") from exc + if not values: + raise argparse.ArgumentTypeError("expected at least one integer") + if any(item <= 0 for item in values): + raise argparse.ArgumentTypeError("all values must be positive") + if len(set(values)) != len(values): + raise argparse.ArgumentTypeError("values must be unique") + return values + + +def _region_ablation_override(value: str) -> tuple[str, int]: + """Parse one ``FAMILY=COUNT`` region-ablation replicate override.""" + + if "=" not in value: + raise argparse.ArgumentTypeError("expected FAMILY=COUNT") + family, raw_count = value.split("=", 1) + family = family.strip() + if family not in _REGION_ABLATION_FAMILIES: + raise argparse.ArgumentTypeError( + f"unknown family {family!r}; expected one of {', '.join(_REGION_ABLATION_FAMILIES)}" + ) + try: + count = int(raw_count) + except ValueError as exc: + raise argparse.ArgumentTypeError("COUNT must be an integer") from exc + if count < 0: + raise argparse.ArgumentTypeError("COUNT must be non-negative") + return family, count + + def _analysis_install_message() -> str: return "This command requires the analysis extra: pip install 'TranscriptML[analysis]'" @@ -246,6 +292,7 @@ def build_parser() -> argparse.ArgumentParser: ("ism", "Run single-nucleotide ISM"), ("window-ism", "Run window-level random-mutagenesis ISM"), ("codon-ism", "Run CDS codon-level ISM"), + ("region-ablation", "Run Saluki transcript-region and junction ablations"), ("motif-ablation", "Run motif ablation"), ("motif-context", "Run motif context scan"), ("epistasis", "Run pairwise motif epistasis"), @@ -273,7 +320,7 @@ def build_parser() -> argparse.ArgumentParser: choices=["random_different", "shuffle", "dinuc_shuffle"], ) p.add_argument("--seed", type=int, default=123) - if name in {"ism", "window-ism", "codon-ism"}: + if name in {"ism", "window-ism", "codon-ism", "region-ablation"}: p.add_argument("--mutation-batch-size", type=int, default=512) if name == "window-ism": p.add_argument("--window-size", type=int, required=True) @@ -316,6 +363,30 @@ def build_parser() -> argparse.ArgumentParser: help="Streaming long-form mutation table format", ) p.add_argument("--rows-per-shard", type=int, default=100_000) + if name == "region-ablation": + p.add_argument("--n-ablations", type=int, default=100) + p.add_argument( + "--n-ablations-for", + type=_region_ablation_override, + action="append", + default=[], + metavar="FAMILY=COUNT", + help="Override replicates for one family; repeat as needed; zero disables it", + ) + p.add_argument( + "--junction-counts", + type=_positive_int_csv, + default=(1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50), + help="Unique positive junction counts in execution order", + ) + p.add_argument("--junction-min-spacing", type=int, default=25) + p.add_argument("--seed", type=int, default=123) + p.add_argument("--cds-channel", help="CDS annotation channel name or integer index") + p.add_argument("--splice-channel", help="Splice annotation channel name or integer index") + p.add_argument("--sequence-start", type=int, help="Inclusive transcript index") + p.add_argument("--sequence-end", type=int, help="Exclusive transcript index") + p.add_argument("--sequence-shard-index", type=int, help="Zero-based transcript shard") + p.add_argument("--sequence-shards", type=int, help="Total transcript shards") if name == "motif-context": p.add_argument("--window-size", type=int, default=5) p.add_argument("--context-width", type=int) @@ -685,7 +756,10 @@ def main(argv: list[str] | None = None) -> None: out_dir = interpret_paths["out_dir"] log_progress(f"{args.command}: loading dataset {dataset}") - bundle = load_bundle(dataset, mmap_mode="r" if args.command in {"codon-ism", "window-ism"} else None) + bundle = load_bundle( + dataset, + mmap_mode="r" if args.command in {"codon-ism", "window-ism", "region-ablation"} else None, + ) log_progress(f"{args.command}: loading checkpoint {checkpoint}") predictor = Predictor.from_checkpoint(checkpoint, device=args.device, batch_size=args.batch_size) cds_channel = _maybe_int(getattr(args, "cds_channel", None)) @@ -744,6 +818,47 @@ def main(argv: list[str] | None = None) -> None: sequence_shards=args.sequence_shards, ) save_codon_ism_result(result, out_dir, save_mutations=False) + elif args.command == "region-ablation": + from transcriptml.interpret.region_ablation import ( + RegionAblationConfig, + region_ablation, + save_region_ablation_result, + ) + + overrides: dict[str, int] = {} + for family, count in args.n_ablations_for: + if family in overrides: + parser.error(f"region-ablation got duplicate --n-ablations-for family {family!r}") + overrides[family] = count + result = region_ablation( + bundle.X, + predictor, + schema=bundle.schema, + sequence_ids=bundle.ids, + metadata=bundle.metadata, + config=RegionAblationConfig( + n_ablations=args.n_ablations, + n_ablations_for=overrides, + junction_counts=tuple(args.junction_counts), + junction_min_spacing=args.junction_min_spacing, + seed=args.seed, + ), + cds_channel=cds_channel, + splice_channel=_maybe_int(args.splice_channel), + reference_batch_size=args.batch_size, + mutation_batch_size=args.mutation_batch_size, + sequence_start=args.sequence_start, + sequence_end=args.sequence_end, + sequence_shard_index=args.sequence_shard_index, + sequence_shards=args.sequence_shards, + storage_dir=out_dir, + ) + save_region_ablation_result( + result, + out_dir, + checkpoint=checkpoint, + dataset=dataset, + ) elif args.command == "motif-ablation": from transcriptml.interpret.ablation import motif_ablation, save_motif_ablation_result diff --git a/src/transcriptml/data/controls.py b/src/transcriptml/data/controls.py index 9e0a2da..23d41ee 100644 --- a/src/transcriptml/data/controls.py +++ b/src/transcriptml/data/controls.py @@ -10,6 +10,17 @@ import numpy as np from transcriptml.data.bundle import DatasetBundle, save_bundle_metadata +from transcriptml.data.region_edits import ( + base_channel_indices as _shared_base_channel_indices, + base_symbols as _shared_base_symbols, + randomize_nucleotides_inplace as _shared_randomize_nucleotides, + region_bounds as _shared_region_bounds, + resolve_cds_channel as _shared_resolve_cds_channel, + shuffle_codons_inplace as _shared_shuffle_codons, + shuffle_nucleotides_inplace as _shared_shuffle_nucleotides, + valid_length_from_bases as _shared_valid_length, + write_base_symbols as _shared_write_base_symbols, +) from transcriptml.data.schemas import SequenceSchema, get_schema from transcriptml.interpret.codon_ism import CDSCodonStarts, find_cds_codon_starts from transcriptml.progress import ProgressReporter, log_progress @@ -443,49 +454,15 @@ def normalize_sequence_control_config(config: object) -> SequenceControlConfig: def _base_channel_indices(schema: SequenceSchema) -> np.ndarray: - indices = [] - letters = [] - for base_name in schema.base_channels: - if base_name not in schema.channels: - raise ValueError(f"Base channel '{base_name}' is not present in schema channels {schema.channels}") - letter = base_name.upper().replace("T", "U") - if letter not in {"A", "C", "G", "U"}: - raise ValueError(f"Unsupported base channel '{base_name}'; expected A/C/G/U/T") - indices.append(schema.channels.index(base_name)) - letters.append(letter) - if set(letters) != {"A", "C", "G", "U"} or len(letters) != 4: - raise ValueError("sequence_controls requires exactly one A, C, G, and U/T base channel") - return np.asarray(indices, dtype=np.int64) + return _shared_base_channel_indices(schema) def _resolve_cds_channel(schema: SequenceSchema, cds_channel: str | int | None) -> int: - if isinstance(cds_channel, int): - if cds_channel < 0 or cds_channel >= schema.n_channels: - raise ValueError(f"cds_channel index {cds_channel} is outside schema with {schema.n_channels} channels") - return int(cds_channel) - if isinstance(cds_channel, str): - try: - return schema.channels.index(cds_channel) - except ValueError as exc: - raise ValueError(f"cds_channel '{cds_channel}' is not in schema channels {schema.channels}") from exc - - preferred = ("CDS_codon_start", "cds_codon_start", "codon_start", "CDS", "cds") - lower_to_index = {name.lower(): i for i, name in enumerate(schema.channels)} - for name in preferred: - if name.lower() in lower_to_index: - return lower_to_index[name.lower()] - for i, name in enumerate(schema.channels): - lowered = name.lower() - if "cds" in lowered or "coding" in lowered or "codon_start" in lowered: - return i - raise ValueError("Could not infer CDS channel from schema; pass cds_channel explicitly") + return _shared_resolve_cds_channel(schema, cds_channel) def _infer_valid_length(x: np.ndarray, base_channels: np.ndarray) -> int: - base = np.asarray(x[base_channels]) - nonzero = np.any(base != 0, axis=0) - idx = np.nonzero(nonzero)[0] - return int(idx[-1] + 1) if idx.size else 0 + return _shared_valid_length(x, base_channels) def _mixed_rng(seed: int, seq_index: int, operation: OperationName, region: RegionName) -> np.random.Generator: @@ -504,14 +481,7 @@ def _mixed_rng(seed: int, seq_index: int, operation: OperationName, region: Regi def _base_symbols(x: np.ndarray, start: int, end: int, base_channels: np.ndarray) -> np.ndarray: - if end <= start: - return np.empty((0,), dtype=np.int16) - region = np.asarray(x[base_channels, int(start) : int(end)]) - called = np.count_nonzero(region, axis=0) == 1 - symbols = np.full(region.shape[1], -1, dtype=np.int16) - if np.any(called): - symbols[called] = np.argmax(region[:, called], axis=0).astype(np.int16, copy=False) - return symbols + return _shared_base_symbols(x, start, end, base_channels) def _write_base_symbols( @@ -521,17 +491,7 @@ def _write_base_symbols( symbols: np.ndarray, base_channels: np.ndarray, ) -> None: - if end <= start: - return - start = int(start) - end = int(end) - x[base_channels, start:end] = 0 - valid = np.asarray(symbols) >= 0 - if not np.any(valid): - return - cols = start + np.nonzero(valid)[0] - channel_offsets = np.asarray(symbols[valid], dtype=np.int64) - x[base_channels[channel_offsets], cols] = 1 + _shared_write_base_symbols(x, start, end, symbols, base_channels) def _region_bounds( @@ -540,17 +500,7 @@ def _region_bounds( valid_length: int, cds: CDSCodonStarts | None, ) -> tuple[int, int] | None: - if region == "transcript": - return 0, int(valid_length) - if cds is None or cds.cds_length < 3 or cds.starts.size == 0: - return None - cds_start = max(0, int(cds.cds_start)) - cds_end = min(int(valid_length), int(cds.cds_end) + 1) - if region == "5utr": - return 0, cds_start - if region == "cds": - return cds_start, cds_end - return cds_end, int(valid_length) + return _shared_region_bounds(region, valid_length=valid_length, cds=cds) def _shuffle_nucleotides( @@ -561,10 +511,13 @@ def _shuffle_nucleotides( base_channels: np.ndarray, rng: np.random.Generator, ) -> None: - symbols = _base_symbols(x, start, end, base_channels) - if symbols.size > 1: - symbols = symbols[rng.permutation(symbols.size)] - _write_base_symbols(x, start, end, symbols, base_channels) + _shared_shuffle_nucleotides( + x, + start=start, + end=end, + base_channels=base_channels, + rng=rng, + ) def _randomize_nucleotides( @@ -575,9 +528,13 @@ def _randomize_nucleotides( base_channels: np.ndarray, rng: np.random.Generator, ) -> None: - length = max(0, int(end) - int(start)) - symbols = rng.integers(0, int(base_channels.size), size=length, dtype=np.int16) - _write_base_symbols(x, start, end, symbols, base_channels) + _shared_randomize_nucleotides( + x, + start=start, + end=end, + base_channels=base_channels, + rng=rng, + ) def _frameshift_cds_channel( @@ -609,15 +566,98 @@ def _shuffle_codons( base_channels: np.ndarray, rng: np.random.Generator, ) -> None: - starts = np.asarray(cds.starts, dtype=np.int64) - starts = starts[(starts >= int(cds.cds_start)) & (starts + 2 <= int(cds.cds_end))] - if starts.size == 0: - return - codons = np.stack([_base_symbols(x, int(start), int(start) + 3, base_channels) for start in starts], axis=0) - if codons.shape[0] > 1: - codons = codons[rng.permutation(codons.shape[0])] - for start, codon in zip(starts.tolist(), codons, strict=True): - _write_base_symbols(x, int(start), int(start) + 3, codon, base_channels) + _shared_shuffle_codons(x, cds=cds, base_channels=base_channels, rng=rng) + + +def sequence_control_base_channels(schema: str | SequenceSchema = "saluki6") -> np.ndarray: + """Return base-channel indices using the sequence-control conventions. + + This public helper lets interpretation analyses use the exact same channel + resolution and edit semantics as training-time sequence controls. + """ + + return _base_channel_indices(get_schema(schema)) + + +def resolve_sequence_control_cds_channel( + schema: str | SequenceSchema = "saluki6", + cds_channel: str | int | None = None, +) -> int: + """Resolve the CDS channel using sequence-control conventions.""" + + return _resolve_cds_channel(get_schema(schema), cds_channel) + + +def sequence_control_valid_length(x: np.ndarray, base_channels: np.ndarray) -> int: + """Infer represented length from the resolved base channels.""" + + return _infer_valid_length(x, np.asarray(base_channels, dtype=np.int64)) + + +def sequence_control_region_bounds( + region: RegionName, + *, + valid_length: int, + cds: CDSCodonStarts | None, +) -> tuple[int, int] | None: + """Return region bounds using the training-time sequence-control rules.""" + + return _region_bounds(region, valid_length=valid_length, cds=cds) + + +def shuffle_region_nucleotides_inplace( + x: np.ndarray, + *, + start: int, + end: int, + base_channels: np.ndarray, + rng: np.random.Generator, +) -> None: + """Shuffle one region with the same semantics as sequence controls.""" + + _shuffle_nucleotides( + x, + start=start, + end=end, + base_channels=np.asarray(base_channels, dtype=np.int64), + rng=rng, + ) + + +def randomize_region_nucleotides_inplace( + x: np.ndarray, + *, + start: int, + end: int, + base_channels: np.ndarray, + rng: np.random.Generator, +) -> None: + """Replace one region with IID bases using sequence-control semantics.""" + + _randomize_nucleotides( + x, + start=start, + end=end, + base_channels=np.asarray(base_channels, dtype=np.int64), + rng=rng, + ) + + +def shuffle_cds_codons_inplace( + x: np.ndarray, + *, + cds: CDSCodonStarts, + base_channels: np.ndarray, + rng: np.random.Generator, +) -> None: + """Shuffle annotated CDS codons using sequence-control semantics.""" + + _shuffle_codons( + x, + cds=cds, + base_channels=np.asarray(base_channels, dtype=np.int64), + rng=rng, + ) def _empty_nested_counts() -> dict[str, dict[str, int]]: diff --git a/src/transcriptml/data/region_edits.py b/src/transcriptml/data/region_edits.py new file mode 100644 index 0000000..9e2284c --- /dev/null +++ b/src/transcriptml/data/region_edits.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np + +from transcriptml.data.schemas import SequenceSchema, get_schema + + +def base_channel_indices(schema: str | SequenceSchema = "saluki6") -> np.ndarray: + """Resolve exactly one A, C, G, and U/T channel.""" + + resolved = get_schema(schema) + indices = [] + letters = [] + for base_name in resolved.base_channels: + if base_name not in resolved.channels: + raise ValueError( + f"Base channel {base_name!r} is not present in schema channels {resolved.channels}" + ) + letter = base_name.upper().replace("T", "U") + if letter not in {"A", "C", "G", "U"}: + raise ValueError(f"Unsupported base channel {base_name!r}; expected A/C/G/U/T") + indices.append(resolved.channels.index(base_name)) + letters.append(letter) + if set(letters) != {"A", "C", "G", "U"} or len(letters) != 4: + raise ValueError("region edits require exactly one A, C, G, and U/T base channel") + return np.asarray(indices, dtype=np.int64) + + +def resolve_cds_channel( + schema: str | SequenceSchema = "saluki6", + cds_channel: str | int | None = None, +) -> int: + """Resolve a CDS channel name/index using TranscriptML conventions.""" + + resolved = get_schema(schema) + if isinstance(cds_channel, int): + if cds_channel < 0 or cds_channel >= resolved.n_channels: + raise ValueError( + f"cds_channel index {cds_channel} is outside schema with " + f"{resolved.n_channels} channels" + ) + return int(cds_channel) + if isinstance(cds_channel, str): + try: + return resolved.channels.index(cds_channel) + except ValueError as exc: + raise ValueError( + f"cds_channel {cds_channel!r} is not in schema channels {resolved.channels}" + ) from exc + + preferred = ("CDS_codon_start", "cds_codon_start", "codon_start", "CDS", "cds") + lower_to_index = {name.lower(): i for i, name in enumerate(resolved.channels)} + for name in preferred: + if name.lower() in lower_to_index: + return lower_to_index[name.lower()] + for i, name in enumerate(resolved.channels): + lowered = name.lower() + if "cds" in lowered or "coding" in lowered or "codon_start" in lowered: + return i + raise ValueError("Could not infer CDS channel from schema; pass cds_channel explicitly") + + +def valid_length_from_bases(x: np.ndarray, base_channels: np.ndarray) -> int: + """Infer represented length from nonzero base-channel columns.""" + + base = np.asarray(x[np.asarray(base_channels, dtype=np.int64)]) + idx = np.nonzero(np.any(base != 0, axis=0))[0] + return int(idx[-1] + 1) if idx.size else 0 + + +def base_symbols(x: np.ndarray, start: int, end: int, base_channels: np.ndarray) -> np.ndarray: + """Read base-channel offsets, retaining ambiguous positions as ``-1``.""" + + if end <= start: + return np.empty((0,), dtype=np.int16) + channels = np.asarray(base_channels, dtype=np.int64) + region = np.asarray(x[channels, int(start) : int(end)]) + called = np.count_nonzero(region, axis=0) == 1 + symbols = np.full(region.shape[1], -1, dtype=np.int16) + if np.any(called): + symbols[called] = np.argmax(region[:, called], axis=0).astype(np.int16, copy=False) + return symbols + + +def write_base_symbols( + x: np.ndarray, + start: int, + end: int, + symbols: np.ndarray, + base_channels: np.ndarray, +) -> None: + """Write base-channel offsets while leaving annotations untouched.""" + + if end <= start: + return + start = int(start) + end = int(end) + channels = np.asarray(base_channels, dtype=np.int64) + x[channels, start:end] = 0 + valid = np.asarray(symbols) >= 0 + if not np.any(valid): + return + columns = start + np.nonzero(valid)[0] + offsets = np.asarray(symbols[valid], dtype=np.int64) + x[channels[offsets], columns] = 1 + + +def region_bounds( + region: str, + *, + valid_length: int, + cds: Any | None, +) -> tuple[int, int] | None: + """Return transcript, UTR, or CDS half-open coordinates.""" + + if region == "transcript": + return 0, int(valid_length) + if cds is None or int(cds.cds_length) < 3 or np.asarray(cds.starts).size == 0: + return None + cds_start = max(0, int(cds.cds_start)) + cds_end = min(int(valid_length), int(cds.cds_end) + 1) + if region == "5utr": + return 0, cds_start + if region == "cds": + return cds_start, cds_end + if region == "3utr": + return cds_end, int(valid_length) + raise ValueError("region must be one of: 5utr, cds, 3utr, transcript") + + +def shuffle_nucleotides_inplace( + x: np.ndarray, + *, + start: int, + end: int, + base_channels: np.ndarray, + rng: np.random.Generator, +) -> None: + """Permute nucleotide/ambiguous symbols within one region.""" + + symbols = base_symbols(x, start, end, base_channels) + if symbols.size > 1: + symbols = symbols[rng.permutation(symbols.size)] + write_base_symbols(x, start, end, symbols, base_channels) + + +def randomize_nucleotides_inplace( + x: np.ndarray, + *, + start: int, + end: int, + base_channels: np.ndarray, + rng: np.random.Generator, +) -> None: + """Replace every position in one region with an IID A/C/G/U base.""" + + length = max(0, int(end) - int(start)) + channels = np.asarray(base_channels, dtype=np.int64) + symbols = rng.integers(0, int(channels.size), size=length, dtype=np.int16) + write_base_symbols(x, start, end, symbols, channels) + + +def shuffle_codons_inplace( + x: np.ndarray, + *, + cds: Any, + base_channels: np.ndarray, + rng: np.random.Generator, +) -> None: + """Permute complete annotated CDS codons as three-base units.""" + + starts = np.asarray(cds.starts, dtype=np.int64) + starts = starts[(starts >= int(cds.cds_start)) & (starts + 2 <= int(cds.cds_end))] + if starts.size == 0: + return + codons = np.stack( + [base_symbols(x, int(start), int(start) + 3, base_channels) for start in starts], + axis=0, + ) + if codons.shape[0] > 1: + codons = codons[rng.permutation(codons.shape[0])] + for start, codon in zip(starts.tolist(), codons, strict=True): + write_base_symbols(x, int(start), int(start) + 3, codon, base_channels) diff --git a/src/transcriptml/interpret/__init__.py b/src/transcriptml/interpret/__init__.py index 8e57d9d..3691d33 100644 --- a/src/transcriptml/interpret/__init__.py +++ b/src/transcriptml/interpret/__init__.py @@ -6,6 +6,7 @@ from transcriptml.interpret.epistasis import motif_epistasis from transcriptml.interpret.ism import compute_ism from transcriptml.interpret.predictor import EnsemblePredictor, Predictor +from transcriptml.interpret.region_ablation import region_ablation from transcriptml.interpret.window_ism import compute_window_ism __all__ = [ @@ -17,4 +18,5 @@ "motif_ablation", "motif_context_scan", "motif_epistasis", + "region_ablation", ] diff --git a/src/transcriptml/interpret/codon_ism.py b/src/transcriptml/interpret/codon_ism.py index 8dddea1..85d2ad3 100644 --- a/src/transcriptml/interpret/codon_ism.py +++ b/src/transcriptml/interpret/codon_ism.py @@ -770,6 +770,27 @@ def _resolve_analysis_indices( return np.arange(n_sequences, dtype=np.int64) +def resolve_analysis_indices( + n_sequences: int, + *, + sequence_indices: Sequence[int] | None = None, + sequence_start: int | None = None, + sequence_end: int | None = None, + sequence_shard_index: int | None = None, + sequence_shards: int | None = None, +) -> np.ndarray: + """Resolve interpretation sequence selectors to original input indices.""" + + return _resolve_analysis_indices( + n_sequences, + sequence_indices=sequence_indices, + sequence_start=sequence_start, + sequence_end=sequence_end, + sequence_shard_index=sequence_shard_index, + sequence_shards=sequence_shards, + ) + + @torch.no_grad() def compute_codon_ism( X: np.ndarray | torch.Tensor, diff --git a/src/transcriptml/interpret/predictor.py b/src/transcriptml/interpret/predictor.py index 2f2bc39..d26600d 100644 --- a/src/transcriptml/interpret/predictor.py +++ b/src/transcriptml/interpret/predictor.py @@ -77,7 +77,10 @@ def predict(self, X: np.ndarray | torch.Tensor, *, batch_size: int | None = None if isinstance(batch, torch.Tensor): xb = batch.to(self.device, dtype=torch.float32) else: - xb = torch.as_tensor(np.asarray(batch), dtype=torch.float32).to(self.device) + batch_array = np.asarray(batch) + if not batch_array.flags.writeable: + batch_array = np.array(batch_array, copy=True) + xb = torch.as_tensor(batch_array, dtype=torch.float32).to(self.device) y = squeeze_prediction(self.model(xb)) outs.append(y.detach().cpu().numpy().astype(np.float32, copy=False).reshape(-1)) return np.concatenate(outs) if outs else np.empty((0,), dtype=np.float32) diff --git a/src/transcriptml/interpret/region_ablation.py b/src/transcriptml/interpret/region_ablation.py new file mode 100644 index 0000000..498f862 --- /dev/null +++ b/src/transcriptml/interpret/region_ablation.py @@ -0,0 +1,842 @@ +from __future__ import annotations + +import hashlib +import json +from collections import Counter +from dataclasses import dataclass, field +from pathlib import Path +from typing import Mapping, Sequence + +import numpy as np + +from transcriptml.data.region_edits import ( + base_channel_indices, + randomize_nucleotides_inplace, + region_bounds, + resolve_cds_channel, + shuffle_codons_inplace, + shuffle_nucleotides_inplace, + valid_length_from_bases, +) +from transcriptml.data.schemas import SequenceSchema, get_schema +from transcriptml.interpret.codon_ism import ( + CDSCodonStarts, + find_cds_codon_starts, + resolve_analysis_indices, +) +from transcriptml.interpret.predictor import Predictor +from transcriptml.interpret.results import save_table +from transcriptml.progress import ProgressReporter, log_progress + + +DEFAULT_JUNCTION_COUNTS: tuple[int, ...] = (1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50) +REGION_ABLATION_FAMILIES: tuple[str, ...] = ( + "5utr_shuffle", + "5utr_random", + "cds_nt_shuffle", + "cds_codon_shuffle", + "cds_random", + "3utr_shuffle", + "3utr_random", + "junction_scatter", +) + +_SEQUENCE_FAMILIES: tuple[tuple[str, str, str], ...] = ( + ("5utr_shuffle", "5utr", "shuffle_nucleotides"), + ("5utr_random", "5utr", "randomize_nucleotides"), + ("cds_nt_shuffle", "cds", "shuffle_nucleotides"), + ("cds_codon_shuffle", "cds", "shuffle_codons"), + ("cds_random", "cds", "randomize_nucleotides"), + ("3utr_shuffle", "3utr", "shuffle_nucleotides"), + ("3utr_random", "3utr", "randomize_nucleotides"), +) +_FAMILY_CODES = {family: i + 1 for i, family in enumerate(REGION_ABLATION_FAMILIES)} + + +@dataclass(frozen=True) +class RegionAblationConfig: + """Configuration for repeated region and junction perturbations.""" + + n_ablations: int = 100 + n_ablations_for: Mapping[str, int] = field(default_factory=dict) + junction_counts: tuple[int, ...] = DEFAULT_JUNCTION_COUNTS + junction_min_spacing: int = 25 + seed: int = 123 + + def normalized(self) -> "RegionAblationConfig": + """Validate values and return an immutable normalized configuration.""" + + default_n = int(self.n_ablations) + if default_n <= 0: + raise ValueError("n_ablations must be positive") + if int(self.seed) < 0: + raise ValueError("seed must be non-negative") + if int(self.junction_min_spacing) <= 0: + raise ValueError("junction_min_spacing must be positive") + + counts = tuple(int(value) for value in self.junction_counts) + if not counts: + raise ValueError("junction_counts must contain at least one count") + if any(value <= 0 for value in counts): + raise ValueError("junction_counts must contain only positive integers") + if len(set(counts)) != len(counts): + raise ValueError("junction_counts must be unique") + + overrides: dict[str, int] = {} + for raw_family, raw_count in dict(self.n_ablations_for).items(): + family = str(raw_family) + if family not in REGION_ABLATION_FAMILIES: + raise ValueError( + f"Unknown region-ablation family {family!r}; expected one of " + f"{', '.join(REGION_ABLATION_FAMILIES)}" + ) + count = int(raw_count) + if count < 0: + raise ValueError("per-family ablation counts must be non-negative") + overrides[family] = count + return RegionAblationConfig( + n_ablations=default_n, + n_ablations_for=overrides, + junction_counts=counts, + junction_min_spacing=int(self.junction_min_spacing), + seed=int(self.seed), + ) + + def replicates_for(self, family: str) -> int: + """Return the requested replicate count for one perturbation family.""" + + return int(self.n_ablations_for.get(family, self.n_ablations)) + + def to_dict(self) -> dict[str, object]: + """Return a JSON-compatible configuration mapping.""" + + return { + "n_ablations": int(self.n_ablations), + "n_ablations_for": {str(k): int(v) for k, v in self.n_ablations_for.items()}, + "junction_counts": [int(value) for value in self.junction_counts], + "junction_min_spacing": int(self.junction_min_spacing), + "seed": int(self.seed), + } + + +@dataclass(frozen=True) +class RegionAblationInstance: + """One transcript-condition row aligned to region-ablation arrays.""" + + instance_index: int + seq_index: int + sequence_id: str + transcript_class: str + operation: str + region: str + valid_length: int + region_start: int + region_end: int + region_length: int + n_replicates: int + junction_count: int | None = None + reference_junction_count: int | None = None + requested_min_spacing: int | None = None + effective_min_spacing: int | None = None + + +@dataclass(frozen=True) +class SkippedRegionAblation: + """A transcript-condition omitted from the scored instance table.""" + + seq_index: int + sequence_id: str + transcript_class: str + operation: str + region: str + junction_count: int | None + reason: str + + +@dataclass +class RegionAblationResult: + """Raw and summarized repeated region-ablation predictions.""" + + instances: list[RegionAblationInstance] + skipped: list[SkippedRegionAblation] + reference_predictions: np.ndarray + ablation_predictions: np.ndarray + effects: np.ndarray + replicate_mask: np.ndarray + mean_effects: np.ndarray + mean_abs_effects: np.ndarray + std_effects: np.ndarray + analysis_indices: np.ndarray + transcript_classes: tuple[str, ...] + config: RegionAblationConfig + input_shape: tuple[int, int, int] + schema_name: str + cds_channel_index: int + splice_channel_index: int + sequence_ids: tuple[str, ...] + storage_dir: Path | None = None + + +def _resolve_splice_channel(schema: SequenceSchema, splice_channel: str | int | None) -> int: + """Resolve a splice-junction channel selector to an integer index.""" + + if isinstance(splice_channel, int): + if splice_channel < 0 or splice_channel >= schema.n_channels: + raise ValueError( + f"splice_channel index {splice_channel} is outside schema with " + f"{schema.n_channels} channels" + ) + return int(splice_channel) + if isinstance(splice_channel, str): + try: + return schema.channels.index(splice_channel) + except ValueError as exc: + raise ValueError( + f"splice_channel {splice_channel!r} is not in schema channels {schema.channels}" + ) from exc + + lower_to_index = {name.lower(): i for i, name in enumerate(schema.channels)} + for name in ("splice_junction", "splice-junction", "splice", "junction"): + if name in lower_to_index: + return lower_to_index[name] + for i, name in enumerate(schema.channels): + lowered = name.lower() + if "splice" in lowered or "junction" in lowered: + return i + raise ValueError("Could not infer splice-junction channel from schema; pass splice_channel explicitly") + + +def _metadata_reports_coding(metadata: Mapping[str, object] | None) -> bool: + """Return whether metadata explicitly reports a positive original CDS length.""" + + if metadata is None or metadata.get("cds_length") is None: + return False + try: + return float(metadata["cds_length"]) > 0 + except (TypeError, ValueError): + return False + + +def effective_junction_spacing(region_length: int, junction_count: int, requested_spacing: int) -> int: + """Return the largest feasible spacing up to the requested soft target.""" + + length = int(region_length) + count = int(junction_count) + requested = int(requested_spacing) + if count <= 0: + raise ValueError("junction_count must be positive") + if requested <= 0: + raise ValueError("requested_spacing must be positive") + if length <= count: + raise ValueError("region_length must be greater than junction_count") + if count == 1: + return requested + maximum = max(1, (length - 2) // (count - 1)) + return min(requested, maximum) + + +def sample_junction_positions( + *, + start: int, + end: int, + junction_count: int, + min_spacing: int, + rng: np.random.Generator, +) -> np.ndarray: + """Uniformly sample spaced junction marks within a half-open region. + + Candidate marks are ``start`` through ``end - 2`` so every mark has at + least one downstream nucleotide. Compressed coordinates provide a + rejection-free bijection to layouts obeying the requested separation. + """ + + region_start = int(start) + region_end = int(end) + count = int(junction_count) + spacing = int(min_spacing) + region_length = region_end - region_start + if region_length <= count: + raise ValueError("region length must be greater than junction_count") + if count <= 0: + raise ValueError("junction_count must be positive") + if spacing <= 0: + raise ValueError("min_spacing must be positive") + + candidate_count = region_length - 1 + if count == 1: + return np.asarray( + [region_start + int(rng.integers(0, candidate_count))], + dtype=np.int64, + ) + maximum_spacing = max(1, (region_length - 2) // (count - 1)) + if spacing > maximum_spacing: + raise ValueError( + f"min_spacing={spacing} is infeasible for {count} junctions in length {region_length}" + ) + compressed_count = candidate_count - (spacing - 1) * (count - 1) + compressed = np.sort(rng.choice(compressed_count, size=count, replace=False)) + expanded = compressed + np.arange(count, dtype=np.int64) * (spacing - 1) + return expanded.astype(np.int64, copy=False) + region_start + + +def scatter_junctions_inplace( + x: np.ndarray, + *, + start: int, + end: int, + splice_channel: int, + junction_count: int, + min_spacing: int, + rng: np.random.Generator, +) -> np.ndarray: + """Clear and replace junction marks in one region, returning new positions.""" + + channel = int(splice_channel) + x[channel, int(start) : int(end)] = 0 + positions = sample_junction_positions( + start=start, + end=end, + junction_count=junction_count, + min_spacing=min_spacing, + rng=rng, + ) + x[channel, positions] = 1 + return positions + + +def _sequence_ids_digest(sequence_ids: Sequence[str]) -> str: + digest = hashlib.sha256() + for sequence_id in sequence_ids: + digest.update(str(sequence_id).encode("utf-8")) + digest.update(b"\0") + return digest.hexdigest() + + +def _allocate_array( + storage_dir: Path | None, + name: str, + shape: tuple[int, ...], + dtype: np.dtype | type, + fill_value: float | bool, +) -> np.ndarray: + if storage_dir is not None and all(dimension > 0 for dimension in shape): + storage_dir.mkdir(parents=True, exist_ok=True) + out = np.lib.format.open_memmap( + storage_dir / f"{name}.npy", + mode="w+", + dtype=dtype, + shape=shape, + ) + else: + out = np.empty(shape, dtype=dtype) + out[...] = fill_value + return out + + +def _predict(predictor: Predictor, X: np.ndarray, *, batch_size: int | None = None) -> np.ndarray: + try: + values = predictor.predict(X, batch_size=batch_size) + except TypeError: + values = predictor.predict(X) + return np.asarray(values, dtype=np.float32).reshape(-1) + + +def _reference_batch(arr: np.ndarray, indices: np.ndarray) -> np.ndarray: + if indices.size == 0: + return np.empty((0, arr.shape[1], arr.shape[2]), dtype=arr.dtype) + first = int(indices[0]) + last = int(indices[-1]) + if np.array_equal(indices, np.arange(first, last + 1, dtype=np.int64)): + return arr[first : last + 1] + return arr[indices] + + +def _replicate_rng( + seed: int, + seq_index: int, + operation: str, + junction_count: int | None, + replicate_index: int, +) -> np.random.Generator: + return np.random.default_rng( + np.random.SeedSequence( + [ + int(seed), + int(seq_index), + int(_FAMILY_CODES[operation]), + 0 if junction_count is None else int(junction_count), + int(replicate_index), + ] + ) + ) + + +def region_ablation( + X: np.ndarray, + predictor: Predictor, + *, + schema: str | SequenceSchema = "saluki6", + sequence_ids: Sequence[str] | None = None, + metadata: Sequence[Mapping[str, object]] | None = None, + config: RegionAblationConfig | None = None, + valid_lengths: Sequence[int] | None = None, + cds_channel: str | int | None = None, + splice_channel: str | int | None = None, + reference_batch_size: int | None = None, + mutation_batch_size: int = 512, + sequence_indices: Sequence[int] | None = None, + sequence_start: int | None = None, + sequence_end: int | None = None, + sequence_shard_index: int | None = None, + sequence_shards: int | None = None, + storage_dir: str | Path | None = None, + progress: bool = True, +) -> RegionAblationResult: + """Run repeated region sequence edits and exon-junction density scans.""" + + arr = np.asarray(X) + if arr.ndim != 3: + raise ValueError(f"Expected X with shape (N, C, L), got {arr.shape}") + resolved_schema = get_schema(schema) + if arr.shape[1] < resolved_schema.n_channels: + raise ValueError( + f"X has {arr.shape[1]} channels, but schema {resolved_schema.name!r} " + f"expects {resolved_schema.n_channels}" + ) + cfg = (config or RegionAblationConfig()).normalized() + if int(mutation_batch_size) <= 0: + raise ValueError("mutation_batch_size must be positive") + + n_sequences = int(arr.shape[0]) + ids = tuple(str(i) for i in range(n_sequences)) if sequence_ids is None else tuple(map(str, sequence_ids)) + if len(ids) != n_sequences: + raise ValueError("sequence_ids length must match X.shape[0]") + if metadata is not None and len(metadata) != n_sequences: + raise ValueError("metadata length must match X.shape[0]") + + base_channels = base_channel_indices(resolved_schema) + cds_channel_index = resolve_cds_channel(resolved_schema, cds_channel) + splice_channel_index = _resolve_splice_channel(resolved_schema, splice_channel) + if splice_channel_index in set(base_channels.tolist()): + raise ValueError("splice_channel must not select a nucleotide base channel") + if splice_channel_index == cds_channel_index: + raise ValueError("splice_channel and cds_channel must select different channels") + + analysis_indices = resolve_analysis_indices( + n_sequences, + sequence_indices=sequence_indices, + sequence_start=sequence_start, + sequence_end=sequence_end, + sequence_shard_index=sequence_shard_index, + sequence_shards=sequence_shards, + ) + full_lengths = None if valid_lengths is None else np.asarray(valid_lengths, dtype=np.int64) + if full_lengths is not None: + if full_lengths.shape != (n_sequences,): + raise ValueError(f"valid_lengths must have shape ({n_sequences},)") + if np.any(full_lengths < 0) or np.any(full_lengths > arr.shape[-1]): + raise ValueError("valid_lengths entries are outside the encoded sequence length") + + instances: list[RegionAblationInstance] = [] + skipped: list[SkippedRegionAblation] = [] + transcript_classes: list[str] = [] + cds_by_sequence: dict[int, CDSCodonStarts] = {} + valid_length_by_sequence: dict[int, int] = {} + + def add_skip( + seq_index: int, + transcript_class: str, + operation: str, + region: str, + reason: str, + junction_count: int | None = None, + ) -> None: + skipped.append( + SkippedRegionAblation( + seq_index=seq_index, + sequence_id=ids[seq_index], + transcript_class=transcript_class, + operation=operation, + region=region, + junction_count=junction_count, + reason=reason, + ) + ) + + reporter = ProgressReporter( + "region-ablation: enumerate conditions", + total=int(analysis_indices.size), + unit="transcripts", + enabled=progress, + ) + for seq_index_raw in analysis_indices: + seq_index = int(seq_index_raw) + x = np.asarray(arr[seq_index]) + valid_length = ( + int(full_lengths[seq_index]) + if full_lengths is not None + else valid_length_from_bases(x, base_channels) + ) + valid_length = min(max(valid_length, 0), int(arr.shape[-1])) + valid_length_by_sequence[seq_index] = valid_length + cds = find_cds_codon_starts( + x, + resolved_schema, + valid_length=valid_length, + cds_channel=cds_channel_index, + ) + has_cds = cds.starts.size > 0 and cds.cds_length >= 3 + row_metadata = None if metadata is None else metadata[seq_index] + if has_cds: + transcript_class = "coding" + cds_by_sequence[seq_index] = cds + elif _metadata_reports_coding(row_metadata): + transcript_class = "coding_unresolved" + else: + transcript_class = "noncoding" + transcript_classes.append(transcript_class) + + if transcript_class == "coding_unresolved": + for family, region, _ in _SEQUENCE_FAMILIES: + if cfg.replicates_for(family) > 0: + add_skip(seq_index, transcript_class, family, region, "represented_cds_missing") + if cfg.replicates_for("junction_scatter") > 0: + for count in cfg.junction_counts: + add_skip( + seq_index, + transcript_class, + "junction_scatter", + "cds", + "represented_cds_missing", + count, + ) + reporter.update() + continue + + if transcript_class == "coding": + for family, region, _ in _SEQUENCE_FAMILIES: + n_replicates = cfg.replicates_for(family) + if n_replicates == 0: + continue + bounds = region_bounds( + region, valid_length=valid_length, cds=cds + ) + if bounds is None or bounds[1] <= bounds[0]: + add_skip(seq_index, transcript_class, family, region, "empty_region") + continue + start, end = bounds + instances.append( + RegionAblationInstance( + instance_index=len(instances), + seq_index=seq_index, + sequence_id=ids[seq_index], + transcript_class=transcript_class, + operation=family, + region=region, + valid_length=valid_length, + region_start=int(start), + region_end=int(end), + region_length=int(end - start), + n_replicates=n_replicates, + ) + ) + junction_region = "cds" + junction_bounds = region_bounds( + "cds", valid_length=valid_length, cds=cds + ) + else: + junction_region = "transcript" + junction_bounds = (0, valid_length) + + junction_replicates = cfg.replicates_for("junction_scatter") + if junction_replicates > 0: + assert junction_bounds is not None + start, end = junction_bounds + region_length = int(end - start) + reference_junction_count = int( + np.count_nonzero(x[splice_channel_index, int(start) : int(end)] > 0) + ) + for count in cfg.junction_counts: + if region_length <= count: + add_skip( + seq_index, + transcript_class, + "junction_scatter", + junction_region, + "region_length_not_greater_than_junction_count", + count, + ) + continue + spacing = effective_junction_spacing( + region_length, + count, + cfg.junction_min_spacing, + ) + instances.append( + RegionAblationInstance( + instance_index=len(instances), + seq_index=seq_index, + sequence_id=ids[seq_index], + transcript_class=transcript_class, + operation="junction_scatter", + region=junction_region, + valid_length=valid_length, + region_start=int(start), + region_end=int(end), + region_length=region_length, + n_replicates=junction_replicates, + junction_count=int(count), + reference_junction_count=reference_junction_count, + requested_min_spacing=cfg.junction_min_spacing, + effective_min_spacing=spacing, + ) + ) + reporter.update() + reporter.close(extra=f"{len(instances)} conditions") + + storage_path = None if storage_dir is None else Path(storage_dir) + n_instances = len(instances) + max_replicates = max((instance.n_replicates for instance in instances), default=0) + reference_out = _allocate_array( + storage_path, "reference_predictions", (n_instances,), np.float32, np.nan + ) + ablation_out = _allocate_array( + storage_path, + "ablation_predictions", + (n_instances, max_replicates), + np.float32, + np.nan, + ) + effects_out = _allocate_array( + storage_path, "effects", (n_instances, max_replicates), np.float32, np.nan + ) + mask_out = _allocate_array( + storage_path, "replicate_mask", (n_instances, max_replicates), np.bool_, False + ) + for instance in instances: + mask_out[instance.instance_index, : instance.n_replicates] = True + + X_reference = _reference_batch(arr, analysis_indices) + log_progress( + f"region-ablation: predicting {analysis_indices.size} reference sequences", + enabled=progress, + ) + selected_reference = _predict(predictor, X_reference, batch_size=reference_batch_size) + if selected_reference.shape != (analysis_indices.size,): + raise ValueError("predictor must return one scalar prediction per reference sequence") + reference_by_sequence = { + int(seq_index): float(prediction) + for seq_index, prediction in zip(analysis_indices.tolist(), selected_reference.tolist()) + } + for instance in instances: + reference_out[instance.instance_index] = reference_by_sequence[instance.seq_index] + + mutant_batch: list[np.ndarray] = [] + pending: list[tuple[int, int]] = [] + + def flush_mutants() -> None: + if not mutant_batch: + return + predictions = _predict( + predictor, + np.stack(mutant_batch, axis=0), + batch_size=mutation_batch_size, + ) + if predictions.shape != (len(pending),): + raise ValueError("predictor must return one scalar prediction per mutant sequence") + for prediction, (instance_index, replicate_index) in zip(predictions, pending): + value = np.float32(prediction) + ablation_out[instance_index, replicate_index] = value + effects_out[instance_index, replicate_index] = np.float32( + value - reference_out[instance_index] + ) + mutant_batch.clear() + pending.clear() + + total_mutants = sum(instance.n_replicates for instance in instances) + mutation_reporter = ProgressReporter( + "region-ablation: predict mutants", + total=total_mutants, + unit="mutants", + enabled=progress, + ) + for instance in instances: + for replicate_index in range(instance.n_replicates): + mutant = np.array(arr[instance.seq_index], copy=True) + rng = _replicate_rng( + cfg.seed, + instance.seq_index, + instance.operation, + instance.junction_count, + replicate_index, + ) + if instance.operation == "junction_scatter": + assert instance.junction_count is not None + assert instance.effective_min_spacing is not None + scatter_junctions_inplace( + mutant, + start=instance.region_start, + end=instance.region_end, + splice_channel=splice_channel_index, + junction_count=instance.junction_count, + min_spacing=instance.effective_min_spacing, + rng=rng, + ) + elif instance.operation == "cds_codon_shuffle": + shuffle_codons_inplace( + mutant, + cds=cds_by_sequence[instance.seq_index], + base_channels=base_channels, + rng=rng, + ) + elif instance.operation.endswith("_shuffle"): + shuffle_nucleotides_inplace( + mutant, + start=instance.region_start, + end=instance.region_end, + base_channels=base_channels, + rng=rng, + ) + elif instance.operation.endswith("_random"): + randomize_nucleotides_inplace( + mutant, + start=instance.region_start, + end=instance.region_end, + base_channels=base_channels, + rng=rng, + ) + else: # pragma: no cover - guarded by the fixed family table. + raise RuntimeError(f"Unsupported region-ablation operation {instance.operation!r}") + mutant_batch.append(mutant) + pending.append((instance.instance_index, replicate_index)) + mutation_reporter.update() + if len(mutant_batch) >= int(mutation_batch_size): + flush_mutants() + flush_mutants() + mutation_reporter.close(extra=f"{total_mutants} mutants") + + mean_effects = _allocate_array( + storage_path, "mean_effects", (n_instances,), np.float32, np.nan + ) + mean_abs_effects = _allocate_array( + storage_path, "mean_abs_effects", (n_instances,), np.float32, np.nan + ) + std_effects = _allocate_array( + storage_path, "std_effects", (n_instances,), np.float32, np.nan + ) + for instance in instances: + values = np.asarray( + effects_out[instance.instance_index, : instance.n_replicates], + dtype=np.float64, + ) + mean_effects[instance.instance_index] = np.float32(values.mean()) + mean_abs_effects[instance.instance_index] = np.float32(np.abs(values).mean()) + std_effects[instance.instance_index] = np.float32(values.std(ddof=0)) + + for values in ( + reference_out, + ablation_out, + effects_out, + mask_out, + mean_effects, + mean_abs_effects, + std_effects, + ): + if hasattr(values, "flush"): + values.flush() + + return RegionAblationResult( + instances=instances, + skipped=skipped, + reference_predictions=reference_out, + ablation_predictions=ablation_out, + effects=effects_out, + replicate_mask=mask_out, + mean_effects=mean_effects, + mean_abs_effects=mean_abs_effects, + std_effects=std_effects, + analysis_indices=analysis_indices.astype(np.int64, copy=False), + transcript_classes=tuple(transcript_classes), + config=cfg, + input_shape=tuple(int(value) for value in arr.shape), + schema_name=resolved_schema.name, + cds_channel_index=cds_channel_index, + splice_channel_index=splice_channel_index, + sequence_ids=tuple(ids[int(index)] for index in analysis_indices), + storage_dir=storage_path, + ) + + +def _persist_array(path: Path, values: np.ndarray) -> None: + if isinstance(values, np.memmap) and getattr(values, "filename", None) is not None: + source = Path(str(values.filename)) + if source.resolve() == path.resolve(): + values.flush() + return + np.save(path, np.asarray(values)) + + +def save_region_ablation_result( + result: RegionAblationResult, + out_dir: str | Path, + *, + checkpoint: str | Path | None = None, + dataset: str | Path | None = None, + progress: bool = True, +) -> None: + """Save raw arrays, condition tables, and reproducibility metadata.""" + + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + log_progress(f"region-ablation: saving results to {out}", enabled=progress) + arrays = { + "reference_predictions": result.reference_predictions, + "ablation_predictions": result.ablation_predictions, + "effects": result.effects, + "replicate_mask": result.replicate_mask, + "mean_effects": result.mean_effects, + "mean_abs_effects": result.mean_abs_effects, + "std_effects": result.std_effects, + } + for name, values in arrays.items(): + _persist_array(out / f"{name}.npy", values) + save_table(out / "instances.csv", result.instances) + save_table(out / "skipped.csv", result.skipped) + + operation_counts = Counter(instance.operation for instance in result.instances) + skip_counts = Counter(row.reason for row in result.skipped) + class_counts = Counter(result.transcript_classes) + summary = { + "analysis": "region_ablation", + "effect_definition": "ablation_prediction - reference_prediction", + "raw_replicates_saved": True, + "inactive_replicate_fill_value": "NaN", + "coordinate_convention": "zero_based_half_open", + "junction_mark_definition": "nucleotide immediately upstream of exon-exon boundary", + "junction_candidate_interval": "[region_start, region_end - 1)", + "junction_spacing_policy": "soft target relaxed to maximum feasible separation", + "junction_sampling": "uniform compressed-coordinate layouts without replacement", + "seed_policy": "SeedSequence(seed, seq_index, family_code, junction_count_or_0, replicate_index)", + "config": result.config.to_dict(), + "checkpoint": str(checkpoint) if checkpoint is not None else None, + "dataset": str(dataset) if dataset is not None else None, + "input_shape": list(result.input_shape), + "schema": result.schema_name, + "cds_channel_index": int(result.cds_channel_index), + "splice_channel_index": int(result.splice_channel_index), + "analysis_sequence_indices": [int(value) for value in result.analysis_indices], + "sequence_ids_sha256": _sequence_ids_digest(result.sequence_ids), + "n_sequences": int(result.analysis_indices.size), + "n_instances": len(result.instances), + "n_mutants": int(result.replicate_mask.sum()), + "n_skipped": len(result.skipped), + "operation_counts": dict(sorted(operation_counts.items())), + "transcript_class_counts": dict(sorted(class_counts.items())), + "skip_reason_counts": dict(sorted(skip_counts.items())), + "arrays": { + name: {"shape": list(values.shape), "dtype": str(values.dtype)} + for name, values in arrays.items() + }, + } + (out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + log_progress("region-ablation: done", enabled=progress) diff --git a/tests/test_cli_analysis.py b/tests/test_cli_analysis.py index 00c96b1..895e8db 100644 --- a/tests/test_cli_analysis.py +++ b/tests/test_cli_analysis.py @@ -6,6 +6,7 @@ from transcriptml.cli.main import _resolve_evaluate_args, _resolve_interpret_args, build_parser, main from transcriptml.data.bundle import DatasetBundle, save_bundle from transcriptml.data.encoding import encode_saluki_transcript +from transcriptml.models.registry import build_model, save_checkpoint def test_models_cli_list_and_show_json(capsys): @@ -139,7 +140,15 @@ def test_interpret_cli_resolves_named_and_legacy_positional_args(): "epistasis": ["--motif", "AUG"], } - for command in ["ism", "window-ism", "codon-ism", "motif-ablation", "motif-context", "epistasis"]: + for command in [ + "ism", + "window-ism", + "codon-ism", + "region-ablation", + "motif-ablation", + "motif-context", + "epistasis", + ]: named = parser.parse_args( [ command, @@ -186,6 +195,114 @@ def test_interpret_cli_resolves_named_and_legacy_positional_args(): assert _resolve_interpret_args(mixed, parser)["out_dir"] == f"interpret/{command}" +def test_region_ablation_cli_defaults_and_overrides(): + parser = build_parser() + args = parser.parse_args( + [ + "region-ablation", + "--checkpoint", + "model.pt", + "--dataset", + "data", + "--out-dir", + "out", + "--junction-counts", + "1,3,8", + "--n-ablations-for", + "cds_random=7", + "--n-ablations-for", + "5utr_shuffle=0", + ] + ) + + assert args.n_ablations == 100 + assert args.junction_counts == (1, 3, 8) + assert args.junction_min_spacing == 25 + assert args.n_ablations_for == [("cds_random", 7), ("5utr_shuffle", 0)] + + +@pytest.mark.parametrize( + "arguments", + [ + ["--junction-counts", "1,1"], + ["--junction-counts", "0,5"], + ["--n-ablations-for", "unknown=3"], + ["--n-ablations-for", "cds_random=-1"], + ], +) +def test_region_ablation_cli_rejects_invalid_condition_configuration(arguments): + parser = build_parser() + with pytest.raises(SystemExit): + parser.parse_args( + [ + "region-ablation", + "--checkpoint", + "model.pt", + "--dataset", + "data", + "--out-dir", + "out", + *arguments, + ] + ) + + +def test_region_ablation_cli_tiny_bundle_and_checkpoint(tmp_path): + model_config = { + "name": "small_cnn", + "params": { + "in_ch": 6, + "n_filters": 2, + "kernel_size": 3, + "n_layers": 1, + "dropout": 0.0, + "head_hidden": 2, + }, + } + checkpoint = tmp_path / "model.pt" + save_checkpoint(checkpoint, build_model(model_config), model_config) + bundle = DatasetBundle( + X=encode_saluki_transcript( + "AACCCGGGUUUA", + length=12, + cds_positions=[2, 5, 8], + splice_positions=[1, 6, 10], + )[None], + ids=["tx1"], + schema="saluki6", + metadata=[{"cds_length": 9}], + ) + dataset = tmp_path / "dataset" + save_bundle(bundle, dataset) + out = tmp_path / "region_ablation" + + main( + [ + "region-ablation", + "--checkpoint", + str(checkpoint), + "--dataset", + str(dataset), + "--out-dir", + str(out), + "--n-ablations", + "1", + "--junction-counts", + "1,5", + "--device", + "cpu", + "--mutation-batch-size", + "4", + ] + ) + + summary = json.loads((out / "summary.json").read_text(encoding="utf-8")) + assert summary["analysis"] == "region_ablation" + assert summary["n_instances"] == 9 + assert summary["n_mutants"] == 9 + assert np.load(out / "effects.npy").shape == (9, 1) + + def test_interpret_cli_rejects_conflicting_named_and_positional_args(): parser = build_parser() args = parser.parse_args( diff --git a/tests/test_region_ablation.py b/tests/test_region_ablation.py new file mode 100644 index 0000000..46cded6 --- /dev/null +++ b/tests/test_region_ablation.py @@ -0,0 +1,304 @@ +import csv +import json + +import numpy as np +import pytest + +from transcriptml.data.encoding import encode_saluki_transcript +from transcriptml.data.schemas import SequenceSchema +from transcriptml.interpret.region_ablation import ( + REGION_ABLATION_FAMILIES, + RegionAblationConfig, + effective_junction_spacing, + region_ablation, + sample_junction_positions, + save_region_ablation_result, +) + + +class RecordingPredictor: + def __init__(self): + self.calls = [] + + def predict(self, X, batch_size=None): + arr = np.asarray(X) + self.calls.append(arr.copy()) + base_score = arr[:, 0].sum(axis=1) + 2 * arr[:, 1].sum(axis=1) + junction_score = 0.25 * arr[:, -1].sum(axis=1) + return (base_score + junction_score).astype(np.float32) + + +def _only(*enabled): + enabled_set = set(enabled) + return {family: int(family in enabled_set) for family in REGION_ABLATION_FAMILIES} + + +def _coding_example(length=24): + return encode_saluki_transcript( + "ACGUAAACCCGGGUUUCGUA", + length=length, + cds_positions=[4, 7, 10, 13], + splice_positions=[2, 8, 17], + ).astype(np.float32) + + +def test_region_sequence_families_preserve_annotations_and_target_bounds(): + X = _coding_example()[None] + predictor = RecordingPredictor() + config = RegionAblationConfig( + n_ablations=1, + n_ablations_for=_only(*(family for family in REGION_ABLATION_FAMILIES if family != "junction_scatter")), + junction_counts=(1,), + seed=7, + ) + result = region_ablation( + X, + predictor, + sequence_ids=["coding"], + metadata=[{"cds_length": 12}], + config=config, + mutation_batch_size=20, + progress=False, + ) + + assert [row.operation for row in result.instances] == list(REGION_ABLATION_FAMILIES[:-1]) + mutants = predictor.calls[1] + assert mutants.shape[0] == 7 + for mutant in mutants: + np.testing.assert_array_equal(mutant[4:], X[0, 4:]) + + before = np.argmax(X[0, :4], axis=0) + bounds = {row.operation: (row.region_start, row.region_end) for row in result.instances} + for row, mutant in zip(result.instances, mutants): + start, end = bounds[row.operation] + if row.operation.endswith("_shuffle") and row.operation != "cds_codon_shuffle": + after = np.argmax(mutant[:4], axis=0) + assert sorted(after[start:end].tolist()) == sorted(before[start:end].tolist()) + if row.operation.endswith("_random"): + assert np.all(mutant[:4, start:end].sum(axis=0) == 1) + + codon_row = next(row for row in result.instances if row.operation == "cds_codon_shuffle") + codon_mutant = mutants[codon_row.instance_index] + original_codons = [X[0, :4, i : i + 3].tobytes() for i in range(4, 16, 3)] + mutant_codons = [codon_mutant[:4, i : i + 3].tobytes() for i in range(4, 16, 3)] + assert sorted(original_codons) == sorted(mutant_codons) + + +def test_junction_scatter_coding_and_noncoding_scope_and_soft_spacing(): + coding = encode_saluki_transcript( + "A" * 60, + length=60, + cds_positions=list(range(5, 56, 3)), + splice_positions=[2, 10, 30, 57], + ).astype(np.float32) + noncoding = encode_saluki_transcript( + "C" * 60, + length=60, + splice_positions=[3, 20, 58], + ).astype(np.float32) + X = np.stack([coding, noncoding]) + predictor = RecordingPredictor() + result = region_ablation( + X, + predictor, + sequence_ids=["coding", "nc"], + metadata=[{"cds_length": 51}, {"cds_length": 0}], + config=RegionAblationConfig( + n_ablations=1, + n_ablations_for=_only("junction_scatter"), + junction_counts=(1, 5, 50), + junction_min_spacing=25, + seed=19, + ), + mutation_batch_size=20, + progress=False, + ) + + assert [(row.seq_index, row.region, row.junction_count) for row in result.instances] == [ + (0, "cds", 1), + (0, "cds", 5), + (0, "cds", 50), + (1, "transcript", 1), + (1, "transcript", 5), + (1, "transcript", 50), + ] + mutants = predictor.calls[1] + for row, mutant in zip(result.instances, mutants): + marks = np.flatnonzero(mutant[5, row.region_start : row.region_end]) + assert marks.size == row.junction_count + assert np.all(marks < row.region_length - 1) + if marks.size > 1: + assert np.diff(marks).min() >= row.effective_min_spacing + np.testing.assert_array_equal(mutant[:5], X[row.seq_index, :5]) + if row.seq_index == 0: + assert mutant[5, 2] == 1 + assert mutant[5, 57] == 1 + assert next(row for row in result.instances if row.seq_index == 0 and row.junction_count == 50).effective_min_spacing == 1 + + +def test_default_grid_has_18_coding_and_11_noncoding_conditions(): + coding = encode_saluki_transcript( + "A" * 70, + length=70, + cds_positions=list(range(5, 62, 3)), + splice_positions=[10, 30], + ) + noncoding = encode_saluki_transcript("C" * 70, length=70, splice_positions=[20]) + result = region_ablation( + np.stack([coding, noncoding]).astype(np.float32), + RecordingPredictor(), + config=RegionAblationConfig(n_ablations=1), + progress=False, + ) + + counts = {seq_index: 0 for seq_index in (0, 1)} + for row in result.instances: + counts[row.seq_index] += 1 + assert counts == {0: 18, 1: 11} + assert int(result.replicate_mask.sum()) == 29 + + +def test_junction_sampling_is_exact_unique_and_rejection_free(): + assert effective_junction_spacing(51, 50, 25) == 1 + assert effective_junction_spacing(100, 5, 25) == 24 + rng = np.random.default_rng(3) + positions = sample_junction_positions( + start=7, + end=107, + junction_count=5, + min_spacing=24, + rng=rng, + ) + assert positions.shape == (5,) + assert len(set(positions.tolist())) == 5 + assert positions.min() >= 7 and positions.max() <= 105 + assert np.diff(positions).min() >= 24 + with pytest.raises(ValueError, match="greater than"): + effective_junction_spacing(5, 5, 25) + + +def test_unresolved_coding_and_infeasible_noncoding_conditions_are_audited(): + X = np.stack( + [ + encode_saluki_transcript("A" * 10, length=10), + encode_saluki_transcript("C" * 3, length=10), + ] + ).astype(np.float32) + result = region_ablation( + X, + RecordingPredictor(), + sequence_ids=["unresolved", "nc"], + metadata=[{"cds_length": 9}, {"cds_length": 0}], + config=RegionAblationConfig( + n_ablations=1, + n_ablations_for=_only("junction_scatter"), + junction_counts=(1, 5), + ), + progress=False, + ) + + assert [(row.seq_index, row.junction_count) for row in result.instances] == [(1, 1)] + assert [(row.seq_index, row.junction_count, row.reason) for row in result.skipped] == [ + (0, 1, "represented_cds_missing"), + (0, 5, "represented_cds_missing"), + (1, 5, "region_length_not_greater_than_junction_count"), + ] + assert result.transcript_classes == ("coding_unresolved", "noncoding") + + +def test_region_ablation_reproducible_across_batching_and_sharding(): + X = np.stack([_coding_example(), _coding_example()]) + config = RegionAblationConfig( + n_ablations=3, + n_ablations_for=_only("cds_random", "junction_scatter"), + junction_counts=(1, 5), + seed=29, + ) + full = region_ablation( + X, + RecordingPredictor(), + sequence_ids=["a", "b"], + config=config, + mutation_batch_size=2, + progress=False, + ) + shard = region_ablation( + X, + RecordingPredictor(), + sequence_ids=["a", "b"], + config=config, + mutation_batch_size=17, + sequence_shard_index=1, + sequence_shards=2, + progress=False, + ) + + full_rows = [row for row in full.instances if row.seq_index == 1] + assert [(row.operation, row.junction_count) for row in full_rows] == [ + (row.operation, row.junction_count) for row in shard.instances + ] + full_indices = [row.instance_index for row in full_rows] + np.testing.assert_array_equal(full.effects[full_indices], shard.effects) + + +def test_storage_serialization_and_custom_channels(tmp_path): + X = _coding_example()[None] + schema = SequenceSchema( + name="custom_saluki", + channels=("A", "C", "G", "U", "coding_marks", "junction_marks"), + ) + out = tmp_path / "region" + result = region_ablation( + X, + RecordingPredictor(), + schema=schema, + sequence_ids=["tx1"], + config=RegionAblationConfig( + n_ablations=2, + n_ablations_for={"junction_scatter": 1}, + junction_counts=(1,), + ), + cds_channel="coding_marks", + splice_channel="junction_marks", + storage_dir=out, + progress=False, + ) + save_region_ablation_result( + result, + out, + checkpoint="fold0/model/best.pt", + dataset="data/saluki", + progress=False, + ) + + expected = { + "reference_predictions.npy", + "ablation_predictions.npy", + "effects.npy", + "replicate_mask.npy", + "mean_effects.npy", + "mean_abs_effects.npy", + "std_effects.npy", + "instances.csv", + "skipped.csv", + "summary.json", + } + assert expected.issubset({path.name for path in out.iterdir()}) + summary = json.loads((out / "summary.json").read_text(encoding="utf-8")) + assert summary["analysis"] == "region_ablation" + assert summary["n_mutants"] == int(result.replicate_mask.sum()) + assert summary["cds_channel_index"] == 4 + assert summary["splice_channel_index"] == 5 + with (out / "instances.csv").open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert rows[0]["sequence_id"] == "tx1" + np.testing.assert_array_equal(np.load(out / "replicate_mask.npy"), result.replicate_mask) + + +def test_region_ablation_config_validation(): + with pytest.raises(ValueError, match="unique"): + RegionAblationConfig(junction_counts=(1, 1)).normalized() + with pytest.raises(ValueError, match="Unknown"): + RegionAblationConfig(n_ablations_for={"bad": 1}).normalized() + with pytest.raises(ValueError, match="non-negative"): + RegionAblationConfig(n_ablations_for={"cds_random": -1}).normalized() From 7aad050e615d199bffcf3ecce35c65a1324b4e84 Mon Sep 17 00:00:00 2001 From: isvock Date: Mon, 17 Aug 2026 14:16:22 -0700 Subject: [PATCH 11/12] Document RBPNet preliminary-ness --- README.md | 5 ++++- docs/api.rst | 12 ++++++++++++ docs/index.rst | 8 ++++++++ docs/installation.md | 9 ++++++++- docs/rbpnet.md | 9 +++++++++ docs/training_configuration.md | 14 ++++++++++++++ docs/usage.md | 11 ++++++++++- scripts/README.md | 7 ++++++- scripts/rbpnet/README.md | 7 +++++++ scripts/rbpnet/rbpnet_config.sh | 3 +++ 10 files changed, 81 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index c912872..5056f25 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,10 @@ TranscriptML currently supports three main workflows: mature-transcript or full-gene coordinate experiment, descriptive windows, explicit selection manifests, and memory-mappable model-ready NumPy bundles, then trains a sequence-only target/control profile model with an optional replicate-aware - enrichment likelihood. + enrichment likelihood. **This entire workflow is experimental:** it has been + minimally tested and has only been confirmed to preprocess data successfully + and train reasonable models on PUM2 eCLIP data. It needs substantially more + validation than the other TranscriptML workflows. In the future, I plan to also support [RiboNN](https://www.nature.com/articles/s41587-025-02712-x) modeling of translation efficiency measurements and extend [RBPNet](https://link.springer.com/article/10.1186/s13059-023-03015-7) interpretation and model variants. diff --git a/docs/api.rst b/docs/api.rst index 89ce318..cbbe1f9 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -46,6 +46,13 @@ Sequence controls RBPNet/eCLIP data ----------------- +.. warning:: + + All RBPNet/eCLIP APIs in this section are experimental. Preprocessing and + modeling have been minimally tested and have only been confirmed to process + data successfully and train reasonable models on PUM2 eCLIP data. They need + substantially more validation than other TranscriptML APIs. + .. automodule:: transcriptml.rbpnet.preprocessing :members: Sample, PipelineConfig, preprocess_eclip :member-order: bysource @@ -109,6 +116,11 @@ Models :members: SmallCNNConfig, SmallCNN :member-order: bysource +.. warning:: + + The RBPNet model API below is experimental and has only received minimal + validation on PUM2 eCLIP data. + .. automodule:: transcriptml.models.rbpnet :members: RBPNetConfig, RBPNetOutput, RBPNet, SamePadConv1d, SameLengthConvTranspose1d, theoretical_receptive_field :member-order: bysource diff --git a/docs/index.rst b/docs/index.rst index 6ee6425..e036f39 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -19,6 +19,14 @@ RBPNet/eCLIP preprocessing, descriptive scanning, region selection, materialized dataset construction, and structured profile/enrichment training are supported. +.. warning:: + + The entire RBPNet/eCLIP workflow is experimental, including preprocessing, + scanning and selection, bundle construction, modeling, and evaluation. It + has been minimally tested and has only been confirmed to preprocess data + successfully and train reasonable models on PUM2 eCLIP data. It needs + substantially more validation than other TranscriptML functionality. + Start here ---------- diff --git a/docs/installation.md b/docs/installation.md index 0e28cf8..a651194 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -40,11 +40,18 @@ TranscriptML use. Optional extras are available for a few heavier workflows: +```{warning} +The `rbpnet` extra enables an experimental workflow. RBPNet/eCLIP preprocessing +and modeling have been minimally tested and have only been confirmed to work +and train reasonable models on PUM2 eCLIP data; they need substantially more +validation than other TranscriptML functionality. +``` + | Use case | Install | | --- | --- | | Write codon-ISM tables as Parquet or Arrow | `python -m pip install -e ".[arrow]"` | | Summarize and plot codon-ISM tables | `python -m pip install -e ".[analysis]"` | -| Preprocess eCLIP and build RBPNet datasets | `python -m pip install -e ".[rbpnet]"` | +| Preprocess eCLIP and build experimental RBPNet datasets | `python -m pip install -e ".[rbpnet]"` | | Run the test suite | `python -m pip install -e ".[dev]"` | Extras can be combined. The `analysis` extra already includes `pyarrow`, so you diff --git a/docs/rbpnet.md b/docs/rbpnet.md index f757875..82c7adc 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -1,5 +1,14 @@ # RBPNet/eCLIP data workflow +```{warning} +Everything in this guide is experimental: eCLIP preprocessing, window scanning +and selection, bundle construction, RBPNet modeling, training, and evaluation. +The workflow has been minimally tested and has only been confirmed to process +data successfully and train reasonable models on PUM2 eCLIP data. It needs +substantially more validation should +not be treated as production-ready or broadly validated. +``` + TranscriptML includes an eCLIP path from ordinary alignments through fixed-shape, memory-mappable arrays and structured RBPNet training. The scanner is descriptive and the selectors prepare model examples; none is intended as a diff --git a/docs/training_configuration.md b/docs/training_configuration.md index 812cbd0..595765c 100644 --- a/docs/training_configuration.md +++ b/docs/training_configuration.md @@ -4,6 +4,13 @@ TranscriptML model training is controlled by a JSON or TOML file. The same top-level training settings are used across Saluki, MPRA-LegNet, and structured RBPNet runs; the model and loss determine the batch contract. +```{warning} +RBPNet preprocessing and modeling are experimental. They have been minimally +tested and have only been confirmed to preprocess data successfully and train +reasonable models on PUM2 eCLIP data. They need substantially more validation +than other TranscriptML functionality. +``` + Create a starter JSON config with: ```bash @@ -108,6 +115,13 @@ section. The Sherlock MPRA workflow has its own editable base config at ## RBPNet Starter Configuration +```{warning} +This starter config is part of the experimental RBPNet workflow. Successful +execution and reasonable training behavior have been checked on PUM2 eCLIP, +but the preprocessing, model, losses, and evaluation require broader +validation before scientific or production use. +``` + `transcriptml init-run --workflow rbpnet` selects the structured RBPNet trainer. Edit the bundle path, output path, and `profile_length` to match the bundle: diff --git a/docs/usage.md b/docs/usage.md index 76ad08a..7fa9453 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -8,6 +8,14 @@ This page walks through the two main TranscriptML workflows: The assay-aware RBPNet/eCLIP preprocessing, window scanning, selection, bundle, and structured training workflow has its own [RBPNet guide](rbpnet.md). +```{warning} +The entire RBPNet/eCLIP workflow is experimental, including preprocessing, +scanning and selection, bundle construction, modeling, and evaluation. It has +been minimally tested and has only been confirmed to preprocess data +successfully and train reasonable models on PUM2 eCLIP data. It needs +substantially more validation than other TranscriptML functionality. +``` + For each workflow, the basic pattern is the same: 1. Build a TranscriptML dataset bundle. @@ -870,4 +878,5 @@ transcriptml init-run --workflow rbpnet --out-dir configs/rbpnet These commands are intentionally small. They are meant to make the first run easier, not to replace project-specific judgment about splits, targets, and -biological grouping. +biological grouping. The RBPNet starter belongs to the experimental workflow +described above and should not be interpreted as a broadly validated default. diff --git a/scripts/README.md b/scripts/README.md index eeddec2..7b1e5d1 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -18,7 +18,12 @@ These scripts are intentionally Sherlock-specific and deliberately small. For a - `motif_ablation_by_fold.sh` and `motif_ablation_all_folds.sh`: motif ablations across the configured motif list. - `motif_epistasis_by_fold.sh` and `motif_epistasis_all_folds.sh`: motif epistasis across the configured motif-pair list. - `mpra/`: MPRA 3-prime UTR insert workflows for building 4-channel LegNet input, training LegNet, and running single-nucleotide ISM. See `mpra/README.md`. -- `rbpnet/`: eCLIP preprocessing, scan/selection/bundle construction, immutable balanced chromosome CV planning, RBPNet training, and structured scientific evaluation per chromosome fold. See `rbpnet/README.md`. +- `rbpnet/`: **experimental** eCLIP preprocessing, scan/selection/bundle construction, immutable balanced chromosome CV planning, RBPNet training, and structured scientific evaluation per chromosome fold. See `rbpnet/README.md`. + +> **Experimental RBPNet status:** Every RBPNet/eCLIP stage is minimally tested. +> The workflow has only been confirmed to preprocess data successfully and train +> reasonable models on PUM2 eCLIP data, and it needs substantially more +> validation than the other TranscriptML workflows. ## Configure A Run diff --git a/scripts/rbpnet/README.md b/scripts/rbpnet/README.md index 44f0ad1..437decc 100644 --- a/scripts/rbpnet/README.md +++ b/scripts/rbpnet/README.md @@ -1,5 +1,12 @@ # RBPNet Sherlock workflow +> **Warning — experimental workflow:** All preprocessing, scanning, selection, +> bundle construction, training, and evaluation described here are minimally +> tested. The workflow has only been confirmed to process data successfully and +> train reasonable models on PUM2 eCLIP data. It needs substantially more +> validation than other TranscriptML functionality and should not be treated as +> production-ready or broadly validated. + This directory implements the staged workflow: ```text diff --git a/scripts/rbpnet/rbpnet_config.sh b/scripts/rbpnet/rbpnet_config.sh index e782858..e939a51 100644 --- a/scripts/rbpnet/rbpnet_config.sh +++ b/scripts/rbpnet/rbpnet_config.sh @@ -2,6 +2,9 @@ # Shared Sherlock defaults for the eCLIP -> RBPNet chromosome-CV workflow. # Copy scripts/rbpnet to a writable run directory and edit this file there. +# EXPERIMENTAL: Every preprocessing/modeling stage has been minimally tested and +# only confirmed on PUM2 eCLIP data. Substantially more validation is required +# than for other TranscriptML workflows. _RBPNET_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" SCRIPT_CONFIG_DIR="${SCRIPT_CONFIG_DIR:-${_RBPNET_SCRIPT_DIR}}" From db4d3411b08b5a8992184943e71926f652b470ce Mon Sep 17 00:00:00 2001 From: isvock Date: Mon, 17 Aug 2026 14:24:16 -0700 Subject: [PATCH 12/12] Fix docs warning --- docs/rbpnet.md | 1 + docs/training_configuration.md | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/rbpnet.md b/docs/rbpnet.md index 82c7adc..86379aa 100644 --- a/docs/rbpnet.md +++ b/docs/rbpnet.md @@ -464,6 +464,7 @@ over an entire transcriptome. The model bundle uses separate `.npy` files because its selected fixed-shape arrays are simple to inspect and memory-map. Changing selection, context, or jitter does not require reprocessing BAMs. +(rbpnet-model-and-training)= ## 5. RBPNet model and training Create a native starter config and train it with the same TranscriptML command diff --git a/docs/training_configuration.md b/docs/training_configuration.md index 595765c..3308eb9 100644 --- a/docs/training_configuration.md +++ b/docs/training_configuration.md @@ -171,8 +171,8 @@ Edit the bundle path, output path, and `profile_length` to match the bundle: Set `enrichment_head_type` to `linear` (or `mlp`) to add the independent replicate-aware enrichment likelihood. The three RBPNet component weights are independent; `lambda_enrichment` has no effect when the head is disabled. See -the [RBPNet guide](rbpnet.md#rbpnet-model-and-training) for the equations, -bundle fields, and jitter semantics. +the {ref}`RBPNet guide ` for the equations, bundle +fields, and jitter semantics. ## Top-Level Training Settings