diff --git a/docs/api.rst b/docs/api.rst index 42307ee..49a722f 100644 --- a/docs/api.rst +++ b/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 diff --git a/docs/usage.md b/docs/usage.md index a7a30a6..0a4b017 100644 --- a/docs/usage.md +++ b/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. diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index d7a38a0..005dc11 100644 --- a/src/transcriptml/cli/main.py +++ b/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 diff --git a/src/transcriptml/training/__init__.py b/src/transcriptml/training/__init__.py index 6ae8706..aa57922 100644 --- a/src/transcriptml/training/__init__.py +++ b/src/transcriptml/training/__init__.py @@ -1,6 +1,6 @@ """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 @@ -8,6 +8,7 @@ __all__ = [ "TrainConfig", "evaluate_checkpoint", + "evaluate_fold_checkpoints", "mse", "pearson_corr", "predict_to_csv", diff --git a/src/transcriptml/training/evaluation.py b/src/transcriptml/training/evaluation.py index 805bf51..ffa6830 100644 --- a/src/transcriptml/training/evaluation.py +++ b/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, diff --git a/src/transcriptml/workflows/__init__.py b/src/transcriptml/workflows/__init__.py index cc7c27c..a833f09 100644 --- a/src/transcriptml/workflows/__init__.py +++ b/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"] diff --git a/src/transcriptml/workflows/cv.py b/src/transcriptml/workflows/cv.py index 1e8f8b6..c2721a4 100644 --- a/src/transcriptml/workflows/cv.py +++ b/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/`` 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() diff --git a/tests/test_evaluation_ensemble.py b/tests/test_evaluation_ensemble.py new file mode 100644 index 0000000..cc6b4ff --- /dev/null +++ b/tests/test_evaluation_ensemble.py @@ -0,0 +1,178 @@ +import csv +import json +from pathlib import Path + +import numpy as np +import pytest +import torch + +from transcriptml.cli.main import main +from transcriptml.data.bundle import DatasetBundle, save_bundle +from transcriptml.training import evaluation +from transcriptml.workflows.cv import find_fold_checkpoints + + +class OffsetModel(torch.nn.Module): + def __init__(self, offset): + super().__init__() + self.offset = float(offset) + + def forward(self, x): + return x[:, 0, 0] + self.offset + + +def _write_dataset(path, *, targets=True): + X = np.zeros((2, 4, 3), dtype=np.float32) + X[:, 0, 0] = [1.0, 2.0] + bundle = DatasetBundle( + X=X, + y=np.array([5.0, 7.0], dtype=np.float32) if targets else None, + ids=["tx1", "tx2"], + schema="rna4", + ) + save_bundle(bundle, path) + + +def _write_checkpoint(path): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"test checkpoint placeholder") + + +def test_evaluate_fold_checkpoints_writes_average_predictions_and_residuals(tmp_path, monkeypatch): + dataset = tmp_path / "dataset" + _write_dataset(dataset) + checkpoints = [tmp_path / "fold0.pt", tmp_path / "fold1.pt"] + for path in checkpoints: + _write_checkpoint(path) + + offsets = {"fold0.pt": 0.0, "fold1.pt": 2.0} + monkeypatch.setattr( + evaluation, + "load_checkpoint", + lambda path, map_location: (OffsetModel(offsets[path.name]), {}), + ) + + out_csv = tmp_path / "ensemble_predictions.csv" + result = evaluation.evaluate_fold_checkpoints( + checkpoints, + dataset, + out_csv, + batch_size=1, + progress=False, + ) + + np.testing.assert_allclose(result["average_predictions"], [2.0, 3.0]) + np.testing.assert_allclose(result["targets"], [5.0, 7.0]) + np.testing.assert_allclose(result["average_residuals"], [3.0, 4.0]) + assert result["fold_count"] == 2 + assert result["mse"] == pytest.approx(12.5) + assert result["pearson"] == pytest.approx(1.0) + assert result["mean_residual"] == pytest.approx(3.5) + + with out_csv.open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert list(rows[0]) == ["index", "id", "target", "average_prediction", "average_residual"] + assert rows == [ + { + "index": "0", + "id": "tx1", + "target": "5.0", + "average_prediction": "2.0", + "average_residual": "3.0", + }, + { + "index": "1", + "id": "tx2", + "target": "7.0", + "average_prediction": "3.0", + "average_residual": "4.0", + }, + ] + + summary = json.loads(out_csv.with_suffix(".summary.json").read_text(encoding="utf-8")) + assert summary["analysis"] == "fold_checkpoint_ensemble" + assert summary["fold_count"] == 2 + assert summary["residual_definition"] == ( + "mean(truth - fold_prediction) = truth - average_prediction" + ) + assert summary["mse"] == pytest.approx(12.5) + assert summary["mean_residual"] == pytest.approx(3.5) + + +def test_evaluate_fold_checkpoints_supports_targetless_dataset(tmp_path, monkeypatch): + dataset = tmp_path / "dataset" + _write_dataset(dataset, targets=False) + checkpoint = tmp_path / "fold0.pt" + _write_checkpoint(checkpoint) + monkeypatch.setattr( + evaluation, + "load_checkpoint", + lambda path, map_location: (OffsetModel(1.0), {}), + ) + + out_csv = tmp_path / "predictions.csv" + result = evaluation.evaluate_fold_checkpoints( + [checkpoint], + dataset, + out_csv, + progress=False, + ) + + np.testing.assert_allclose(result["average_predictions"], [2.0, 3.0]) + assert "average_residuals" not in result + assert "targets" not in result + with out_csv.open(newline="", encoding="utf-8") as handle: + assert next(csv.reader(handle)) == ["index", "id", "average_prediction"] + + summary = json.loads(out_csv.with_suffix(".summary.json").read_text(encoding="utf-8")) + assert summary["target_available"] is False + assert "mean_residual" not in summary + + +def test_evaluate_fold_checkpoints_validates_checkpoint_inputs(tmp_path): + dataset = tmp_path / "dataset" + _write_dataset(dataset) + + with pytest.raises(ValueError, match="at least one"): + evaluation.evaluate_fold_checkpoints([], dataset, progress=False) + with pytest.raises(FileNotFoundError, match="does not exist"): + evaluation.evaluate_fold_checkpoints([tmp_path / "missing.pt"], dataset, progress=False) + + +def test_cv_ensemble_predict_discovers_checkpoints_in_natural_order(tmp_path, monkeypatch, capsys): + dataset = tmp_path / "dataset" + _write_dataset(dataset) + cv_root = tmp_path / "cv" + for fold in (10, 2): + _write_checkpoint(cv_root / f"fold{fold}" / "model" / "best.pt") + + checkpoint_paths = find_fold_checkpoints(cv_root) + assert [path.parent.parent.name for path in checkpoint_paths] == ["fold2", "fold10"] + + def load_checkpoint(path, map_location): + offset = 0.0 if path.parent.parent.name == "fold2" else 2.0 + return OffsetModel(offset), {} + + monkeypatch.setattr(evaluation, "load_checkpoint", load_checkpoint) + out_csv = tmp_path / "ensemble.csv" + main( + [ + "cv", + "ensemble-predict", + "--cv-root", + str(cv_root), + "--dataset", + str(dataset), + "--out-csv", + str(out_csv), + "--device", + "cpu", + ] + ) + + assert capsys.readouterr().out.strip() == str(out_csv) + with out_csv.open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert [float(row["average_prediction"]) for row in rows] == [2.0, 3.0] + summary = json.loads(out_csv.with_suffix(".summary.json").read_text(encoding="utf-8")) + assert [Path(path).parent.parent.name for path in summary["checkpoint_paths"]] == ["fold2", "fold10"]