Add a prediction and residual calculating tool

docs/api.rst

@@ -78,7 +78,7 @@ Training and evaluation
    :member-order: bysource
 
 .. automodule:: transcriptml.training.evaluation
-   :members: predict_array, evaluate_model, predict_to_csv, evaluate_checkpoint
+   :members: predict_array, evaluate_model, predict_to_csv, evaluate_checkpoint, evaluate_fold_checkpoints
    :member-order: bysource
 
 .. automodule:: transcriptml.training.splits

docs/usage.md

@@ -188,6 +188,60 @@ and test metrics. The `eval/test_predictions.csv` file contains held-out
 predictions for that fold. Concatenate the ten fold-level prediction tables to
 measure performance across the full dataset.
 
+To make an ensemble prediction instead, apply every fold checkpoint to the
+same complete dataset and average their outputs:
+
+```bash
+transcriptml cv ensemble-predict \
+  --cv-root runs/saluki_cv10 \
+  --dataset data/saluki \
+  --out-csv runs/saluki_cv10/ensemble_predictions.csv \
+  --batch-size 128 \
+  --device auto
+```
+
+This discovers `fold*/model/best.pt` under the CV root in numeric fold order.
+Use `--checkpoint-name last.pt` to select a different checkpoint filename.
+Models are loaded and evaluated one at a time, so they do not all need to fit
+in device memory simultaneously.
+
+The same operation is available from Python when checkpoint paths are managed
+outside the standard CV directory layout:
+
+```python
+from transcriptml.training.evaluation import evaluate_fold_checkpoints
+
+result = evaluate_fold_checkpoints(
+    [
+        "runs/saluki_cv10/fold0/model/best.pt",
+        "runs/saluki_cv10/fold1/model/best.pt",
+    ],
+    "data/saluki",
+    "runs/saluki_cv10/ensemble_predictions.csv",
+    device="auto",
+)
+average_predictions = result["average_predictions"]
+average_residuals = result["average_residuals"]
+```
+
+The output has one row per dataset example:
+
+```text
+index,id,target,average_prediction,average_residual
+```
+
+Here `average_prediction` is the mean prediction from all discovered fold
+checkpoints, and `average_residual` is `target - average_prediction`. If the
+dataset has no `y.npy`, the prediction columns are still written but target and
+residual columns are omitted. A sibling `ensemble_predictions.summary.json`
+records the checkpoints, fold count, ensemble MSE and Pearson correlation, and
+mean residual.
+
+Do not use the disjoint fold-level `test_predictions.csv` files as inputs to
+this averaging operation. Those rows are out-of-fold estimates and should be
+concatenated, whereas ensemble averaging requires every fold model to score the
+same examples.
+
 The built-in fold assignment is random and transcript-level. If related
 isoforms, homologous transcripts, or other biological groups must stay
 together, create grouped folds upstream and run one predefined split per fold.

src/transcriptml/cli/main.py

@@ -114,6 +114,20 @@ 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_ensemble = cv_sub.add_parser(
+        "ensemble-predict",
+        help="Average predictions from fold checkpoints on one shared dataset",
+    )
+    p_ensemble.add_argument("--cv-root", required=True, help="Directory containing fold*/model checkpoints")
+    p_ensemble.add_argument("--dataset", required=True, help="Shared dataset bundle scored by every fold")
+    p_ensemble.add_argument("--out-csv", required=True, help="Averaged prediction CSV output path")
+    p_ensemble.add_argument(
+        "--checkpoint-name",
+        default="best.pt",
+        help="Checkpoint filename under each fold's model directory",
+    )
+    p_ensemble.add_argument("--batch-size", type=int, default=128)
+    p_ensemble.add_argument("--device", default="cpu")
 
     p = sub.add_parser("build-mpra", help="Build an RNA4 MPRA dataset bundle")
     p.add_argument("table")
@@ -350,7 +364,7 @@ def main(argv: list[str] | None = None) -> None:
                     print(f"{key}\t{value}")
             return
     if args.command == "cv":
-        from transcriptml.workflows import prepare_cv_fold
+        from transcriptml.workflows import find_fold_checkpoints, prepare_cv_fold
 
         if args.cv_command == "prepare-fold":
             config_path = prepare_cv_fold(
@@ -365,6 +379,26 @@ def main(argv: list[str] | None = None) -> None:
             )
             print(config_path)
             return
+        if args.cv_command == "ensemble-predict":
+            from transcriptml.training.evaluation import evaluate_fold_checkpoints
+
+            checkpoint_paths = find_fold_checkpoints(
+                args.cv_root,
+                checkpoint_name=args.checkpoint_name,
+            )
+            if not checkpoint_paths:
+                raise SystemExit(
+                    f"No fold*/model/{args.checkpoint_name} checkpoints found under {args.cv_root}"
+                )
+            evaluate_fold_checkpoints(
+                checkpoint_paths,
+                args.dataset,
+                args.out_csv,
+                batch_size=args.batch_size,
+                device=args.device,
+            )
+            print(args.out_csv)
+            return
     if args.command == "plot-ism":
         from transcriptml.plotting.single_nt_ism import plot_ism_from_args
 

