feat: Add Prov-GigaPath linear probe test workflows and prediction-map utilities

configs/experiment/ml/final_linear_provgigapath_adamw.yaml

@@ -0,0 +1,21 @@
+# @package _global_
+
+defaults:
+  - /experiment/ml/final_linear_virchow2_adamw
+  - _self_
+
+embedding_model_name: ProvGigaPath
+embedding_dim: 1536
+embedding_run_id: 410c8672471348ceb4c58817f70fa097
+kfold_strategy: stratified_group
+kfold_run_id: ${dataset.mlflow_artifacts.stratified_group_kfold_run_id}
+mlflow_artifact_path: linear_classifier_final_provgigapath
+
+# Set after Stage 1 from ProvGigaPath's own AdamW sweep selected by
+# validation/f1_macro.
+model:
+  weight_decay: 1.0e-4
+
+metadata:
+  run_name: Final Linear Classifier AdamW ProvGigaPath ${dataset.name}
+  description: "Final AdamW linear probe over frozen ProvGigaPath embeddings, trained on all training folds with the ProvGigaPath-selected weight decay."

configs/experiment/ml/final_linear_provgigapath_lbfgs.yaml

@@ -0,0 +1,21 @@
+# @package _global_
+
+defaults:
+  - /experiment/ml/final_linear_virchow2_lbfgs
+  - _self_
+
+embedding_model_name: ProvGigaPath
+embedding_dim: 1536
+embedding_run_id: 410c8672471348ceb4c58817f70fa097
+kfold_strategy: stratified_group
+kfold_run_id: ${dataset.mlflow_artifacts.stratified_group_kfold_run_id}
+mlflow_artifact_path: linear_classifier_final_provgigapath
+
+# Set after Stage 1 from ProvGigaPath's own LBFGS sweep selected by
+# validation/f1_macro.
+model:
+  weight_decay: 1.0e-4
+
+metadata:
+  run_name: Final Linear Classifier LBFGS ProvGigaPath ${dataset.name}
+  description: "Final LBFGS linear probe over frozen ProvGigaPath embeddings, exact full-batch solve with the ProvGigaPath-selected weight decay."

configs/experiment/ml/test_linear_provgigapath_adamw.yaml

@@ -0,0 +1,16 @@
+# @package _global_
+
+defaults:
+  - /experiment/ml/final_linear_provgigapath_adamw
+  - _self_
+
+# Held-out test for the final ProvGigaPath AdamW checkpoint. Uses the same
+# filtered labeled test split, thresholds, metrics, and checkpoint convention as
+# the Virchow2 test config.
+mode: test
+final_train_run_id: fe172ccd8c1140269f7f3d1fdbd351ea
+checkpoint: mlflow-artifacts:/104/${final_train_run_id}/artifacts/checkpoints/last/checkpoint.ckpt
+checkpoint_weights_only: false
+
+data:
+  num_workers: 0

configs/experiment/ml/test_linear_provgigapath_lbfgs.yaml

@@ -0,0 +1,18 @@
+# @package _global_
+
+defaults:
+  - /experiment/ml/final_linear_provgigapath_lbfgs
+  - override /ml/trainer: early_stopping
+  - _self_
+
+# Held-out test for the final ProvGigaPath LBFGS checkpoint. Uses the same
+# filtered labeled test split, thresholds, metrics, and checkpoint convention as
+# the Virchow2 test config.
+mode: test
+final_train_run_id: 067b08dcbdb54d9187fbd4dd8d5599a1
+checkpoint: mlflow-artifacts:/104/${final_train_run_id}/artifacts/checkpoints/last/checkpoint.ckpt
+checkpoint_weights_only: false
+
+data:
+  train_batch_size: 1024
+  num_workers: 0

configs/experiment/ml/linear_classifier_test_adamw.yaml → configs/experiment/ml/test_linear_virchow2_adamw.yaml

@@ -1,7 +1,7 @@
 # @package _global_
 
 defaults:
-  - /experiment/ml/linear_classifier_final_adamw
+  - /experiment/ml/final_linear_virchow2_adamw
   - _self_
 
 # Test the AdamW final checkpoint on the held-out test split. Same model
