diff --git a/docs/usage.md b/docs/usage.md index 0a4b017..7e1d86a 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -205,6 +205,25 @@ 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. +For an out-of-fold report on the training dataset, add `--test-only`: + +```bash +transcriptml cv ensemble-predict \ + --cv-root runs/saluki_cv10 \ + --dataset data/saluki \ + --out-csv runs/saluki_cv10/out_of_fold_predictions.csv \ + --test-only \ + --device auto +``` + +This reads each checkpoint's sibling `foldN/dataset/splits.json` and scores +only that fold's test indices. The fold test splits must cover every dataset +example exactly once. Consequently, a standard CV run contributes one +held-out prediction per row rather than averaging predictions from models that +trained on that row. Use this mode for unbiased CV metrics and residual +diagnostics; use the default full-dataset mode when ensembling fold models on a +new external dataset. + The same operation is available from Python when checkpoint paths are managed outside the standard CV directory layout: @@ -230,17 +249,19 @@ The output has one row per dataset example: 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. +By default, `average_prediction` is the mean prediction from all discovered +fold checkpoints. With `--test-only`, it is the single held-out prediction for +that row. In both modes, `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 prediction scope, +checkpoints, fold count, 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 default full-dataset averaging operation. Those rows are out-of-fold +estimates and should be concatenated, or reproduced directly with +`--test-only`, 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 diff --git a/src/transcriptml/cli/main.py b/src/transcriptml/cli/main.py index 005dc11..fcf277a 100644 --- a/src/transcriptml/cli/main.py +++ b/src/transcriptml/cli/main.py @@ -128,6 +128,11 @@ def build_parser() -> argparse.ArgumentParser: ) p_ensemble.add_argument("--batch-size", type=int, default=128) p_ensemble.add_argument("--device", default="cpu") + p_ensemble.add_argument( + "--test-only", + action="store_true", + help="Use only each checkpoint's fold test split for an out-of-fold report", + ) p = sub.add_parser("build-mpra", help="Build an RNA4 MPRA dataset bundle") p.add_argument("table") @@ -380,6 +385,7 @@ def main(argv: list[str] | None = None) -> None: print(config_path) return if args.cv_command == "ensemble-predict": + from transcriptml.progress import log_progress from transcriptml.training.evaluation import evaluate_fold_checkpoints checkpoint_paths = find_fold_checkpoints( @@ -390,12 +396,20 @@ def main(argv: list[str] | None = None) -> None: raise SystemExit( f"No fold*/model/{args.checkpoint_name} checkpoints found under {args.cv_root}" ) + log_progress( + ( + f"ensemble-predict: discovered {len(checkpoint_paths)} checkpoints " + f"under {args.cv_root}; " + f"scope={'fold test sets' if args.test_only else 'full dataset'}" + ) + ) evaluate_fold_checkpoints( checkpoint_paths, args.dataset, args.out_csv, batch_size=args.batch_size, device=args.device, + test_only=args.test_only, ) print(args.out_csv) return diff --git a/src/transcriptml/training/evaluation.py b/src/transcriptml/training/evaluation.py index ffa6830..39531a7 100644 --- a/src/transcriptml/training/evaluation.py +++ b/src/transcriptml/training/evaluation.py @@ -2,6 +2,7 @@ import csv import json +import time from pathlib import Path from typing import Sequence @@ -213,6 +214,19 @@ def _fold_ensemble_to_csv( writer.writerow(row) +def _format_duration(seconds: float) -> str: + """Format a short human-readable duration for progress messages.""" + + seconds = max(0.0, float(seconds)) + if seconds < 60: + return f"{seconds:.1f}s" + minutes, remaining_seconds = divmod(int(round(seconds)), 60) + if minutes < 60: + return f"{minutes}m {remaining_seconds:02d}s" + hours, remaining_minutes = divmod(minutes, 60) + return f"{hours}h {remaining_minutes:02d}m" + + def evaluate_fold_checkpoints( checkpoint_paths: Sequence[str | Path], dataset_path: str | Path, @@ -221,12 +235,14 @@ def evaluate_fold_checkpoints( batch_size: int = 128, device: str | torch.device = "cpu", progress: bool = True, + test_only: bool = False, ) -> 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. + By default every checkpoint scores every example in the dataset, producing + an ensemble prediction. With ``test_only=True``, each checkpoint instead + scores its own fold test split, producing one out-of-fold prediction per + example. Args: checkpoint_paths: Non-empty sequence of TranscriptML checkpoints. @@ -236,6 +252,9 @@ def evaluate_fold_checkpoints( 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. + test_only: Whether each checkpoint should score only the test indices + from its sibling fold ``dataset/splits.json``. Test splits must + cover every dataset example exactly once. Returns: A dictionary containing ``average_predictions``, example identifiers @@ -254,36 +273,118 @@ def evaluate_fold_checkpoints( raise FileNotFoundError(f"Checkpoint does not exist: {missing[0]}") resolved_device = resolve_device(device) + ensemble_start = time.monotonic() + prediction_scope = "fold test sets" if test_only else "full dataset" + log_progress( + ( + f"ensemble: starting {len(paths)} checkpoints; " + f"scope={prediction_scope}, batch_size={int(batch_size)}, " + f"device={resolved_device}" + ), + enabled=progress, + ) + test_indices_by_checkpoint = None + if test_only: + from transcriptml.workflows.cv import load_fold_test_indices + + log_progress("ensemble: loading fold test split assignments", enabled=progress) + test_indices_by_checkpoint = load_fold_test_indices(paths) + 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) + if test_indices_by_checkpoint is None: + checkpoint_indices = [indices] * len(paths) + predictions_per_example = len(paths) + else: + checkpoint_indices = [] + coverage = np.zeros(indices.shape[0], dtype=np.int64) + for checkpoint_path, raw_indices in zip(paths, test_indices_by_checkpoint): + fold_indices = np.asarray(raw_indices, dtype=int) + if fold_indices.ndim != 1: + raise ValueError(f"Test indices for {checkpoint_path} must be one-dimensional") + if np.unique(fold_indices).shape[0] != fold_indices.shape[0]: + raise ValueError(f"Test split contains duplicate indices for {checkpoint_path}") + if np.any(fold_indices < 0) or np.any(fold_indices >= indices.shape[0]): + raise ValueError( + f"Test split for {checkpoint_path} contains an index outside " + f"[0, {indices.shape[0]})" + ) + coverage[fold_indices] += 1 + checkpoint_indices.append(fold_indices) + missing_count = int(np.count_nonzero(coverage == 0)) + repeated_count = int(np.count_nonzero(coverage > 1)) + if missing_count or repeated_count: + raise ValueError( + "Fold test splits must cover every dataset example exactly once; " + f"missing={missing_count}, repeated={repeated_count}" + ) + predictions_per_example = 1 + + target_status = "targets available" if bundle.y is not None else "no targets" + log_progress( + ( + f"ensemble: dataset ready: {indices.shape[0]:,} examples; " + f"{target_status}; {predictions_per_example} prediction(s) per example" + ), + enabled=progress, + ) - for fold_number, checkpoint_path in enumerate(paths, start=1): + for fold_number, (checkpoint_path, fold_indices) in enumerate( + zip(paths, checkpoint_indices), + start=1, + ): + checkpoint_start = time.monotonic() log_progress( - f"ensemble: loading checkpoint {fold_number}/{len(paths)}: {checkpoint_path}", + ( + f"ensemble: checkpoint {fold_number}/{len(paths)}: loading {checkpoint_path}; " + f"scoring {fold_indices.shape[0]:,} examples" + ), enabled=progress, ) model, _ = load_checkpoint(checkpoint_path, map_location=resolved_device) predictions = _predict_indexed_array( model, bundle.X, - indices, + fold_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: + if predictions.shape != fold_indices.shape: raise ValueError( f"Checkpoint {checkpoint_path} returned {predictions.shape[0]} predictions; " - f"expected {prediction_sum.shape[0]}" + f"expected {fold_indices.shape[0]}" ) - prediction_sum += predictions + if test_only: + prediction_sum[fold_indices] += predictions + else: + prediction_sum += predictions del model - average_predictions64 = prediction_sum / len(paths) + checkpoint_elapsed = time.monotonic() - checkpoint_start + total_elapsed = time.monotonic() - ensemble_start + remaining_checkpoints = len(paths) - fold_number + timing = f"completed in {_format_duration(checkpoint_elapsed)}" + if remaining_checkpoints: + estimated_remaining = (total_elapsed / fold_number) * remaining_checkpoints + timing += f"; estimated remaining {_format_duration(estimated_remaining)}" + log_progress( + f"ensemble: checkpoint {fold_number}/{len(paths)} complete; {timing}", + enabled=progress, + ) + + log_progress( + ( + f"ensemble: combining {prediction_scope} predictions " + f"({predictions_per_example} per example)" + ), + enabled=progress, + ) + average_predictions64 = prediction_sum / predictions_per_example average_predictions = average_predictions64.astype(np.float32) result: dict[str, object] = { "average_predictions": average_predictions, @@ -291,6 +392,8 @@ def evaluate_fold_checkpoints( "ids": [str(identifier) for identifier in bundle.ids], "fold_count": len(paths), "checkpoint_paths": [str(path) for path in paths], + "prediction_scope": "test_only" if test_only else "full_dataset", + "predictions_per_example": predictions_per_example, } targets = None @@ -315,7 +418,16 @@ def evaluate_fold_checkpoints( ), } ) + log_progress( + ( + f"ensemble: metrics: mse={result['mse']:.6g}, " + f"pearson={result['pearson']:.6g}, " + f"mean_residual={result['mean_residual']:.6g}" + ), + enabled=progress, + ) + output_message = "" if out_csv is not None: out_path = Path(out_csv) log_progress(f"ensemble: writing predictions to {out_path}", enabled=progress) @@ -333,6 +445,8 @@ def evaluate_fold_checkpoints( "fold_count": len(paths), "checkpoint_paths": [str(path) for path in paths], "n_examples": int(indices.shape[0]), + "prediction_scope": result["prediction_scope"], + "predictions_per_example": predictions_per_example, "target_available": targets is not None, "residual_definition": "mean(truth - fold_prediction) = truth - average_prediction", "output_csv": str(out_path), @@ -348,8 +462,15 @@ def evaluate_fold_checkpoints( 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") + output_message = f"; predictions={out_path}, summary={summary_path}" - log_progress("ensemble: done", enabled=progress) + log_progress( + ( + f"ensemble: done: {indices.shape[0]:,} examples across {len(paths)} checkpoints " + f"in {_format_duration(time.monotonic() - ensemble_start)}{output_message}" + ), + enabled=progress, + ) return result diff --git a/src/transcriptml/workflows/__init__.py b/src/transcriptml/workflows/__init__.py index a833f09..f232762 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 find_fold_checkpoints, prepare_cv_fold +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", "prepare_cv_fold"] +__all__ = ["find_fold_checkpoints", "init_run", "load_fold_test_indices", "prepare_cv_fold"] diff --git a/src/transcriptml/workflows/cv.py b/src/transcriptml/workflows/cv.py index c2721a4..a2adc8a 100644 --- a/src/transcriptml/workflows/cv.py +++ b/src/transcriptml/workflows/cv.py @@ -3,7 +3,7 @@ import json import re from pathlib import Path -from typing import Any, Mapping +from typing import Any, Mapping, Sequence import numpy as np @@ -30,6 +30,30 @@ def find_fold_checkpoints( return [path for _, path in paths] +def load_fold_test_indices(checkpoint_paths: Sequence[str | Path]) -> list[list[int]]: + """Load each checkpoint's sibling ``dataset/splits.json`` test indices.""" + + test_indices: list[list[int]] = [] + for checkpoint in checkpoint_paths: + checkpoint_path = Path(checkpoint) + split_path = checkpoint_path.parent.parent / "dataset" / "splits.json" + if not split_path.is_file(): + raise FileNotFoundError( + f"Test-only ensemble prediction requires fold split file: {split_path}" + ) + splits = json.loads(split_path.read_text(encoding="utf-8")) + if not isinstance(splits, Mapping) or "test" not in splits: + raise ValueError(f"Fold split file has no 'test' split: {split_path}") + raw_indices = splits["test"] + if not isinstance(raw_indices, list): + raise ValueError(f"Fold 'test' split must be a list: {split_path}") + try: + test_indices.append([int(index) for index in raw_indices]) + except (TypeError, ValueError) as exc: + raise ValueError(f"Fold 'test' split contains a non-integer index: {split_path}") from exc + return test_indices + + 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 index cc6b4ff..8420b97 100644 --- a/tests/test_evaluation_ensemble.py +++ b/tests/test_evaluation_ensemble.py @@ -9,7 +9,7 @@ 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 +from transcriptml.workflows.cv import find_fold_checkpoints, load_fold_test_indices class OffsetModel(torch.nn.Module): @@ -38,6 +38,15 @@ def _write_checkpoint(path): path.write_bytes(b"test checkpoint placeholder") +def _write_fold_test_split(checkpoint_path, test_indices): + split_path = checkpoint_path.parent.parent / "dataset" / "splits.json" + split_path.parent.mkdir(parents=True, exist_ok=True) + split_path.write_text( + json.dumps({"train": [], "val": [], "test": list(test_indices)}), + encoding="utf-8", + ) + + def test_evaluate_fold_checkpoints_writes_average_predictions_and_residuals(tmp_path, monkeypatch): dataset = tmp_path / "dataset" _write_dataset(dataset) @@ -65,6 +74,8 @@ def test_evaluate_fold_checkpoints_writes_average_predictions_and_residuals(tmp_ 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["prediction_scope"] == "full_dataset" + assert result["predictions_per_example"] == 2 assert result["mse"] == pytest.approx(12.5) assert result["pearson"] == pytest.approx(1.0) assert result["mean_residual"] == pytest.approx(3.5) @@ -92,6 +103,8 @@ def test_evaluate_fold_checkpoints_writes_average_predictions_and_residuals(tmp_ 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["prediction_scope"] == "full_dataset" + assert summary["predictions_per_example"] == 2 assert summary["residual_definition"] == ( "mean(truth - fold_prediction) = truth - average_prediction" ) @@ -170,9 +183,90 @@ def load_checkpoint(path, map_location): ] ) - assert capsys.readouterr().out.strip() == str(out_csv) + captured = capsys.readouterr() + assert captured.out.strip() == str(out_csv) + assert "ensemble-predict: discovered 2 checkpoints" in captured.err + assert "ensemble: dataset ready: 2 examples; targets available" in captured.err + assert "ensemble: checkpoint 1/2 complete" in captured.err + assert "estimated remaining" in captured.err + assert "ensemble: metrics: mse=" in captured.err + assert "ensemble: done: 2 examples across 2 checkpoints" in captured.err 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"] + + +def test_cv_ensemble_predict_test_only_uses_each_folds_test_indices(tmp_path, monkeypatch, capsys): + dataset = tmp_path / "dataset" + _write_dataset(dataset) + cv_root = tmp_path / "cv" + fold2 = cv_root / "fold2" / "model" / "best.pt" + fold10 = cv_root / "fold10" / "model" / "best.pt" + for checkpoint, test_indices in [(fold2, [0]), (fold10, [1])]: + _write_checkpoint(checkpoint) + _write_fold_test_split(checkpoint, test_indices) + + assert load_fold_test_indices([fold2, fold10]) == [[0], [1]] + + 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 / "test_only.csv" + main( + [ + "cv", + "ensemble-predict", + "--cv-root", + str(cv_root), + "--dataset", + str(dataset), + "--out-csv", + str(out_csv), + "--test-only", + ] + ) + + captured = capsys.readouterr() + assert "scope=fold test sets" in captured.err + assert "1 prediction(s) per example" in captured.err + with out_csv.open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert [float(row["average_prediction"]) for row in rows] == [1.0, 4.0] + assert [float(row["average_residual"]) for row in rows] == [4.0, 3.0] + + summary = json.loads(out_csv.with_suffix(".summary.json").read_text(encoding="utf-8")) + assert summary["prediction_scope"] == "test_only" + assert summary["predictions_per_example"] == 1 + + +def test_test_only_requires_complete_nonoverlapping_fold_coverage(tmp_path): + dataset = tmp_path / "dataset" + _write_dataset(dataset) + cv_root = tmp_path / "cv" + checkpoints = [ + cv_root / "fold0" / "model" / "best.pt", + cv_root / "fold1" / "model" / "best.pt", + ] + for checkpoint in checkpoints: + _write_checkpoint(checkpoint) + _write_fold_test_split(checkpoint, [0]) + + with pytest.raises(ValueError, match=r"missing=1, repeated=1"): + evaluation.evaluate_fold_checkpoints( + checkpoints, + dataset, + test_only=True, + progress=False, + ) + + +def test_load_fold_test_indices_requires_split_file(tmp_path): + checkpoint = tmp_path / "cv" / "fold0" / "model" / "best.pt" + _write_checkpoint(checkpoint) + + with pytest.raises(FileNotFoundError, match="requires fold split file"): + load_fold_test_indices([checkpoint])