src/transcriptml/training/__init__.py

@@ -1,13 +1,14 @@
 """Training and evaluation utilities."""
 
-from transcriptml.training.evaluation import evaluate_checkpoint, predict_to_csv
+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.trainer import TrainConfig, train_from_config, train_model
 
 __all__ = [
     "TrainConfig",
     "evaluate_checkpoint",
+    "evaluate_fold_checkpoints",
     "mse",
     "pearson_corr",
     "predict_to_csv",

src/transcriptml/training/evaluation.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import csv
+import json
 from pathlib import Path
 from typing import Sequence
 
@@ -175,6 +176,183 @@ def predict_to_csv(
             writer.writerow(row)
 
 
+def _fold_ensemble_to_csv(
+    path: str | Path,
+    *,
+    ids: Sequence[str],
+    average_predictions: Sequence[float],
+    targets: Sequence[float] | None = None,
+    average_residuals: Sequence[float] | None = None,
+    indices: Sequence[int] | None = None,
+) -> None:
+    """Write per-example mean predictions and optional mean residuals."""
+
+    if (targets is None) != (average_residuals is None):
+        raise ValueError("targets and average_residuals must either both be provided or both be omitted")
+
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    with Path(path).open("w", newline="", encoding="utf-8") as handle:
+        fieldnames = ["index", "id"]
+        if targets is not None:
+            fieldnames.append("target")
+        fieldnames.append("average_prediction")
+        if average_residuals is not None:
+            fieldnames.append("average_residual")
+        writer = csv.DictWriter(handle, fieldnames=fieldnames)
+        writer.writeheader()
+        idx = list(range(len(average_predictions))) if indices is None else list(indices)
+        for j, prediction in enumerate(average_predictions):
+            row: dict[str, object] = {
+                "index": int(idx[j]),
+                "id": str(ids[j]),
+                "average_prediction": float(prediction),
+            }
+            if targets is not None and average_residuals is not None:
+                row["target"] = float(targets[j])
+                row["average_residual"] = float(average_residuals[j])
+            writer.writerow(row)
+
+
+def evaluate_fold_checkpoints(
+    checkpoint_paths: Sequence[str | Path],
+    dataset_path: str | Path,
+    out_csv: str | Path | None = None,
+    *,
+    batch_size: int = 128,
+    device: str | torch.device = "cpu",
+    progress: bool = True,
+) -> dict[str, object]:
+    """Average predictions from fold checkpoints evaluated on one shared dataset.
+
+    Every checkpoint scores every example in the dataset. This produces an
+    ensemble prediction rather than an out-of-fold CV prediction; disjoint
+    fold-level test prediction tables should be concatenated instead.
+
+    Args:
+        checkpoint_paths: Non-empty sequence of TranscriptML checkpoints.
+        dataset_path: Shared dataset bundle scored by every checkpoint.
+        out_csv: Optional destination for per-example ensemble predictions and
+            residuals. A sibling ``.summary.json`` file is also written.
+        batch_size: Number of examples to score per prediction batch.
+        device: Torch device used to load and run each model.
+        progress: Whether to emit progress messages while evaluating.
+
+    Returns:
+        A dictionary containing ``average_predictions``, example identifiers
+        and indices, fold provenance, and, when targets are available,
+        ``targets``, ``average_residuals`` (truth minus prediction), MSE,
+        Pearson correlation, and mean residual.
+    """
+
+    paths = [Path(path) for path in checkpoint_paths]
+    if not paths:
+        raise ValueError("checkpoint_paths must contain at least one checkpoint")
+    if int(batch_size) <= 0:
+        raise ValueError("batch_size must be positive")
+    missing = [path for path in paths if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(f"Checkpoint does not exist: {missing[0]}")
+
+    resolved_device = resolve_device(device)
+    log_progress(f"ensemble: loading dataset {dataset_path}", enabled=progress)
+    bundle = load_bundle(dataset_path, mmap_mode="r")
+    indices = np.arange(int(bundle.X.shape[0]), dtype=int)
+    prediction_sum = np.zeros(indices.shape[0], dtype=np.float64)
+
+    for fold_number, checkpoint_path in enumerate(paths, start=1):
+        log_progress(
+            f"ensemble: loading checkpoint {fold_number}/{len(paths)}: {checkpoint_path}",
+            enabled=progress,
+        )
+        model, _ = load_checkpoint(checkpoint_path, map_location=resolved_device)
+        predictions = _predict_indexed_array(
+            model,
+            bundle.X,
+            indices,
+            batch_size=int(batch_size),
+            device=resolved_device,
+            progress=progress,
+            progress_label=f"ensemble: checkpoint {fold_number}/{len(paths)}",
+        )
+        predictions = np.asarray(predictions, dtype=np.float64).reshape(-1)
+        if predictions.shape != prediction_sum.shape:
+            raise ValueError(
+                f"Checkpoint {checkpoint_path} returned {predictions.shape[0]} predictions; "
+                f"expected {prediction_sum.shape[0]}"
+            )
+        prediction_sum += predictions
+        del model
+
+    average_predictions64 = prediction_sum / len(paths)
+    average_predictions = average_predictions64.astype(np.float32)
+    result: dict[str, object] = {
+        "average_predictions": average_predictions,
+        "indices": indices.tolist(),
+        "ids": [str(identifier) for identifier in bundle.ids],
+        "fold_count": len(paths),
+        "checkpoint_paths": [str(path) for path in paths],
+    }
+
+    targets = None
+    average_residuals = None
+    if bundle.y is not None:
+        targets = np.asarray(bundle.y, dtype=np.float32).reshape(-1)
+        if targets.shape != average_predictions.shape:
+            raise ValueError(
+                f"Dataset targets have shape {targets.shape}; expected {average_predictions.shape}"
+            )
+        average_residuals = (targets.astype(np.float64) - average_predictions64).astype(np.float32)
+        result.update(
+            {
+                "targets": targets,
+                "average_residuals": average_residuals,
+                "mse": mse(targets, average_predictions),
+                "pearson": pearson_corr(targets, average_predictions),
+                "mean_residual": (
+                    float(np.mean(average_residuals, dtype=np.float64))
+                    if average_residuals.size
+                    else float("nan")
+                ),
+            }
+        )
+
+    if out_csv is not None:
+        out_path = Path(out_csv)
+        log_progress(f"ensemble: writing predictions to {out_path}", enabled=progress)
+        _fold_ensemble_to_csv(
+            out_path,
+            ids=result["ids"],
+            average_predictions=average_predictions,
+            targets=targets,
+            average_residuals=average_residuals,
+            indices=result["indices"],
+        )
+        summary: dict[str, object] = {
+            "analysis": "fold_checkpoint_ensemble",
+            "dataset": str(dataset_path),
+            "fold_count": len(paths),
+            "checkpoint_paths": [str(path) for path in paths],
+            "n_examples": int(indices.shape[0]),
+            "target_available": targets is not None,
+            "residual_definition": "mean(truth - fold_prediction) = truth - average_prediction",
+            "output_csv": str(out_path),
+        }
+        if targets is not None:
+            summary.update(
+                {
+                    "mse": result["mse"],
+                    "pearson": result["pearson"],
+                    "mean_residual": result["mean_residual"],
+                }
+            )
+        summary_path = out_path.with_suffix(".summary.json")
+        log_progress(f"ensemble: writing summary to {summary_path}", enabled=progress)
+        summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+    log_progress("ensemble: done", enabled=progress)
+    return result
+
+
 def evaluate_checkpoint(
     checkpoint_path: str | Path,
     dataset_path: str | Path,

src/transcriptml/workflows/__init__.py

@@ -1,6 +1,6 @@
 """Workflow template helpers for TranscriptML."""
 
-from transcriptml.workflows.cv import prepare_cv_fold
+from transcriptml.workflows.cv import find_fold_checkpoints, prepare_cv_fold
 from transcriptml.workflows.init_run import init_run
 
-__all__ = ["init_run", "prepare_cv_fold"]
+__all__ = ["find_fold_checkpoints", "init_run", "prepare_cv_fold"]

src/transcriptml/workflows/cv.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import re
 from pathlib import Path
 from typing import Any, Mapping
 
@@ -10,6 +11,25 @@
 BUNDLE_FILES = ("X.npy", "y.npy", "ids.txt", "schema.json", "metadata.json", "config.json")
 
 
+def find_fold_checkpoints(
+    cv_root: str | Path,
+    *,
+    checkpoint_name: str = "best.pt",
+) -> list[Path]:
+    """Find ``foldN/model/<checkpoint_name>`` paths in natural fold order."""
+
+    root = Path(cv_root)
+    paths: list[tuple[int, Path]] = []
+    for fold_dir in root.glob("fold*"):
+        match = re.fullmatch(r"fold(\d+)", fold_dir.name)
+        if match and fold_dir.is_dir():
+            checkpoint_path = fold_dir / "model" / checkpoint_name
+            if checkpoint_path.is_file():
+                paths.append((int(match.group(1)), checkpoint_path))
+    paths.sort(key=lambda item: item[0])
+    return [path for _, path in paths]
+
+
 def _replace_link(src: Path, dst: Path) -> None:
     if dst.exists() or dst.is_symlink():
         dst.unlink()