@@ -20,13 +20,3 @@ checkpoint_weights_only: false
 # before the first test batch. final_embedding_tiles defaults to 4; override here.
 data:
   num_workers: 0
-
-trainer:
-  callbacks:
-    tiff_prediction_maps:
-      _target_: ml.callbacks.TiffPredictionMapWriter
-      slides_uri: runs:/${dataset.mlflow_artifacts.tiling_run_id}/test_split/slides.parquet
-      artifact_path: prediction_maps_tiff
-      draw_region: central_stride
-      slide_selection: all
-      max_slides: null

configs/experiment/ml/linear_classifier_test_lbfgs.yaml → configs/experiment/ml/test_linear_virchow2_lbfgs.yaml

@@ -1,7 +1,7 @@
 # @package _global_
 
 defaults:
-  - /experiment/ml/linear_classifier_final_lbfgs
+  - /experiment/ml/final_linear_virchow2_lbfgs
   - override /ml/trainer: early_stopping
   - _self_
 

configs/experiment/ml/train_linear_provgigapath_adamw_group_kfold.yaml

@@ -0,0 +1,14 @@
+# @package _global_
+
+defaults:
+  - /experiment/ml/train_linear_virchow2_adamw_group_kfold
+  - _self_
+
+embedding_model_name: ProvGigaPath
+embedding_dim: 1536
+embedding_run_id: 410c8672471348ceb4c58817f70fa097
+mlflow_artifact_path: linear_classifier_provgigapath
+
+metadata:
+  run_name: Linear Classifier ProvGigaPath ${dataset.name} ${kfold_strategy} fold=${val_fold} opt=${model.optimizer} wd=${model.weight_decay}
+  description: "Linear probe over frozen ProvGigaPath embeddings (run ${embedding_run_id}), ${kfold_strategy} kfold metadata ${kfold_run_id}."

configs/experiment/ml/linear_classifier_lbfgs_stratified_kfold.yaml → configs/experiment/ml/train_linear_provgigapath_lbfgs_group_kfold.yaml

@@ -1,7 +1,7 @@
 # @package _global_
 
 defaults:
-  - /experiment/ml/linear_classifier_stratified_kfold
+  - /experiment/ml/train_linear_provgigapath_adamw_group_kfold
   - _self_
 
 trainer:
@@ -11,6 +11,7 @@ data:
   train_batch_size: 1000000000
   train_shuffle: false
   train_drop_last: false
+  num_workers: 0
 
 model:
   optimizer: lbfgs

configs/experiment/ml/linear_classifier_lbfgs_stratified_group_kfold.yaml → configs/experiment/ml/train_linear_virchow2_lbfgs_group_kfold.yaml

@@ -1,7 +1,7 @@
 # @package _global_
 
 defaults:
-  - /experiment/ml/linear_classifier_stratified_group_kfold
+  - /experiment/ml/train_linear_virchow2_adamw_group_kfold
   - _self_
 
 trainer:

configs/experiment/preprocessing/embeddings_virchow2_tissue_tiles_0_5mpp.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+
+defaults:
+  - /experiment/preprocessing/embeddings_virchow2_0_5mpp
+  - _self_
+
+# Embeddings for a deterministic sampled subset of test slides whose tiles
+# intersect the tissue mask. The sample is capped by slide_sample_max_tiles and
+# selected with slide_sample_seed for doctor-review prediction maps.
+splits:
+  - test
+tile_source_run_id: ${dataset.mlflow_artifacts.tissue_stats_run_id}
+tile_source_artifact_template: "tissue_stats/{split}_tiles.parquet"
+tile_filter_column: tile_tissue_coverage
+slide_sample_max_tiles: 2000000
+slide_sample_seed: 0
+
+metadata:
+  run_name: "Embeddings: ${model} tissue tiles"
+  description: "Tile embeddings using ${model} over a sampled held-out test slide subset with tile_tissue_coverage > 0, capped by slide_sample_max_tiles=${slide_sample_max_tiles} and selected with slide_sample_seed=${slide_sample_seed}."

configs/ml/data/final_embedding_tiles.yaml

@@ -21,3 +21,4 @@ data:
     class_indices: ${class_indices}
     thresholds: ${thresholds}
     tissue_prop_min: ${tissue_prop_min}
+    slide_metadata_uri: ${test_slide_metadata_uri}

configs/ml/model/linear_classifier.yaml

@@ -6,7 +6,7 @@ model:
 
   decode_head:
     _target_: torch.nn.Linear
-    in_features: 2560
+    in_features: ${embedding_dim}
     out_features: ${len:${class_indices}}
 
   class_indices: ${class_indices}

configs/ml/task/final_linear_classifier.yaml

@@ -10,6 +10,8 @@ defaults:
 
 mode: fit
 
+embedding_model_name: Virchow2
+embedding_dim: 2560
 embedding_run_id: ${dataset.mlflow_artifacts.embedding_run_id}
 kfold_strategy: stratified
 kfold_run_id: ${dataset.mlflow_artifacts.stratified_kfold_run_id}
@@ -19,6 +21,7 @@ train_embedding_uri: runs:/${embedding_run_id}/train/tiles
 test_embedding_uri: runs:/${embedding_run_id}/test/tiles
 train_metadata_uri: runs:/${kfold_run_id}/kfold_split/kfold_tiles.parquet
 test_metadata_uri: runs:/${filter_tiles_run_id}/filter_tiles/test_tiles.parquet
+test_slide_metadata_uri: runs:/${embedding_run_id}/test/slides.parquet
 
 tissue_prop_min: 0.2
 thresholds:
@@ -34,8 +37,10 @@ mlflow_artifact_path: linear_classifier_final
 
 metadata:
   run_name: Final Linear Classifier ${dataset.name}
-  description: "Final linear probe over frozen Virchow2 embeddings trained on all training folds for ${trainer.max_epochs} epochs."
+  description: "Final linear probe over frozen ${embedding_model_name} embeddings trained on all training folds for ${trainer.max_epochs} epochs."
   hyperparams:
+    embedding_model_name: ${embedding_model_name}
+    embedding_dim: ${embedding_dim}
     embedding_run_id: ${embedding_run_id}
     kfold_strategy: ${kfold_strategy}
     kfold_run_id: ${kfold_run_id}

configs/ml/task/kfold_linear_classifier.yaml

@@ -10,6 +10,8 @@ defaults:
 
 mode: fit
 
+embedding_model_name: Virchow2
+embedding_dim: 2560
 embedding_run_id: ${dataset.mlflow_artifacts.embedding_run_id}
 kfold_strategy: stratified
 kfold_run_id: ${dataset.mlflow_artifacts.stratified_kfold_run_id}
@@ -36,8 +38,10 @@ mlflow_artifact_path: linear_classifier
 
 metadata:
   run_name: Linear Classifier ${dataset.name} ${kfold_strategy} fold=${val_fold} opt=${model.optimizer} wd=${model.weight_decay}
-  description: "Linear probe over frozen Virchow2 embeddings (run ${embedding_run_id}), ${kfold_strategy} kfold metadata ${kfold_run_id}."
+  description: "Linear probe over frozen ${embedding_model_name} embeddings (run ${embedding_run_id}), ${kfold_strategy} kfold metadata ${kfold_run_id}."
   hyperparams:
+    embedding_model_name: ${embedding_model_name}
+    embedding_dim: ${embedding_dim}
     embedding_run_id: ${embedding_run_id}
     kfold_strategy: ${kfold_strategy}
     kfold_run_id: ${kfold_run_id}

configs/ml/trainer/early_stopping.yaml

@@ -19,6 +19,7 @@ trainer:
       _target_: lightning.pytorch.callbacks.ModelCheckpoint
       monitor: train/loss_epoch
       mode: min
+      save_last: true
       save_top_k: 1
       filename: "epoch={epoch}-train_loss={train/loss_epoch:.4f}"
       auto_insert_metric_name: false

configs/preprocessing/embeddings.yaml

@@ -11,6 +11,8 @@ splits:
 tile_source_run_id: ${dataset.mlflow_artifacts.filter_tiles_run_id}
 tile_source_artifact_template: "filter_tiles/{split}_tiles.parquet"
 tile_filter_column: null
+slide_sample_max_tiles: null
+slide_sample_seed: 0
 
 metadata:
   run_name: "Embeddings: ${model}"
@@ -23,3 +25,5 @@ metadata:
     tile_source_run_id: ${tile_source_run_id}
     tile_source_artifact_template: ${tile_source_artifact_template}
     tile_filter_column: ${tile_filter_column}
+    slide_sample_max_tiles: ${slide_sample_max_tiles}
+    slide_sample_seed: ${slide_sample_seed}

ml/callbacks/tiff_prediction_map_writer.py

@@ -1,6 +1,7 @@
 """Write tile predictions as WSI-aligned BigTIFF masks."""
 
 from collections.abc import Mapping
+from hashlib import blake2b
 from pathlib import Path
 from re import sub
 from tempfile import TemporaryDirectory
@@ -517,7 +518,10 @@ def _safe_filename(value: str) -> str:
 
 
 def _slide_prediction_filename(path: str | Path) -> str:
-    return Path(str(path)).with_suffix(".tiff").name
+    path_str = str(path)
+    stem = Path(path_str).stem
+    digest = blake2b(path_str.encode("utf-8"), digest_size=4).hexdigest()
+    return _safe_filename(f"{stem}-{digest}.tiff")
 
 
 def _spread_lut(n_classes: int) -> np.ndarray:

ml/data/datasets/embedding_tiles.py

@@ -27,6 +27,7 @@ def __init__(
         embedding_uri: str | Path,
         meta_df: pd.DataFrame,
         diag: Callable[[str], None],
+        slide_metadata_uri: str | Path | None = None,
     ) -> None:
         diag(f"metadata filtered: {len(meta_df)} rows; reading embeddings")
         joined_keys, embeddings = _load_embeddings_and_join(
@@ -37,6 +38,9 @@ def __init__(
         self.slide_ids = joined_keys.column("slide_id").to_pandas().to_numpy()
         self.xs = joined_keys.column("x").to_pandas().to_numpy(dtype=np.int64)
         self.ys = joined_keys.column("y").to_pandas().to_numpy(dtype=np.int64)
+        self.slide_names_by_id = (
+            _load_slide_names(slide_metadata_uri) if slide_metadata_uri else {}
+        )
         diag(f"dataset ready: {len(self.labels)} samples, dim={embeddings.shape[1]}")
 
     def __len__(self) -> int:
@@ -78,6 +82,7 @@ def __init__(
         tissue_prop_min: float,
         include_folds: list[int] | None = None,
         exclude_folds: list[int] | None = None,
+        slide_metadata_uri: str | Path | None = None,
     ) -> None:
         self.class_indices = class_indices
         diag = _make_diag(type(self).__name__)
@@ -89,7 +94,7 @@ def __init__(
             include_folds,
             exclude_folds,
         )
-        super().__init__(embedding_uri, meta_df, diag)
+        super().__init__(embedding_uri, meta_df, diag, slide_metadata_uri)
 
     def _labels_from_joined_keys(self, joined_keys: pa.Table) -> np.ndarray:
         labels = joined_keys.column("label").to_pandas()
@@ -233,12 +238,13 @@ def __init__(
         tissue_column: str = "tile_tissue_coverage",
         tissue_min: float = 0.0,
         label_value: int = -1,
+        slide_metadata_uri: str | Path | None = None,
     ) -> None:
         self.label_value = label_value
         diag = _make_diag(type(self).__name__)
         diag("filtering metadata")
         meta_df = self._filter_metadata(metadata_uri, tissue_column, tissue_min)
-        super().__init__(embedding_uri, meta_df, diag)
+        super().__init__(embedding_uri, meta_df, diag, slide_metadata_uri)
 
     def _labels_from_joined_keys(self, joined_keys: pa.Table) -> np.ndarray:
         return np.full(joined_keys.num_rows, self.label_value, dtype=np.int64)
@@ -268,6 +274,12 @@ def _resolve_uri(path_or_uri: str | Path) -> str:
     return _resolve_uri_cached(str(path_or_uri))
 
 
+def _load_slide_names(slide_metadata_uri: str | Path) -> dict[str, str]:
+    local = _resolve_uri(slide_metadata_uri)
+    df = pd.read_parquet(local, columns=["id", "path"])
+    return {str(row.id): Path(str(row.path)).name for row in df.itertuples(index=False)}
+
+
 def _make_diag(dataset_name: str) -> Callable[[str], None]:
     t0 = perf_counter()